improvement: Improve polling mechanism

Author: Alex Wang

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

@@ -21,6 +21,7 @@
 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.JitterStrategy;
 import software.amazon.lambda.durable.serde.JacksonSerDes;
 import software.amazon.lambda.durable.serde.SerDes;
 
@@ -78,6 +79,8 @@ public final class DurableConfig {
     private final ExecutorService executorService;
     private final LoggerConfig loggerConfig;
     private final Duration pollingInterval;
+    private final JitterStrategy pollingJitter;
+    private final double pollingBackoffRate;
     private final Duration checkpointDelay;
 
     private DurableConfig(Builder builder) {
@@ -88,6 +91,8 @@ private DurableConfig(Builder builder) {
         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.pollingJitter = builder.pollingJitter != null ? builder.pollingJitter : JitterStrategy.FULL;
+        this.pollingBackoffRate = builder.pollingBackoffRate != 0.0 ? builder.pollingBackoffRate : 2.0;
         this.checkpointDelay = builder.checkpointDelay != null ? builder.checkpointDelay : Duration.ofSeconds(0);
     }
 
@@ -154,6 +159,24 @@ public Duration getPollingInterval() {
         return pollingInterval;
     }
 
+    /**
+     * Gets the polling backoff rate.
+     *
+     * @return polling backoff rate.
+     */
+    public double getPollingBackoffRate() {
+        return pollingBackoffRate;
+    }
+
+    /**
+     * Gets the polling jitter.
+     *
+     * @return polling jitter.
+     */
+    public JitterStrategy getPollingJitter() {
+        return pollingJitter;
+    }
+
     /**
      * Gets the configured checkpoint delay.
      *
@@ -249,6 +272,8 @@ public static final class Builder {
         private ExecutorService executorService;
         private LoggerConfig loggerConfig;
         private Duration pollingInterval;
+        private JitterStrategy pollingJitter;
+        private double pollingBackoffRate;
         private Duration checkpointDelay;
 
         private Builder() {}
@@ -347,6 +372,31 @@ public Builder withPollingInterval(Duration duration) {
             return this;
         }
 
+        /**
+         * Sets the polling backoff rate.
+         *
+         * @param backoffRate the backoff rate (must be positive)
+         * @return This builder
+         */
+        public Builder withPollingBackoffRate(double backoffRate) {
+            if (backoffRate <= 0) {
+                throw new IllegalArgumentException("backoffRate must be positive");
+            }
+            this.pollingBackoffRate = backoffRate;
+            return this;
+        }
+
+        /**
+         * Sets the polling jitter strategy.
+         *
+         * @param jitter the jitter strategy to use
+         * @return This builder
+         */
+        public Builder withPollingJitter(JitterStrategy jitter) {
+            this.pollingJitter = jitter;
+            return this;
+        }
+
         /**
          * Sets how often the SDK checkpoints updates to backend. If not set, defaults to 0, which disables checkpoint
          * batching.

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

@@ -57,11 +57,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.getPollingInterval(), false);
     }
 
     /** Polls for updates of the specified operation with specified delay */
     CompletableFuture<Operation> pollForUpdate(String operationId, Duration delay) {
+        return pollForUpdate(operationId, delay, true);
+    }
+
+    /** Polls for updates of the specified operation with custom delay and interval behavior. */
+    CompletableFuture<Operation> pollForUpdate(String operationId, Duration delay, boolean isFixedInterval) {
         logger.debug("Polling request received: operation id {}", operationId);
         var future = new CompletableFuture<Operation>();
         synchronized (pollingFutures) {
@@ -70,16 +75,24 @@ CompletableFuture<Operation> pollForUpdate(String operationId, Duration delay) {
                     .computeIfAbsent(operationId, k -> Collections.synchronizedList(new ArrayList<>()))
                     .add(future);
         }
-        pollForUpdateInternal(future, delay);
+        pollForUpdateInternal(future, delay, isFixedInterval, 0);
         return future;
     }
 
-    private CompletableFuture<Void> pollForUpdateInternal(CompletableFuture<Operation> future, Duration delay) {
-        return checkpointApiRequestBatcher.submit(null, delay).thenCompose(v -> {
+    private CompletableFuture<Void> pollForUpdateInternal(
+            CompletableFuture<Operation> future, Duration initialDelay, boolean isFixedInterval, int attempt) {
+        Duration actualDelay = initialDelay;
+        if (!isFixedInterval) {
+            double delayMilli = initialDelay.toMillis();
+            delayMilli = delayMilli * Math.pow(config.getPollingBackoffRate(), attempt);
+            delayMilli = config.getPollingJitter().apply(delayMilli);
+            actualDelay = Duration.ofMillis(Math.round(delayMilli));
+        }
+        return checkpointApiRequestBatcher.submit(null, actualDelay).thenCompose(v -> {
             if (future.isDone()) {
                 return CompletableFuture.completedFuture(null);
             }
-            return pollForUpdateInternal(future, delay);
+            return pollForUpdateInternal(future, initialDelay, isFixedInterval, attempt + 1);
         });
     }
 

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/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

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

@@ -12,13 +12,15 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 
+import java.time.Duration;
 import java.util.concurrent.ExecutorService;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
 import software.amazon.awssdk.services.lambda.LambdaClient;
 import software.amazon.lambda.durable.client.DurableExecutionClient;
 import software.amazon.lambda.durable.client.LambdaDurableFunctionsClient;
+import software.amazon.lambda.durable.retry.JitterStrategy;
 import software.amazon.lambda.durable.serde.JacksonSerDes;
 import software.amazon.lambda.durable.serde.SerDes;
 
@@ -294,4 +296,148 @@ void testAddUserAgentSuffix_ReturnsSameBuilderInstance() {
 
         assertSame(lambdaClientBuilder, result);
     }
+
+    // --- Polling interval tests ---
+
+    @Test
+    void testDefaultConfig_PollingIntervalDefaults() {
+        var config = DurableConfig.defaultConfig();
+
+        assertEquals(Duration.ofMillis(1000), config.getPollingInterval());
+        assertEquals(JitterStrategy.FULL, config.getPollingJitter());
+        assertEquals(2.0, config.getPollingBackoffRate());
+        assertEquals(Duration.ofSeconds(0), config.getCheckpointDelay());
+    }
+
+    @Test
+    void testBuilder_WithCustomPollingInterval() {
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withPollingInterval(Duration.ofMillis(500))
+                .build();
+
+        assertEquals(Duration.ofMillis(500), config.getPollingInterval());
+    }
+
+    @Test
+    void testBuilder_WithCustomPollingBackoffRate() {
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withPollingBackoffRate(3.0)
+                .build();
+
+        assertEquals(3.0, config.getPollingBackoffRate());
+    }
+
+    @Test
+    void testBuilder_WithPollingBackoffRateNotSet_UsesDefault() {
+        var config =
+                DurableConfig.builder().withDurableExecutionClient(mockClient).build();
+
+        assertEquals(2.0, config.getPollingBackoffRate());
+    }
+
+    @Test
+    void testBuilder_WithZeroPollingBackoffRate_ThrowsException() {
+        var builder = DurableConfig.builder();
+
+        var exception = assertThrows(IllegalArgumentException.class, () -> builder.withPollingBackoffRate(0.0));
+
+        assertEquals("backoffRate must be positive", exception.getMessage());
+    }
+
+    @Test
+    void testBuilder_WithNegativePollingBackoffRate_ThrowsException() {
+        var builder = DurableConfig.builder();
+
+        var exception = assertThrows(IllegalArgumentException.class, () -> builder.withPollingBackoffRate(-1.0));
+
+        assertEquals("backoffRate must be positive", exception.getMessage());
+    }
+
+    @Test
+    void testBuilder_WithCustomPollingJitter() {
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withPollingJitter(JitterStrategy.NONE)
+                .build();
+
+        assertEquals(JitterStrategy.NONE, config.getPollingJitter());
+    }
+
+    @Test
+    void testBuilder_WithPollingJitterHalf() {
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withPollingJitter(JitterStrategy.HALF)
+                .build();
+
+        assertEquals(JitterStrategy.HALF, config.getPollingJitter());
+    }
+
+    @Test
+    void testBuilder_WithPollingJitterNull_UsesDefault() {
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withPollingJitter(null)
+                .build();
+
+        assertEquals(JitterStrategy.FULL, config.getPollingJitter());
+    }
+
+    @Test
+    void testBuilder_WithCustomCheckpointDelay() {
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withCheckpointDelay(Duration.ofSeconds(5))
+                .build();
+
+        assertEquals(Duration.ofSeconds(5), config.getCheckpointDelay());
+    }
+
+    @Test
+    void testBuilder_WithAllPollingSettings() {
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withPollingInterval(Duration.ofMillis(200))
+                .withPollingBackoffRate(1.5)
+                .withPollingJitter(JitterStrategy.HALF)
+                .withCheckpointDelay(Duration.ofSeconds(2))
+                .build();
+
+        assertEquals(Duration.ofMillis(200), config.getPollingInterval());
+        assertEquals(1.5, config.getPollingBackoffRate());
+        assertEquals(JitterStrategy.HALF, config.getPollingJitter());
+        assertEquals(Duration.ofSeconds(2), config.getCheckpointDelay());
+    }
+
+    @Test
+    void testBuilder_FluentAPI_PollingMethods() {
+        var builder = DurableConfig.builder();
+
+        assertSame(builder, builder.withPollingInterval(Duration.ofMillis(500)));
+        assertSame(builder, builder.withPollingBackoffRate(2.0));
+        assertSame(builder, builder.withPollingJitter(JitterStrategy.FULL));
+        assertSame(builder, builder.withCheckpointDelay(Duration.ofSeconds(1)));
+    }
+
+    @Test
+    void testBuilder_PollingIntervalNull_UsesDefault() {
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withPollingInterval(null)
+                .build();
+
+        assertEquals(Duration.ofMillis(1000), config.getPollingInterval());
+    }
+
+    @Test
+    void testBuilder_CheckpointDelayNull_UsesDefault() {
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withCheckpointDelay(null)
+                .build();
+
+        assertEquals(Duration.ofSeconds(0), config.getCheckpointDelay());
+    }
 }