make initialState optional in config

Author: zhongkechen

sdk/src/main/java/software/amazon/lambda/durable/config/WaitForConditionConfig.java

@@ -16,10 +16,12 @@
 public class WaitForConditionConfig<T> {
     private final WaitForConditionWaitStrategy<T> waitStrategy;
     private final SerDes serDes;
+    private final T initialState;
 
     private WaitForConditionConfig(Builder<T> builder) {
         this.waitStrategy = builder.waitStrategy;
         this.serDes = builder.serDes;
+        this.initialState = builder.initialState;
     }
 
     /**
@@ -35,6 +37,11 @@ public SerDes serDes() {
         return serDes;
     }
 
+    /** Returns the initial state object, or null if not specified. */
+    public T initialState() {
+        return initialState;
+    }
+
     /**
      * Returns a new builder initialized with the values from this config. Useful internally for injecting default
      * SerDes.
@@ -45,6 +52,7 @@ public Builder<T> toBuilder() {
         var b = new Builder<T>();
         b.waitStrategy = this.waitStrategy;
         b.serDes = this.serDes;
+        b.initialState = this.initialState;
         return b;
     }
 
@@ -61,6 +69,7 @@ public static <T> Builder<T> builder() {
     public static class Builder<T> {
         private WaitForConditionWaitStrategy<T> waitStrategy;
         private SerDes serDes;
+        private T initialState;
 
         private Builder() {}
 
@@ -91,6 +100,17 @@ public Builder<T> serDes(SerDes serDes) {
             return this;
         }
 
+        /**
+         * Sets the initial state for the waitForCondition operation. The initial state will be null if it's not set.
+         *
+         * @param initialState the initial state object to pass to the condition function
+         * @return this builder for method chaining
+         */
+        public Builder<T> initialState(T initialState) {
+            this.initialState = initialState;
+            return this;
+        }
+
         public WaitForConditionConfig<T> build() {
             return new WaitForConditionConfig<>(this);
         }

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

@@ -61,7 +61,7 @@ public WaitForConditionOperation(
 
     @Override
     protected void start() {
-        executeCheckLogic(null, 0);
+        executeCheckLogic(config.initialState(), 0);
     }
 
     @Override
@@ -105,7 +105,7 @@ private void resumeCheckLoop(Operation existing) {
         if (checkpointData != null) {
             currentState = deserializeResult(checkpointData);
         } else {
-            currentState = null;
+            currentState = config.initialState();
         }
         executeCheckLogic(currentState, attempt);
     }

sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java

@@ -343,8 +343,11 @@ void replayStartedWithNullCheckpointDataUsesInitialState() throws Exception {
                 .build();
         when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op);
 
-        var receivedState = new java.util.concurrent.atomic.AtomicReference<Integer>(-1);
-        var config = WaitForConditionConfig.<Integer>builder().serDes(SERDES).build();
+        var receivedState = new java.util.concurrent.atomic.AtomicInteger(-1);
+        var config = WaitForConditionConfig.<Integer>builder()
+                .serDes(SERDES)
+                .initialState(42)
+                .build();
         var operation = createOperation(
                 (state, ctx) -> {
                     receivedState.set(state);
@@ -355,7 +358,7 @@ void replayStartedWithNullCheckpointDataUsesInitialState() throws Exception {
         operation.execute();
 
         Thread.sleep(200);
-        assertNull(receivedState.get(), "Should use initialState when checkpoint data is null");
+        assertEquals(42, receivedState.get(), "Should use initialState when checkpoint data is null");
     }
 
     // ===== resumeCheckLoop checkpoint deserialize exception =====