improvement: improve polling mechanism

- Refactor jitter
- Add jitter, polling interval, backoff rate and max interval
- unit tests
- example

Author: Alex Wang

examples/src/main/java/software/amazon/lambda/durable/examples/CustomPollingExample.java

@@ -0,0 +1,62 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples;
+
+import java.time.Duration;
+import software.amazon.lambda.durable.DurableConfig;
+import software.amazon.lambda.durable.DurableContext;
+import software.amazon.lambda.durable.DurableHandler;
+import software.amazon.lambda.durable.InvokeConfig;
+import software.amazon.lambda.durable.retry.JitterStrategy;
+import software.amazon.lambda.durable.retry.PollingStrategies;
+
+/**
+ * Example demonstrating custom polling strategy configuration.
+ *
+ * <p>The polling strategy controls how the SDK polls for async operation results. By default, the SDK uses exponential
+ * backoff (1s base, 2x rate, full jitter). This example shows how to customize the polling behavior.
+ *
+ * <p>This example configures:
+ *
+ * <ul>
+ *   <li>Exponential backoff with 500ms base interval
+ *   <li>1.5x backoff rate for gentler growth
+ *   <li>Half jitter to balance between consistency and thundering herd avoidance
+ * </ul>
+ */
+public class CustomPollingExample extends DurableHandler<GreetingRequest, String> {
+
+    @Override
+    protected DurableConfig createConfiguration() {
+        return DurableConfig.builder()
+                .withPollingStrategy(PollingStrategies.exponentialBackoff(
+                        Duration.ofMillis(500), 1.5, JitterStrategy.HALF, Duration.ofSeconds(5)))
+                .build();
+    }
+
+    @Override
+    public String handleRequest(GreetingRequest input, DurableContext context) {
+        context.getLogger().info("Starting workflow with input: {}", input);
+
+        // Step 1: lower case the input
+        var lowered = context.stepAsync("validate", String.class, () -> {
+            try {
+                // prevent the execution from suspension
+                Thread.sleep(5000);
+            } catch (InterruptedException e) {
+                e.printStackTrace();
+            }
+            return input.getName().toLowerCase();
+        });
+
+        // Step 2: Invoke async
+        var future = context.invokeAsync(
+                "call-greeting",
+                "simple-step-example" + input.getName() + ":$LATEST",
+                input,
+                String.class,
+                InvokeConfig.builder().build());
+
+        return future.get() + lowered.get();
+    }
+}

examples/src/test/java/software/amazon/lambda/durable/examples/CustomPollingExampleTest.java

@@ -0,0 +1,31 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.examples;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
+
+class CustomPollingExampleTest {
+
+    @Test
+    void testCustomPollingExample() {
+        var handler = new CustomPollingExample();
+        var runner = LocalDurableTestRunner.create(GreetingRequest.class, handler);
+
+        // First run: executes validate step, then pending at wait
+        var input = new GreetingRequest("world");
+        var output1 = runner.run(input);
+
+        assertEquals(ExecutionStatus.PENDING, output1.getStatus());
+
+        // Second run:
+        runner.completeChainedInvoke("call-greeting", "\"hello\"");
+        var output2 = runner.run(input);
+
+        assertEquals(ExecutionStatus.SUCCEEDED, output2.getStatus());
+        assertEquals("helloworld", output2.getResult(String.class));
+    }
+}

sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java

@@ -155,6 +155,9 @@ public TestResult<O> run(I input) {
                     .withDurableExecutionClient(storage)
                     .withSerDes(customerConfig.getSerDes())
                     .withExecutorService(customerConfig.getExecutorService())
+                    .withPollingStrategy(customerConfig.getPollingStrategy())
+                    .withCheckpointDelay(customerConfig.getCheckpointDelay())
+                    .withLoggerConfig(customerConfig.getLoggerConfig())
                     .build();
         } else {
             // Fallback to default config with in-memory client

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

@@ -21,6 +21,8 @@
 import software.amazon.lambda.durable.client.DurableExecutionClient;
 import software.amazon.lambda.durable.client.LambdaDurableFunctionsClient;
 import software.amazon.lambda.durable.logging.LoggerConfig;
+import software.amazon.lambda.durable.retry.PollingStrategies;
+import software.amazon.lambda.durable.retry.PollingStrategy;
 import software.amazon.lambda.durable.serde.JacksonSerDes;
 import software.amazon.lambda.durable.serde.SerDes;
 
@@ -77,7 +79,7 @@ public final class DurableConfig {
     private final SerDes serDes;
     private final ExecutorService executorService;
     private final LoggerConfig loggerConfig;
-    private final Duration pollingInterval;
+    private final PollingStrategy pollingStrategy;
     private final Duration checkpointDelay;
 
     private DurableConfig(Builder builder) {
@@ -87,7 +89,8 @@ 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.ofMillis(1000);
+        this.pollingStrategy =
+                builder.pollingStrategy != null ? builder.pollingStrategy : PollingStrategies.Presets.DEFAULT;
         this.checkpointDelay = builder.checkpointDelay != null ? builder.checkpointDelay : Duration.ofSeconds(0);
     }
 
@@ -146,12 +149,12 @@ public LoggerConfig getLoggerConfig() {
     }
 
     /**
-     * Gets the configured polling interval.
+     * Gets the polling strategy.
      *
-     * @return polling interval in Duration.
+     * @return PollingStrategy instance (never null)
      */
-    public Duration getPollingInterval() {
-        return pollingInterval;
+    public PollingStrategy getPollingStrategy() {
+        return pollingStrategy;
     }
 
     /**
@@ -248,7 +251,7 @@ public static final class Builder {
         private SerDes serDes;
         private ExecutorService executorService;
         private LoggerConfig loggerConfig;
-        private Duration pollingInterval;
+        private PollingStrategy pollingStrategy;
         private Duration checkpointDelay;
 
         private Builder() {}
@@ -336,14 +339,14 @@ public Builder withLoggerConfig(LoggerConfig loggerConfig) {
         }
 
         /**
-         * Sets how often the SDK polls updates from backend.
+         * Sets the polling strategy. If not set, defaults to 1 second with full jitter and 2x backoff.
          *
-         * @param duration the polling interval in Duration
+         * @param pollingStrategy Custom PollingStrategy instance
          * @return This builder
          */
-        public Builder withPollingInterval(Duration duration) {
+        public Builder withPollingStrategy(PollingStrategy pollingStrategy) {
             // No validation - polling intervals can be less than 1 second (e.g., 200ms with backoff)
-            this.pollingInterval = duration;
+            this.pollingStrategy = pollingStrategy;
             return this;
         }
 

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

@@ -17,6 +17,8 @@
 import software.amazon.awssdk.services.lambda.model.Operation;
 import software.amazon.awssdk.services.lambda.model.OperationUpdate;
 import software.amazon.lambda.durable.DurableConfig;
+import software.amazon.lambda.durable.retry.PollingStrategies;
+import software.amazon.lambda.durable.retry.PollingStrategy;
 
 /**
  * Package-private checkpoint manager for batching and queueing checkpoint API calls.
@@ -57,11 +59,16 @@ CompletableFuture<Void> checkpoint(OperationUpdate update) {
 
     /** Polls for updates of the specified operation with preconfigured intervals */
     CompletableFuture<Operation> pollForUpdate(String operationId) {
-        return pollForUpdate(operationId, config.getPollingInterval());
+        return pollForUpdate(operationId, config.getPollingStrategy());
     }
 
     /** Polls for updates of the specified operation with specified delay */
     CompletableFuture<Operation> pollForUpdate(String operationId, Duration delay) {
+        return pollForUpdate(operationId, PollingStrategies.fixedDelay(delay));
+    }
+
+    /** Polls for updates of the specified operation with specified polling strategy */
+    CompletableFuture<Operation> pollForUpdate(String operationId, PollingStrategy pollingStrategy) {
         logger.debug("Polling request received: operation id {}", operationId);
         var future = new CompletableFuture<Operation>();
         synchronized (pollingFutures) {
@@ -70,17 +77,20 @@ CompletableFuture<Operation> pollForUpdate(String operationId, Duration delay) {
                     .computeIfAbsent(operationId, k -> Collections.synchronizedList(new ArrayList<>()))
                     .add(future);
         }
-        pollForUpdateInternal(future, delay);
+        pollForUpdateInternal(future, 0, pollingStrategy);
         return future;
     }
 
-    private CompletableFuture<Void> pollForUpdateInternal(CompletableFuture<Operation> future, Duration delay) {
-        return checkpointApiRequestBatcher.submit(null, delay).thenCompose(v -> {
-            if (future.isDone()) {
-                return CompletableFuture.completedFuture(null);
-            }
-            return pollForUpdateInternal(future, delay);
-        });
+    private CompletableFuture<Void> pollForUpdateInternal(
+            CompletableFuture<Operation> future, int attempt, PollingStrategy pollingStrategy) {
+        return checkpointApiRequestBatcher
+                .submit(null, pollingStrategy.computeDelay(attempt))
+                .thenCompose(v -> {
+                    if (future.isDone()) {
+                        return CompletableFuture.completedFuture(null);
+                    }
+                    return pollForUpdateInternal(future, attempt + 1, pollingStrategy);
+                });
     }
 
     /** Cancels all polling futures and waits for all pending checkpoint requests to complete */

sdk/src/main/java/software/amazon/lambda/durable/retry/JitterStrategy.java

@@ -13,17 +13,32 @@ public enum JitterStrategy {
     /**
      * No jitter - use exact calculated delay. This provides predictable timing but may cause thundering herd issues.
      */
-    NONE,
-
+    NONE {
+        @Override
+        public double apply(double baseDelay) {
+            return baseDelay;
+        }
+    },
     /**
      * Full jitter - random delay between 0 and calculated delay. This provides maximum spread but may result in very
      * short delays.
      */
-    FULL,
-
+    FULL {
+        @Override
+        public double apply(double baseDelay) {
+            return Math.random() * baseDelay;
+        }
+    },
     /**
      * Half jitter - random delay between 50% and 100% of calculated delay. This provides good spread while maintaining
      * reasonable minimum delays.
      */
-    HALF
+    HALF {
+        @Override
+        public double apply(double baseDelay) {
+            return baseDelay / 2 + Math.random() * (baseDelay / 2);
+        }
+    };
+
+    public abstract double apply(double baseDelay);
 }

sdk/src/main/java/software/amazon/lambda/durable/retry/PollingStrategies.java

@@ -0,0 +1,70 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.retry;
+
+import java.time.Duration;
+import java.util.Objects;
+
+/** Factory class for creating common polling strategies. */
+public class PollingStrategies {
+
+    /** Preset polling strategies for common use cases. */
+    public static class Presets {
+
+        /**
+         * Default polling strategy: - Base interval: 1 second - Backoff rate: 2x - Jitter: FULL - Max interval 10
+         * second
+         */
+        public static final PollingStrategy DEFAULT =
+                exponentialBackoff(Duration.ofMillis(1000), 2.0, JitterStrategy.FULL, Duration.ofSeconds(10));
+    }
+
+    /**
+     * Creates an exponential backoff polling strategy.
+     *
+     * <p>The delay calculation follows the formula: delay = jitter(baseInterval × backoffRate^attempt)
+     *
+     * @param baseInterval Base delay before first poll
+     * @param backoffRate Multiplier for exponential backoff (must be positive)
+     * @param jitter Jitter strategy to apply to delays
+     * @param maxInterval Maximum delay between polls
+     * @return PollingStrategy implementing exponential backoff with jitter
+     */
+    public static PollingStrategy exponentialBackoff(
+            Duration baseInterval, double backoffRate, JitterStrategy jitter, Duration maxInterval) {
+        Objects.requireNonNull(jitter, "jitter must not be null");
+        Objects.requireNonNull(baseInterval, "base interval must not be null");
+        Objects.requireNonNull(maxInterval, "max interval must not be null");
+        if (backoffRate <= 0) {
+            throw new IllegalArgumentException("backoffRate must be positive");
+        }
+
+        if (baseInterval.isNegative() || baseInterval.isZero()) {
+            throw new IllegalArgumentException("baseInterval must be positive");
+        }
+
+        if (maxInterval.isNegative() || maxInterval.isZero()) {
+            throw new IllegalArgumentException("maxInterval must be positive");
+        }
+
+        return (attempt) -> {
+            double delayMs = baseInterval.toMillis() * Math.pow(backoffRate, attempt);
+            delayMs = Math.min(jitter.apply(delayMs), maxInterval.toMillis());
+            return Duration.ofMillis(Math.round(delayMs));
+        };
+    }
+
+    /**
+     * Creates a fixed-delay polling strategy that uses the same interval for every attempt.
+     *
+     * @param interval Fixed delay between polls
+     * @return PollingStrategy with fixed delay
+     */
+    public static PollingStrategy fixedDelay(Duration interval) {
+        Objects.requireNonNull(interval, "interval must not be null");
+        if (interval.isNegative() || interval.isZero()) {
+            throw new IllegalArgumentException("interval must be positive");
+        }
+        return (attempt) -> interval;
+    }
+}

sdk/src/main/java/software/amazon/lambda/durable/retry/PollingStrategy.java

@@ -0,0 +1,18 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.retry;
+
+import java.time.Duration;
+
+/** Functional interface for computing polling delays between attempts. */
+@FunctionalInterface
+public interface PollingStrategy {
+
+    /**
+     * Computes the delay before the next polling attempt.
+     *
+     * @param attempt The current attempt number (0-based)
+     * @return Duration to wait before the next poll
+     */
+    Duration computeDelay(int attempt);
+}

sdk/src/main/java/software/amazon/lambda/durable/retry/RetryStrategies.java

@@ -69,12 +69,7 @@ public static RetryStrategy exponentialBackoff(
             double baseDelay = Math.min(initialDelaySeconds * Math.pow(backoffRate, attemptNumber), maxDelaySeconds);
 
             // Apply jitter
-            double delayWithJitter =
-                    switch (jitter) {
-                        case NONE -> baseDelay;
-                        case FULL -> Math.random() * baseDelay;
-                        case HALF -> baseDelay / 2 + Math.random() * (baseDelay / 2);
-                    };
+            double delayWithJitter = jitter.apply(baseDelay);
 
             // Round to nearest second, minimum 1
             // Same rounding logic as TS SDK: https://tinyurl.com/4ntxsefu