fix runUserHandler in BaseDurableOperation

Author: zhongkechen

sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java

@@ -224,9 +224,11 @@ protected Operation waitForOperationCompletion() {
         return op;
     }
 
-    protected void runUserHandler(Runnable runnable, String contextId, ThreadType threadType) {
+    protected void runUserHandler(Runnable runnable, ThreadType threadType) {
+        String operationId = getOperationId();
+        logger.debug("Starting user handler for operation {} ({})", operationId, threadType);
         Runnable wrapped = () -> {
-            executionManager.setCurrentThreadContext(new ThreadContext(contextId, threadType));
+            executionManager.setCurrentThreadContext(new ThreadContext(operationId, threadType));
             try {
                 runnable.run();
             } catch (Throwable throwable) {
@@ -239,11 +241,11 @@ protected void runUserHandler(Runnable runnable, String contextId, ThreadType th
                             "An unhandled exception is thrown from user function: " + throwable);
                 }
             } finally {
-                if (contextId != null) {
+                if (operationId != null) {
                     try {
                         // if this is a child context or a step context, we need to
                         // deregister the context's thread from the execution manager
-                        executionManager.deregisterActiveThread(contextId);
+                        executionManager.deregisterActiveThread(operationId);
                     } catch (SuspendExecutionException e) {
                         // Expected when this is the last active thread. Must catch here because:
                         // 1/ This runs in a worker thread detached from handlerFuture
@@ -257,8 +259,10 @@ protected void runUserHandler(Runnable runnable, String contextId, ThreadType th
         };
 
         // runUserHandler is used to ensure that only one user handler is running at a time
-        if (runningUserHandler.get() != null) {
-            throw new IllegalStateException("User handler already running");
+        if (runningUserHandler.get() != null && !runningUserHandler.get().isDone()) {
+            logger.error("User handler already running for operation {} ({})", getOperationId(), threadType);
+            throw terminateExecutionWithIllegalDurableOperationException(
+                    "User handler already running: " + getOperationId());
         }
 
         // Thread registration is intentionally split across two threads:
@@ -267,14 +271,10 @@ protected void runUserHandler(Runnable runnable, String contextId, ThreadType th
         // 2. setCurrentContext on the CHILD thread — sets the ThreadLocal so operations inside
         //    the child context know which context they belong to.
         // registerActiveThread is idempotent (no-op if already registered).
-        registerActiveThread(contextId);
+        registerActiveThread(operationId);
 
-        if (!runningUserHandler.compareAndSet(
-                null,
-                CompletableFuture.runAsync(
-                        wrapped, getContext().getDurableConfig().getExecutorService()))) {
-            throw new IllegalStateException("User handler already running");
-        }
+        runningUserHandler.set(CompletableFuture.runAsync(
+                wrapped, getContext().getDurableConfig().getExecutorService()));
     }
 
     /**

sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java

@@ -129,7 +129,7 @@ private void executeChildContext() {
         };
 
         // Execute user provided child context code in user-configured executor
-        runUserHandler(userHandler, contextId, ThreadType.CONTEXT);
+        runUserHandler(userHandler, ThreadType.CONTEXT);
     }
 
     private void handleChildContextSuccess(T result) {

sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java

@@ -210,7 +210,7 @@ protected void executeItems() {
             }
         };
         // run consumer in the user thread pool, although it's not a real user thread
-        runUserHandler(consumer, getOperationId(), ThreadType.CONTEXT);
+        runUserHandler(consumer, ThreadType.CONTEXT);
     }
 
     private void handleException(Throwable ex) {

sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java

@@ -112,7 +112,7 @@ private void executeStepLogic(int attempt) {
         };
 
         // Execute user provided step code in user-configured executor
-        runUserHandler(userHandler, getOperationId(), ThreadType.STEP);
+        runUserHandler(userHandler, ThreadType.STEP);
     }
 
     private void checkpointStarted() {

sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java

@@ -167,7 +167,7 @@ private void executeCheckLogic(T currentState, int attempt) {
             }
         };
 
-        runUserHandler(userHandler, getOperationId(), ThreadType.STEP);
+        runUserHandler(userHandler, ThreadType.STEP);
     }
 
     private void handleCheckFailure(Throwable exception) {