fix thread pool leak

Author: zhongkechen

examples/src/main/java/software/amazon/lambda/durable/examples/callback/WaitForCallbackFailedExample.java

@@ -9,6 +9,7 @@
 import software.amazon.lambda.durable.config.WaitForCallbackConfig;
 import software.amazon.lambda.durable.examples.types.ApprovalRequest;
 import software.amazon.lambda.durable.exception.SerDesException;
+import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.serde.JacksonSerDes;
 
 public class WaitForCallbackFailedExample extends DurableHandler<ApprovalRequest, String> {
@@ -32,6 +33,10 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
                                     .build())
                             .build());
         } catch (Exception ex) {
+            // not to swallow the SuspendExecutionException
+            if (ex instanceof SuspendExecutionException suspendExecutionException) {
+                throw suspendExecutionException;
+            }
             return ex.getClass().getSimpleName() + ":" + ex.getMessage();
         }
 

sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java

@@ -74,6 +74,20 @@ public final class DurableConfig {
     private static final String PROJECT_VERSION = getProjectVersion(VERSION_FILE);
     private static final String USER_AGENT_SUFFIX = "@aws/durable-execution-sdk-java/" + PROJECT_VERSION;
 
+    /**
+     * A default ExecutorService for running user-defined operations. Uses a cached thread pool with daemon threads by
+     * default.
+     *
+     * <p>This executor is used exclusively for user operations. Internal SDK coordination uses the
+     * InternalExecutor::INSTANCE
+     */
+    private static final ExecutorService DEFAULT_USER_THREAD_POOL = Executors.newCachedThreadPool(r -> {
+        Thread t = new Thread(r);
+        t.setName("durable-exec-" + t.getId());
+        t.setDaemon(true);
+        return t;
+    });
+
     private final DurableExecutionClient durableExecutionClient;
     private final SerDes serDes;
     private final ExecutorService executorService;
@@ -250,12 +264,7 @@ private static String getProjectVersion(String versionFile) {
      */
     private static ExecutorService createDefaultExecutor() {
         logger.debug("Creating default ExecutorService");
-        return Executors.newCachedThreadPool(r -> {
-            Thread t = new Thread(r);
-            t.setName("durable-exec-" + t.getId());
-            t.setDaemon(true);
-            return t;
-        });
+        return DEFAULT_USER_THREAD_POOL;
     }
 
     /** Builder for DurableConfig. Provides fluent API for configuring SDK components. */

sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java

@@ -280,17 +280,23 @@ public static boolean isTerminalStatus(OperationStatus status) {
      * @param exception the unrecoverable exception that caused termination
      */
     public void terminateExecution(UnrecoverableDurableExecutionException exception) {
+        stopAllOperations(exception);
         executionExceptionFuture.completeExceptionally(exception);
         throw exception;
     }
 
     /** Suspends the execution by completing the execution exception future with a {@link SuspendExecutionException}. */
     public void suspendExecution() {
         var ex = new SuspendExecutionException();
+        stopAllOperations(ex);
         executionExceptionFuture.completeExceptionally(ex);
         throw ex;
     }
 
+    private void stopAllOperations(Exception cause) {
+        registeredOperations.values().forEach(op -> op.getCompletionFuture().completeExceptionally(cause));
+    }
+
     /**
      * return a future that completes when userFuture completes successfully or the execution is terminated or
      * suspended.

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

@@ -22,6 +22,7 @@
 import software.amazon.lambda.durable.execution.ThreadType;
 import software.amazon.lambda.durable.model.OperationIdentifier;
 import software.amazon.lambda.durable.model.OperationSubType;
+import software.amazon.lambda.durable.util.ExceptionHelper;
 
 /**
  * Base class for all durable operations (STEP, WAIT, etc.).
@@ -208,7 +209,11 @@ protected Operation waitForOperationCompletion() {
         }
 
         // Block until operation completes. No-op if the future is already completed.
-        completionFuture.join();
+        try {
+            completionFuture.join();
+        } catch (Throwable throwable) {
+            ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(throwable));
+        }
 
         // Get result based on status
         var op = getOperation();

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

@@ -14,11 +14,14 @@
 import software.amazon.lambda.durable.config.CompletionConfig;
 import software.amazon.lambda.durable.config.MapConfig;
 import software.amazon.lambda.durable.context.DurableContextImpl;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
+import software.amazon.lambda.durable.execution.SuspendExecutionException;
 import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus;
 import software.amazon.lambda.durable.model.MapResult;
 import software.amazon.lambda.durable.model.OperationIdentifier;
 import software.amazon.lambda.durable.model.OperationSubType;
 import software.amazon.lambda.durable.serde.SerDes;
+import software.amazon.lambda.durable.util.ExceptionHelper;
 
 /**
  * Executes a map operation: applies a function to each item in a collection concurrently, with each item running in its
@@ -153,8 +156,18 @@ protected void handleCompletion(ConcurrencyCompletionStatus concurrencyCompletio
             } else {
                 try {
                     resultItems.set(i, MapResult.MapResultItem.succeeded(branch.get()));
-                } catch (Exception e) {
-                    resultItems.set(i, MapResult.MapResultItem.failed(MapResult.MapError.of(e)));
+                } catch (Throwable exception) {
+                    Throwable throwable = ExceptionHelper.unwrapCompletableFuture(exception);
+                    if (throwable instanceof SuspendExecutionException suspendExecutionException) {
+                        // Rethrow Error immediately — do not checkpoint
+                        throw suspendExecutionException;
+                    }
+                    if (throwable
+                            instanceof UnrecoverableDurableExecutionException unrecoverableDurableExecutionException) {
+                        // terminate the execution and throw the exception if it's not recoverable
+                        throw terminateExecution(unrecoverableDurableExecutionException);
+                    }
+                    resultItems.set(i, MapResult.MapResultItem.failed(MapResult.MapError.of(throwable)));
                 }
             }
         }

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

@@ -4,7 +4,6 @@
 
 import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutorService;
 import java.util.function.BiFunction;
 import software.amazon.awssdk.services.lambda.model.Operation;
 import software.amazon.awssdk.services.lambda.model.OperationAction;
@@ -40,7 +39,6 @@ public class WaitForConditionOperation<T> extends SerializableDurableOperation<T
 
     private final BiFunction<T, StepContext, WaitForConditionResult<T>> checkFunc;
     private final WaitForConditionConfig<T> config;
-    private final ExecutorService userExecutor;
 
     public WaitForConditionOperation(
             String operationId,
@@ -57,7 +55,6 @@ public WaitForConditionOperation(
 
         this.checkFunc = checkFunc;
         this.config = config;
-        this.userExecutor = durableContext.getDurableConfig().getExecutorService();
     }
 
     @Override

sdk/src/test/java/software/amazon/lambda/durable/DurableConfigTest.java

@@ -190,7 +190,7 @@ void testBuilder_MultipleBuilds_CreateIndependentInstances() {
         assertEquals(config1.getDurableExecutionClient(), config2.getDurableExecutionClient());
 
         // ExecutorService should be different instances (each gets its own)
-        assertNotSame(config1.getExecutorService(), config2.getExecutorService());
+        assertSame(config1.getExecutorService(), config2.getExecutorService());
     }
 
     @Test
@@ -210,7 +210,7 @@ void testDefaultConfig_CreatesNewInstancesEachTime() {
         var config2 = DurableConfig.defaultConfig();
 
         assertNotSame(config1, config2);
-        assertNotSame(config1.getExecutorService(), config2.getExecutorService());
+        assertSame(config1.getExecutorService(), config2.getExecutorService());
     }
 
     @Test