optimize pollers

Author: zhongkechen

sdk/src/main/java/com/amazonaws/lambda/durable/DurableConfig.java

@@ -7,6 +7,7 @@
 import com.amazonaws.lambda.durable.logging.LoggerConfig;
 import com.amazonaws.lambda.durable.serde.JacksonSerDes;
 import com.amazonaws.lambda.durable.serde.SerDes;
+import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -67,6 +68,7 @@ public final class DurableConfig {
     private final SerDes serDes;
     private final ExecutorService executorService;
     private final LoggerConfig loggerConfig;
+    private final Duration pollingInterval;
 
     private DurableConfig(Builder builder) {
         this.durableExecutionClient = builder.durableExecutionClient != null
@@ -75,6 +77,7 @@ private DurableConfig(Builder builder) {
         this.serDes = builder.serDes != null ? builder.serDes : new JacksonSerDes();
         this.executorService = builder.executorService != null ? builder.executorService : createDefaultExecutor();
         this.loggerConfig = builder.loggerConfig != null ? builder.loggerConfig : LoggerConfig.defaults();
+        this.pollingInterval = builder.pollingInterval != null ? builder.pollingInterval : Duration.ofSeconds(1);
     }
 
     /**
@@ -131,6 +134,15 @@ public LoggerConfig getLoggerConfig() {
         return loggerConfig;
     }
 
+    /**
+     * Gets the configured polling interval.
+     *
+     * @return polling interval in Duration.
+     */
+    public Duration getPollingInterval() {
+        return pollingInterval;
+    }
+
     /**
      * Creates a default DurableExecutionClient with production LambdaClient. Uses
      * EnvironmentVariableCredentialsProvider and region from AWS_REGION. If AWS_REGION is not set, defaults to
@@ -195,6 +207,7 @@ public static final class Builder {
         private SerDes serDes;
         private ExecutorService executorService;
         private LoggerConfig loggerConfig;
+        public Duration pollingInterval;
 
         private Builder() {}
 
@@ -280,6 +293,17 @@ public Builder withLoggerConfig(LoggerConfig loggerConfig) {
             return this;
         }
 
+        /**
+         * Sets how often the SDK polls updates from backend.
+         *
+         * @param duration the polling interval in Duration
+         * @return This builder
+         */
+        public Builder withPollingInterval(Duration duration) {
+            this.pollingInterval = duration;
+            return this;
+        }
+
         /**
          * Builds the DurableConfig instance.
          *

sdk/src/main/java/com/amazonaws/lambda/durable/DurableExecutor.java

@@ -59,10 +59,7 @@ public static <I, O> DurableExecutionOutput execute(
         }
 
         var executionManager = new ExecutionManager(
-                input.durableExecutionArn(),
-                input.checkpointToken(),
-                input.initialExecutionState(),
-                config.getDurableExecutionClient());
+                input.durableExecutionArn(), input.checkpointToken(), input.initialExecutionState(), config);
 
         logger.debug(
                 "EXECUTION operation found: {}",

sdk/src/main/java/com/amazonaws/lambda/durable/DurableHandler.java

@@ -170,7 +170,7 @@ public final void handleRequest(InputStream inputStream, OutputStream outputStre
      *
      * <p>Customer-facing serialization uses SerDes from DurableConfig.
      *
-     * @return Configured ObjectMapper for DAR backend communication
+     * @return Configured ObjectMapper for durable backend communication
      */
     public static ObjectMapper createObjectMapper() {
         var dateModule = new SimpleModule();

sdk/src/main/java/com/amazonaws/lambda/durable/client/LambdaDurableFunctionsClient.java

@@ -35,7 +35,7 @@ public CheckpointDurableExecutionResponse checkpoint(String arn, String token, L
                 .checkpointToken(token)
                 .updates(updates)
                 .build();
-        logger.debug("Calling DAR backend with {} updates: {}", updates.size(), request);
+        logger.debug("Calling durable backend with {} updates: {}", updates.size(), request);
 
         return lambdaClient.checkpointDurableExecution(request);
     }

sdk/src/main/java/com/amazonaws/lambda/durable/execution/CheckpointBatcher.java

@@ -2,12 +2,17 @@
 // SPDX-License-Identifier: Apache-2.0
 package com.amazonaws.lambda.durable.execution;
 
+import com.amazonaws.lambda.durable.DurableConfig;
 import com.amazonaws.lambda.durable.client.DurableExecutionClient;
+import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
@@ -35,41 +40,57 @@ class CheckpointBatcher {
     private final String durableExecutionArn;
     private final DurableExecutionClient client;
     private final BlockingQueue<CheckpointRequest> queue = new LinkedBlockingQueue<>();
-    private final AtomicBoolean isProcessing = new AtomicBoolean(false);
+    private final AtomicBoolean isRunning = new AtomicBoolean(true);
+    private final Duration pollingInterval;
+    private final Map<String, List<CompletableFuture<Operation>>> pollingFutures = new ConcurrentHashMap<>();
     private String checkpointToken;
 
     record CheckpointRequest(OperationUpdate update, CompletableFuture<Void> completion) {}
 
     CheckpointBatcher(
-            DurableExecutionClient client,
+            DurableConfig config,
             String durableExecutionArn,
             String checkpointToken,
             Consumer<List<Operation>> callback) {
-        this.client = client;
+        this.client = config.getDurableExecutionClient();
         this.durableExecutionArn = durableExecutionArn;
         this.callback = callback;
         this.checkpointToken = checkpointToken;
+        this.pollingInterval = config.getPollingInterval();
+
+        InternalExecutor.INSTANCE.execute(this::processQueue);
     }
 
     CompletableFuture<Void> checkpoint(OperationUpdate update) {
-        logger.debug(
-                "Checkpoint request received: Action {}",
-                update != null ? update.action() : "NULL (Checkpoint request)");
+        logger.debug("Checkpoint request received: Action {}", update.action());
         var future = new CompletableFuture<Void>();
         queue.add(new CheckpointRequest(update, future));
+        return future;
+    }
 
-        if (isProcessing.compareAndSet(false, true)) {
-            InternalExecutor.INSTANCE.execute(this::processQueue);
-        }
-
+    CompletableFuture<Operation> pollForUpdate(String operationId) {
+        var future = new CompletableFuture<Operation>();
+        pollingFutures
+                .computeIfAbsent(operationId, k -> Collections.synchronizedList(new ArrayList<>()))
+                .add(future);
         return future;
     }
 
     void shutdown() {
         var remaining = new ArrayList<CheckpointRequest>();
+        isRunning.set(false);
         queue.drainTo(remaining);
+
+        // fail the checkpoint requests
         remaining.forEach(
                 req -> req.completion().completeExceptionally(new IllegalStateException("CheckpointManager shutdown")));
+
+        // fail the pollers
+        for (var operationId : pollingFutures.keySet()) {
+            pollingFutures
+                    .remove(operationId)
+                    .forEach(f -> f.completeExceptionally(new IllegalStateException("CheckpointManager shutdown")));
+        }
     }
 
     public List<Operation> fetchAllPages(List<Operation> initialOperations, String nextMarker) {
@@ -79,25 +100,32 @@ public List<Operation> fetchAllPages(List<Operation> initialOperations, String n
         }
         while (nextMarker != null && !nextMarker.isEmpty()) {
             var response = client.getExecutionState(durableExecutionArn, checkpointToken, nextMarker);
-            logger.debug("DAR getExecutionState called: {}.", response);
+            logger.debug("Durable API getExecutionState called: {}.", response);
             operations.addAll(response.operations());
             nextMarker = response.nextMarker();
         }
         return operations;
     }
 
     private void processQueue() {
-        try {
-            var batch = collectBatch();
-            if (!batch.isEmpty()) {
+        while (isRunning.get()) {
+            if (queue.isEmpty() && pollingFutures.isEmpty()) {
+                // nothing to process
+                try {
+                    Thread.sleep(pollingInterval.toMillis());
+                } catch (InterruptedException ignored) {
+                }
+            }
+            try {
+                var batch = collectBatch();
                 // Filter out null updates (empty checkpoints for polling)
                 var updates = batch.stream()
                         .map(CheckpointRequest::update)
                         .filter(Objects::nonNull)
                         .toList();
 
                 var response = client.checkpoint(durableExecutionArn, checkpointToken, updates);
-                logger.debug("DAR checkpointDurableExecution called: {}.", response);
+                logger.debug("Durable API checkpointDurableExecution called: {}.", response);
 
                 // Notify callback of completion
                 // TODO: sam local backend returns no new execution state when called with zero
@@ -112,19 +140,21 @@ private void processQueue() {
                     if (!operations.isEmpty()) {
                         callback.accept(operations);
                     }
+
+                    for (var operation : operations) {
+                        var pollers = pollingFutures.remove(operation.id());
+                        if (pollers != null) {
+                            pollers.forEach(poller -> poller.complete(operation));
+                        }
+                    }
                 }
 
+                // checkpoint operation completed
                 batch.forEach(req -> req.completion().complete(null));
-            }
-        } catch (Exception e) {
-            var batch = new ArrayList<CheckpointRequest>();
-            queue.drainTo(batch);
-            batch.forEach(req -> req.completion().completeExceptionally(e));
-        } finally {
-            isProcessing.set(false);
-
-            if (!queue.isEmpty() && isProcessing.compareAndSet(false, true)) {
-                InternalExecutor.INSTANCE.execute(this::processQueue);
+            } catch (Throwable e) {
+                var batch = new ArrayList<CheckpointRequest>();
+                queue.drainTo(batch);
+                batch.forEach(req -> req.completion().completeExceptionally(e));
             }
         }
     }
@@ -134,14 +164,16 @@ private List<CheckpointRequest> collectBatch() {
         var currentSize = 0;
 
         CheckpointRequest req;
-        while ((req = queue.poll()) != null) {
+        while ((req = queue.peek()) != null) {
             var itemSize = estimateSize(req.update());
 
+            // will include the big item in the batch if the batch is empty
             if (currentSize + itemSize > MAX_BATCH_SIZE_BYTES && !batch.isEmpty()) {
-                queue.add(req);
                 break;
             }
 
+            queue.remove();
+
             batch.add(req);
             currentSize += itemSize;
         }