improvement: Improve polling mechanism

- Refactor jitter
- Add polling strategies class
- Tests

Author: Alex Wang

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/CustomConfigIntegrationTest.java

@@ -6,11 +6,16 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.time.Duration;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.lambda.LambdaClient;
 import software.amazon.lambda.durable.client.LambdaDurableFunctionsClient;
+import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.retry.JitterStrategy;
+import software.amazon.lambda.durable.retry.PollingStrategies;
+import software.amazon.lambda.durable.retry.RetryStrategies;
 import software.amazon.lambda.durable.serde.JacksonSerDes;
 import software.amazon.lambda.durable.serde.SerDes;
 import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
@@ -147,14 +152,14 @@ void testCustomConfig_WithRetry() {
                                 throw new RuntimeException("Simulated failure attempt " + currentAttempt);
                             },
                             StepConfig.builder()
-                                    .retryStrategy(software.amazon.lambda.durable.retry.RetryStrategies.Presets.DEFAULT)
+                                    .retryStrategy(RetryStrategies.Presets.DEFAULT)
                                     .build());
                 },
                 customConfig);
 
         // First run should return PENDING (retry scheduled) - matching existing RetryIntegrationTest pattern
         var result = runner.run("test");
-        assertEquals(software.amazon.lambda.durable.model.ExecutionStatus.PENDING, result.getStatus());
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
         assertEquals(1, attemptCount.get());
 
         // Verify custom SerDes was used during retry operations
@@ -284,4 +289,94 @@ void testCustomDurableExecutionClient_Configuration() {
             lambdaClient.close();
         }
     }
+
+    // --- Polling strategy integration tests ---
+
+    @Test
+    void testCustomPollingStrategy_ExponentialBackoff_WithSteps() {
+        var customConfig = DurableConfig.builder()
+                .withPollingStrategy(
+                        PollingStrategies.exponentialBackoff(Duration.ofMillis(100), 2.0, JitterStrategy.FULL))
+                .build();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> {
+                    var step1 = context.step("step1", String.class, () -> "first");
+                    var step2 = context.step("step2", String.class, () -> "second");
+                    return step1 + "," + step2;
+                },
+                customConfig);
+
+        var result = runner.run("test");
+
+        assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
+        assertEquals("first,second", result.getResult(String.class));
+        assertNotNull(result.getOperation("step1"));
+        assertNotNull(result.getOperation("step2"));
+    }
+
+    @Test
+    void testCustomPollingStrategy_FixedDelay_WithWait() {
+        var customConfig = DurableConfig.builder()
+                .withPollingStrategy(PollingStrategies.fixedDelay(Duration.ofMillis(200)))
+                .build();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> {
+                    var step1 = context.step("before-wait", String.class, () -> "before");
+                    context.wait("my-wait", Duration.ofSeconds(5));
+                    var step2 = context.step("after-wait", String.class, () -> "after");
+                    return step1 + "," + step2;
+                },
+                customConfig);
+
+        // First run suspends at wait
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Advance time and complete
+        var finalResult = runner.runUntilComplete("test");
+        assertEquals(ExecutionStatus.SUCCEEDED, finalResult.getStatus());
+        assertEquals("before,after", finalResult.getResult(String.class));
+    }
+
+    @Test
+    void testCustomPollingStrategy_WithRetry() {
+        var attemptCount = new AtomicInteger(0);
+
+        var customConfig = DurableConfig.builder()
+                .withPollingStrategy(
+                        PollingStrategies.exponentialBackoff(Duration.ofMillis(50), 1.5, JitterStrategy.HALF))
+                .build();
+
+        var runner = LocalDurableTestRunner.create(
+                String.class,
+                (input, context) -> {
+                    return context.step(
+                            "retry-step",
+                            String.class,
+                            () -> {
+                                int attempt = attemptCount.incrementAndGet();
+                                if (attempt < 3) {
+                                    throw new RuntimeException("Fail attempt " + attempt);
+                                }
+                                return "success on attempt " + attempt;
+                            },
+                            StepConfig.builder()
+                                    .retryStrategy(RetryStrategies.Presets.DEFAULT)
+                                    .build());
+                },
+                customConfig);
+
+        // First run fails, schedules retry → PENDING
+        var result = runner.run("test");
+        assertEquals(ExecutionStatus.PENDING, result.getStatus());
+
+        // Run until complete — retries until attempt 3 succeeds
+        var finalResult = runner.runUntilComplete("test");
+        assertEquals(ExecutionStatus.SUCCEEDED, finalResult.getStatus());
+        assertEquals("success on attempt 3", finalResult.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,60 @@
+// 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 */
+        public static final PollingStrategy DEFAULT =
+                exponentialBackoff(Duration.ofMillis(1000), 2.0, JitterStrategy.FULL);
+    }
+
+    /**
+     * 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
+     * @return PollingStrategy implementing exponential backoff with jitter
+     */
+    public static PollingStrategy exponentialBackoff(Duration baseInterval, double backoffRate, JitterStrategy jitter) {
+        Objects.requireNonNull(jitter, "jitter must not be null");
+        Objects.requireNonNull(baseInterval, "base 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");
+        }
+
+        return (attempt) -> {
+            double delayMs = baseInterval.toMillis() * Math.pow(backoffRate, attempt);
+            delayMs = jitter.apply(delayMs);
+            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