fix negative polling delay (aws#260)

Author: zhongkechen

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

@@ -72,12 +72,14 @@ CompletableFuture<Operation> pollForUpdate(String operationId) {
     }
 
     /**
-     * Polls for updates of the specified operation with specified delay
+     * Polls for updates of the specified operation at the specified time. If the give time is at the past, SDK will
+     * immediately make a polling call.
      *
+     * @param at the time to poll for the update
      * @return a future that completes when the operation is updated
      */
-    CompletableFuture<Operation> pollForUpdate(String operationId, Duration delay) {
-        return pollForUpdate(operationId, PollingStrategies.fixedDelay(delay));
+    CompletableFuture<Operation> pollForUpdate(String operationId, Instant at) {
+        return pollForUpdate(operationId, PollingStrategies.at(at));
     }
 
     /**

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

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 package software.amazon.lambda.durable.execution;
 
-import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -247,8 +247,15 @@ public CompletableFuture<Operation> pollForOperationUpdates(String operationId)
         return checkpointManager.pollForUpdate(operationId);
     }
 
-    public CompletableFuture<Operation> pollForOperationUpdates(String operationId, Duration delay) {
-        return checkpointManager.pollForUpdate(operationId, delay);
+    /**
+     * Pools for operation updates at a specific time
+     *
+     * @param operationId the operation id to poll for updates
+     * @param at the time to poll for updates
+     * @return a completable future that completes with the operation update
+     */
+    public CompletableFuture<Operation> pollForOperationUpdates(String operationId, Instant at) {
+        return checkpointManager.pollForUpdate(operationId, at);
     }
 
     // ===== Utilities =====

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

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 package software.amazon.lambda.durable.operation;
 
-import java.time.Duration;
+import java.time.Instant;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
@@ -341,13 +341,13 @@ protected CompletableFuture<Operation> pollForOperationUpdates() {
     }
 
     /**
-     * Polls the backend for updates to this operation after the specified delay.
+     * Polls the backend for updates to this operation at a specific time.
      *
-     * @param delay the delay before polling
+     * @param at the time to poll for updates
      * @return a future that completes with the updated operation
      */
-    protected CompletableFuture<Operation> pollForOperationUpdates(Duration delay) {
-        return executionManager.pollForOperationUpdates(getOperationId(), delay);
+    protected CompletableFuture<Operation> pollForOperationUpdates(Instant at) {
+        return executionManager.pollForOperationUpdates(getOperationId(), at);
     }
 
     /** Sends an operation update synchronously (blocks until the update is acknowledged). */

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

@@ -2,8 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0
 package software.amazon.lambda.durable.operation;
 
-import java.time.Duration;
-import java.time.Instant;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 import software.amazon.awssdk.services.lambda.model.ErrorObject;
@@ -87,10 +85,10 @@ protected void replay(Operation existing) {
 
     private CompletableFuture<Void> pollReadyAndExecuteStepLogic(Operation existing, int attempt) {
         var nextAttemptInstant = existing.stepDetails().nextAttemptTimestamp();
-        return pollForOperationUpdates(Duration.between(Instant.now(), nextAttemptInstant))
+        return pollForOperationUpdates(nextAttemptInstant)
                 .thenCompose(op -> op.status() == OperationStatus.READY
                         ? CompletableFuture.completedFuture(op)
-                        : pollForOperationUpdates())
+                        : pollForOperationUpdates(nextAttemptInstant))
                 .thenRun(() -> executeStepLogic(attempt));
     }
 

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

@@ -60,21 +60,14 @@ protected void replay(Operation existing) {
     }
 
     private void pollForWaitExpiration() {
-        // Always calculate remaining time from scheduledEndTimestamp if scheduledEndTimestamp exists
-        var remainingWaitTime = duration;
+        var scheduledEndTimestamp = Instant.now().plusMillis(duration.toMillis());
         var existing = getOperation();
         if (existing != null
                 && existing.waitDetails() != null
                 && existing.waitDetails().scheduledEndTimestamp() != null) {
-            remainingWaitTime =
-                    Duration.between(Instant.now(), existing.waitDetails().scheduledEndTimestamp());
-            // If the wait has already elapsed, poll immediately with a minimal positive interval
-            if (remainingWaitTime.isNegative() || remainingWaitTime.isZero()) {
-                remainingWaitTime = Duration.ofMillis(1);
-            }
+            scheduledEndTimestamp = existing.waitDetails().scheduledEndTimestamp();
         }
-        logger.debug("Remaining wait time: {} ms", remainingWaitTime.toMillis());
-        pollForOperationUpdates(remainingWaitTime);
+        pollForOperationUpdates(scheduledEndTimestamp);
     }
 
     @Override

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

@@ -3,6 +3,7 @@
 package software.amazon.lambda.durable.retry;
 
 import java.time.Duration;
+import java.time.Instant;
 import java.util.Objects;
 
 /** Factory class for creating common polling strategies. */
@@ -67,4 +68,22 @@ public static PollingStrategy fixedDelay(Duration interval) {
         }
         return (attempt) -> interval;
     }
+
+    /**
+     * Creates a polling strategy that polls at a specific instant in time.
+     *
+     * @param instant The instant to poll at
+     * @return PollingStrategy that calculates delay until the specified instant
+     */
+    public static PollingStrategy at(Instant instant) {
+        Objects.requireNonNull(instant, "instant must not be null");
+        return (attempt) -> {
+            var duration = Duration.between(Instant.now(), instant);
+            if (duration.isNegative()) {
+                // as soon as possible
+                return Duration.ZERO;
+            }
+            return duration;
+        };
+    }
 }

sdk/src/test/java/software/amazon/lambda/durable/execution/CheckpointManagerTest.java

@@ -7,6 +7,7 @@
 import static org.mockito.Mockito.*;
 
 import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
@@ -276,7 +277,7 @@ void pollForUpdate_withCustomDelay() throws Exception {
                                 .build())
                         .build());
 
-        var future = batcher.pollForUpdate("op-1", Duration.ofMillis(100));
+        var future = batcher.pollForUpdate("op-1", Instant.now().plusMillis(100));
 
         var result = future.get(300, TimeUnit.MILLISECONDS);
 
@@ -410,7 +411,7 @@ void pollForUpdate_withCustomDelay_ignoresBackoffConfig() throws Exception {
         });
 
         // Use explicit delay (fixed interval) — should NOT apply backoff
-        var future = backoffBatcher.pollForUpdate("op-1", Duration.ofMillis(20));
+        var future = backoffBatcher.pollForUpdate("op-1", Instant.now().plusMillis(20));
         var result = future.get(1000, TimeUnit.MILLISECONDS);
 
         assertEquals(operation, result);