fix(testing-sdk): fix race condition where concurrent updates can cause duplicate invocations

Author: anthonyting

packages/aws-durable-execution-sdk-js-examples/src/examples/force-checkpointing/multiple-wait/force-checkpointing-multiple-wait.ts

@@ -15,6 +15,7 @@ export const handler = withDurableExecution(
     const results = await ctx.parallel([
       // Branch 1: Long-running operation that blocks termination
       async (branchCtx: DurableContext) => {
+        await branchCtx.wait("wait-1-1", { seconds: 1 });
         return await branchCtx.step("long-running-step", async () => {
           await new Promise((resolve) => setTimeout(resolve, 10000));
           return "long-complete";

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/handlers/__tests__/callbacks.test.ts

@@ -14,6 +14,7 @@ import { CheckpointManager } from "../../storage/checkpoint-manager";
 import {
   createExecutionId,
   createCallbackId,
+  createInvocationId,
 } from "../../utils/tagged-strings";
 
 describe("callbacks handlers", () => {
@@ -28,6 +29,7 @@ describe("callbacks handlers", () => {
     executionManager.startExecution({
       executionId: mockExecutionId,
       payload: '{"test": "data"}',
+      invocationId: createInvocationId(),
     });
 
     const storage = executionManager.getCheckpointsByExecution(mockExecutionId);

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/handlers/__tests__/checkpoint-handlers.test.ts

@@ -12,7 +12,10 @@ import {
   processCheckpointDurableExecution,
 } from "../checkpoint-handlers";
 import { ExecutionManager } from "../../storage/execution-manager";
-import { createExecutionId } from "../../utils/tagged-strings";
+import {
+  createExecutionId,
+  createInvocationId,
+} from "../../utils/tagged-strings";
 import { encodeCheckpointToken } from "../../utils/checkpoint-token";
 
 // Mock only external dependencies we can't control
@@ -28,6 +31,7 @@ describe("checkpoint handlers", () => {
     const invocationResult = executionManager.startExecution({
       executionId,
       payload: '{"test": "data"}',
+      invocationId: createInvocationId(),
     });
 
     const storage = executionManager.getCheckpointsByExecution(executionId);

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/handlers/__tests__/execution-handlers.test.ts

@@ -26,17 +26,25 @@ describe("execution handlers", () => {
       const startExecutionSpy = jest.spyOn(executionManager, "startExecution");
 
       const payload = '{"test": "execution data"}';
-      const result = processStartDurableExecution(payload, executionManager);
+      const invocationId = "mock-invocation-id";
+      const result = processStartDurableExecution(
+        {
+          payload,
+          invocationId: createInvocationId(invocationId),
+        },
+        executionManager,
+      );
 
       expect(startExecutionSpy).toHaveBeenCalledWith({
         payload,
+        invocationId: invocationId,
         executionId: expect.any(String),
       });
 
       expect(result).toEqual({
         checkpointToken: expect.any(String),
         executionId: expect.any(String),
-        invocationId: expect.any(String),
+        invocationId: invocationId,
         operationEvents: expect.any(Array),
       });
     });
@@ -47,6 +55,7 @@ describe("execution handlers", () => {
       // First create an execution
       const payload = '{"test": "data"}';
       executionManager.startExecution({
+        invocationId: createInvocationId(),
         executionId: createExecutionId("test-execution-id"),
         payload,
       });
@@ -56,14 +65,20 @@ describe("execution handlers", () => {
         "startInvocation",
       );
 
+      const executionId = createExecutionId("test-execution-id");
+      const invocationId = createInvocationId("test-invocation-id");
       const result = processStartInvocation(
-        "test-execution-id",
+        {
+          executionId,
+          invocationId,
+        },
         executionManager,
       );
 
-      expect(startInvocationSpy).toHaveBeenCalledWith(
-        createExecutionId("test-execution-id"),
-      );
+      expect(startInvocationSpy).toHaveBeenCalledWith({
+        executionId,
+        invocationId,
+      });
       expect(result).toEqual({
         checkpointToken: expect.any(String),
         executionId: createExecutionId("test-execution-id"),
@@ -79,13 +94,23 @@ describe("execution handlers", () => {
           throw new Error("Execution not found");
         });
 
+      const executionId = createExecutionId("non-existent-execution");
+      const invocationId = createInvocationId("test-invocation-id");
+
       expect(() =>
-        processStartInvocation("non-existent-execution", executionManager),
+        processStartInvocation(
+          {
+            executionId,
+            invocationId,
+          },
+          executionManager,
+        ),
       ).toThrow("Execution not found");
 
-      expect(startInvocationSpy).toHaveBeenCalledWith(
-        createExecutionId("non-existent-execution"),
-      );
+      expect(startInvocationSpy).toHaveBeenCalledWith({
+        executionId,
+        invocationId,
+      });
     });
   });
 
@@ -94,7 +119,11 @@ describe("execution handlers", () => {
       // First create an execution and invocation
       const executionId = createExecutionId("test-execution-id");
       const invocationId = createInvocationId("test-invocation-id");
-      executionManager.startExecution({ executionId, payload: "{}" });
+      executionManager.startExecution({
+        executionId,
+        payload: "{}",
+        invocationId: createInvocationId(),
+      });
 
       const completeInvocationSpy = jest.spyOn(
         executionManager,

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/handlers/__tests__/state-handlers.test.ts

@@ -1,7 +1,10 @@
 import { OperationType, OperationAction } from "@aws-sdk/client-lambda";
 import { processGetDurableExecutionState } from "../state-handlers";
 import { ExecutionManager } from "../../storage/execution-manager";
-import { createExecutionId } from "../../utils/tagged-strings";
+import {
+  createExecutionId,
+  createInvocationId,
+} from "../../utils/tagged-strings";
 
 describe("state handlers", () => {
   let executionManager: ExecutionManager;
@@ -21,6 +24,7 @@ describe("state handlers", () => {
       executionManager.startExecution({
         executionId,
         payload: '{"test": "data"}',
+        invocationId: createInvocationId(),
       });
 
       const storage = executionManager.getCheckpointsByExecution(executionId);

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/handlers/execution-handlers.ts

@@ -8,17 +8,21 @@ import {
   ExecutionManager,
   InvocationResult,
 } from "../storage/execution-manager";
+import {
+  StartDurableExecutionRequest,
+  StartInvocationRequest,
+} from "../worker-api/worker-api-request";
 
 /**
  * Starts a durable execution. Returns the data needed for the handler invocation event.
  */
 export function processStartDurableExecution(
-  payload: string | undefined,
+  params: StartDurableExecutionRequest,
   executionManager: ExecutionManager,
 ): InvocationResult {
   return executionManager.startExecution({
-    payload,
     executionId: createExecutionId(),
+    ...params,
   });
 }
 
@@ -27,10 +31,10 @@ export function processStartDurableExecution(
  * in-progress execution.
  */
 export function processStartInvocation(
-  executionIdParam: string,
+  params: StartInvocationRequest,
   executionManager: ExecutionManager,
 ): InvocationResult {
-  return executionManager.startInvocation(createExecutionId(executionIdParam));
+  return executionManager.startInvocation(params);
 }
 
 export function processCompleteInvocation(

packages/aws-durable-execution-sdk-js-testing/src/checkpoint-server/storage/__tests__/execution-manager.test.ts

@@ -115,9 +115,11 @@ describe("execution-manager", () => {
       it("should create a new execution with the provided parameters", () => {
         // Test data
         const executionId = createExecutionId("test-execution-id");
+        const invocationId = createInvocationId("test-invocation-id");
         const params: StartExecutionParams = {
           executionId,
           payload: '{"key":"value"}',
+          invocationId,
         };
 
         // Expected mock operation with properly typed id
@@ -144,9 +146,9 @@ describe("execution-manager", () => {
 
         // Verify results
         expect(result).toEqual({
-          checkpointToken: `encoded-{"executionId":"test-execution-id","token":"mocked-uuid","invocationId":"mocked-uuid"}`,
+          checkpointToken: `encoded-{"executionId":"${executionId}","token":"mocked-uuid","invocationId":"${invocationId}"}`,
           executionId,
-          invocationId: "mocked-uuid",
+          invocationId,
           operationEvents: [
             {
               operation: mockInitialOperation,
@@ -167,6 +169,7 @@ describe("execution-manager", () => {
         const executionId = createExecutionId("test-execution-id");
         const params: StartExecutionParams = {
           executionId,
+          invocationId: createInvocationId("test-invocation-id"),
         };
 
         const initializeSpy = jest.spyOn(
@@ -185,38 +188,48 @@ describe("execution-manager", () => {
       it("should throw error if execution ID doesn't exist", () => {
         const nonExistentId = createExecutionId("non-existent");
 
-        expect(() => executionManager.startInvocation(nonExistentId)).toThrow(
+        expect(() =>
+          executionManager.startInvocation({
+            executionId: nonExistentId,
+            invocationId: createInvocationId("test-invocation-id"),
+          }),
+        ).toThrow(
           "Could not start invocation for invalid execution non-existent",
         );
       });
 
       it("should start a new invocation for an existing execution", () => {
         // First create an execution
         const executionId = createExecutionId("test-execution-id");
-        executionManager.startExecution({ executionId });
+        const invocationId = createInvocationId("test-invocation-id");
+        executionManager.startExecution({ executionId, invocationId });
 
         // Reset mocks to track new calls
         jest.clearAllMocks();
-        (randomUUID as jest.Mock).mockReturnValue("new-invocation-uuid");
 
+        const newInvocationId = createInvocationId("new-invocation-id");
         // Start a new invocation
-        const result = executionManager.startInvocation(executionId);
+        const result = executionManager.startInvocation({
+          executionId,
+          invocationId: newInvocationId,
+        });
 
         expect(result).toBeDefined();
         expect(result.executionId).toBe(executionId);
-        expect(result.invocationId).toBe("new-invocation-uuid");
+        expect(result.invocationId).toBe(newInvocationId);
         expect(result.operationEvents).toBeInstanceOf(Array);
         expect(encodeCheckpointToken).toHaveBeenCalledWith({
           executionId,
-          token: "new-invocation-uuid",
-          invocationId: "new-invocation-uuid",
+          token: expect.any(String),
+          invocationId: newInvocationId,
         });
       });
 
       it("should include all operations from the checkpoint storage", () => {
         // Create an execution with a mock operation
         const executionId = createExecutionId("test-execution-id");
-        executionManager.startExecution({ executionId });
+        const invocationId = createInvocationId("test-invocation-id");
+        executionManager.startExecution({ executionId, invocationId });
 
         // Get the storage and add some operations to it
         const storage = executionManager.getCheckpointsByExecution(executionId);
@@ -247,7 +260,11 @@ describe("execution-manager", () => {
         });
 
         // Start a new invocation
-        const result = executionManager.startInvocation(executionId);
+        const newInvocationId = createInvocationId("new-invocation-id");
+        const result = executionManager.startInvocation({
+          executionId,
+          invocationId: newInvocationId,
+        });
 
         // Check that we got all operations
         expect(result.operationEvents).toHaveLength(2);
@@ -264,13 +281,20 @@ describe("execution-manager", () => {
       it("should throw error if execution is completed already", () => {
         // First create an execution
         const executionId = createExecutionId("test-execution-id");
-        executionManager.startExecution({ executionId });
+        const invocationId = createInvocationId("test-invocation-id");
+        executionManager.startExecution({ executionId, invocationId });
 
         jest
           .spyOn(CheckpointManager.prototype, "isExecutionCompleted")
           .mockReturnValue(true);
 
-        expect(() => executionManager.startInvocation(executionId)).toThrow(
+        const newInvocationId = createInvocationId("new-invocation-id");
+        expect(() =>
+          executionManager.startInvocation({
+            executionId,
+            invocationId: newInvocationId,
+          }),
+        ).toThrow(
           `Could not start invocation for completed execution ${executionId}`,
         );
       });
@@ -286,7 +310,8 @@ describe("execution-manager", () => {
 
       it("should return the checkpoint storage for an existing execution ID", () => {
         const executionId = createExecutionId("test-execution-id");
-        executionManager.startExecution({ executionId });
+        const invocationId = createInvocationId("test-invocation-id");
+        executionManager.startExecution({ executionId, invocationId });
 
         const storage = executionManager.getCheckpointsByExecution(executionId);
 
@@ -321,7 +346,8 @@ describe("execution-manager", () => {
       it("should return storage and token data for a valid token", () => {
         // Create an execution
         const executionId = createExecutionId("test-execution-id");
-        executionManager.startExecution({ executionId });
+        const invocationId = createInvocationId("test-invocation-id");
+        executionManager.startExecution({ executionId, invocationId });
 
         // Create token data that references the execution
         const tokenData: CheckpointTokenData = {
@@ -378,7 +404,8 @@ describe("execution-manager", () => {
       it("should return CheckpointManager for valid callback ID with existing execution", () => {
         // Create an execution first
         const executionId = createExecutionId("test-execution-id");
-        executionManager.startExecution({ executionId });
+        const invocationId = createInvocationId("test-invocation-id");
+        executionManager.startExecution({ executionId, invocationId });
 
         const callbackId = createCallbackId("valid-callback-id");
 
@@ -398,7 +425,8 @@ describe("execution-manager", () => {
       it("should return the same CheckpointManager instance as getCheckpointsByExecution", () => {
         // Create an execution first
         const executionId = createExecutionId("test-execution-id");
-        executionManager.startExecution({ executionId });
+        const invocationId = createInvocationId("test-invocation-id");
+        executionManager.startExecution({ executionId, invocationId });
 
         const callbackId = createCallbackId("valid-callback-id");
 
@@ -451,7 +479,7 @@ describe("execution-manager", () => {
         // Create an execution first
         const executionId = createExecutionId("test-execution-id");
         const invocationId = createInvocationId("test-invocation-id");
-        executionManager.startExecution({ executionId });
+        executionManager.startExecution({ executionId, invocationId });
 
         const storage = executionManager.getCheckpointsByExecution(executionId);
         expect(storage).toBeDefined();
@@ -517,7 +545,7 @@ describe("execution-manager", () => {
         // Create an execution first
         const executionId = createExecutionId("test-execution-id");
         const invocationId = createInvocationId("test-invocation-id");
-        executionManager.startExecution({ executionId });
+        executionManager.startExecution({ executionId, invocationId });
 
         const storage = executionManager.getCheckpointsByExecution(executionId);
         expect(storage).toBeDefined();
@@ -590,7 +618,7 @@ describe("execution-manager", () => {
         // Create an execution first
         const executionId = createExecutionId("test-execution-id");
         const invocationId = createInvocationId("test-invocation-id");
-        executionManager.startExecution({ executionId });
+        executionManager.startExecution({ executionId, invocationId });
 
         const storage = executionManager.getCheckpointsByExecution(executionId);
         expect(storage).toBeDefined();