feat(sdk): add checkpoint queue completion mechanism (#335)

Implement waitForQueueCompletion to ensure all pending checkpoint
operations complete before function termination, preventing data loss.
Add clearQueue method for graceful termination handling.

- Add waitForQueueCompletion() method to CheckpointManager
- Add clearQueue() public method to CheckpointManager
- Integrate queue completion in withDurableExecution with 10s timeout
- Clear queue on termination without waiting
- Update context validation tests to reflect new behavior
- Add unit tests for queue completion functionality

Author: ParidelPooya

packages/aws-durable-execution-sdk-js-examples/src/examples/context-validation/parent-context-in-child/parent-context-in-child.test.ts

@@ -22,8 +22,12 @@ createTests({
 
       const operations = execution.getOperations();
 
-      // Based on actual behavior: no operations are created when validation fails immediately
-      expect(operations).toHaveLength(0);
+      // The child-context operation is now properly persisted before error occurs,
+      // thanks to waitForQueueCompletion. This ensures the checkpoint is saved
+      // regardless of the events that happens after that
+      expect(operations).toHaveLength(1);
+      expect(operations[0].getName()).toBe("child-context");
+      expect(operations[0].getStatus()).toBe("STARTED");
     });
   },
 });

packages/aws-durable-execution-sdk-js-examples/src/examples/context-validation/parent-context-in-wait-condition/parent-context-in-wait-condition.test.ts

@@ -26,25 +26,22 @@ createTests({
 
       const operations = execution.getOperations();
 
-      // Note: runInChildContext doesn't await the checkpoint for STARTING the context.
-      // If termination happens very fast (like in this case where validation fails immediately),
-      // the child-context operation may not be persisted to the checkpoint at all.
-      // This is expected behavior - the validation catches the error before any operations complete.
+      // With waitForQueueCompletion, operations that start are now properly persisted
+      // even when validation fails later, which is the correct behavior
 
-      // The child-context operation may or may not exist depending on checkpoint timing
+      // The child-context operation should exist and be in STARTED state
       const childContextOp = operations.find(
         (op: any) => op.getName() === "child-context",
       );
-      // If it exists, it should be in STARTED state
-      if (childContextOp) {
-        expect(childContextOp.getStatus()).toBe(OperationStatus.STARTED);
-      }
+      expect(childContextOp).toBeDefined();
+      expect(childContextOp!.getStatus()).toBe(OperationStatus.STARTED);
 
-      // The wrong wait condition should NOT exist - validation prevented it
+      // The wrong wait condition should exist since it started before validation failed
       const wrongWaitOp = operations.find(
         (op: any) => op.getName() === "wrong-wait-condition",
       );
-      expect(wrongWaitOp).toBeUndefined();
+      expect(wrongWaitOp).toBeDefined();
+      expect(wrongWaitOp!.getStatus()).toBe(OperationStatus.STARTED);
 
       // The nested wrong step should NOT exist anywhere
       // 1. Should not exist at root level
@@ -66,8 +63,8 @@ createTests({
       );
       expect(shouldNotStartWait).toBeUndefined();
 
-      // Should have at most the child-context operation (0 or 1 operations)
-      expect(operations.length).toBeLessThanOrEqual(1);
+      // Should have exactly 2 operations: child-context and wrong-wait-condition
+      expect(operations.length).toBe(2);
     });
   },
 });

packages/aws-durable-execution-sdk-js/src/testing/create-test-durable-context.ts

@@ -108,6 +108,7 @@ export function createTestDurableContext(options?: {
     force: jest.fn().mockResolvedValue(undefined),
     setTerminating: jest.fn(),
     hasPendingAncestorCompletion: jest.fn().mockReturnValue(false),
+    waitForQueueCompletion: jest.fn().mockResolvedValue(undefined),
   };
 
   const mockDurableExecution = {

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

@@ -7,6 +7,7 @@ export interface CheckpointFunction extends Checkpoint {
   force(): Promise<void>;
   setTerminating(): void;
   hasPendingAncestorCompletion(stepId: string): boolean;
+  waitForQueueCompletion(): Promise<void>;
 }
 
 export const createMockCheckpoint = (
@@ -18,12 +19,13 @@ export const createMockCheckpoint = (
   const mockFn = jest.fn(
     mockImplementation || jest.fn().mockResolvedValue(undefined),
   );
-  
+
   const mockCheckpoint = Object.assign(mockFn, {
     checkpoint: mockFn, // Same function so calls are tracked together
     force: jest.fn().mockResolvedValue(undefined),
     setTerminating: jest.fn(),
     hasPendingAncestorCompletion: jest.fn().mockReturnValue(false),
+    waitForQueueCompletion: jest.fn().mockResolvedValue(undefined),
   }) as jest.MockedFunction<CheckpointFunction>;
 
   return mockCheckpoint;

packages/aws-durable-execution-sdk-js/src/utils/checkpoint/checkpoint-helper.ts

@@ -6,4 +6,5 @@ export interface Checkpoint {
   force?(): Promise<void>;
   setTerminating?(): void;
   hasPendingAncestorCompletion?(stepId: string): boolean;
+  waitForQueueCompletion(): Promise<void>;
 }

packages/aws-durable-execution-sdk-js/src/utils/checkpoint/checkpoint-manager.ts

@@ -40,6 +40,8 @@ export class CheckpointManager implements Checkpoint {
     resolve: () => void;
     reject: (error: Error) => void;
   }> = [];
+  private queueCompletionResolver: (() => void) | null = null;
+  private queueCompletionTimeout: NodeJS.Timeout | null = null;
   private readonly MAX_PAYLOAD_SIZE = 750 * 1024; // 750KB in bytes
   private isTerminating = false;
   private static textEncoder = new TextEncoder();
@@ -98,6 +100,33 @@ export class CheckpointManager implements Checkpoint {
     });
   }
 
+  async waitForQueueCompletion(): Promise<void> {
+    if (this.queue.length === 0 && !this.isProcessing) {
+      return;
+    }
+
+    return new Promise<void>((resolve, reject) => {
+      this.queueCompletionResolver = resolve;
+
+      // Set a timeout to prevent infinite waiting
+      this.queueCompletionTimeout = setTimeout(() => {
+        this.queueCompletionResolver = null;
+        this.queueCompletionTimeout = null;
+        // Clear the queue since it's taking too long
+        this.clearQueue();
+        reject(new Error("Timeout waiting for checkpoint queue completion"));
+      }, 3000); // 3 second timeout
+    });
+  }
+
+  public clearQueue(): void {
+    // Silently clear queue - we're terminating so no need to reject promises
+    this.queue = [];
+    this.forceCheckpointPromises = [];
+    // Resolve any waiting queue completion promises since we're clearing
+    this.notifyQueueCompletion();
+  }
+
   // Alias for backward compatibility with Checkpoint interface
   async force(): Promise<void> {
     return this.forceCheckpoint();
@@ -321,6 +350,9 @@ export class CheckpointManager implements Checkpoint {
 
       const checkpointError = this.classifyCheckpointError(error);
 
+      // Clear remaining queue silently - we're terminating
+      this.clearQueue();
+
       this.terminationManager.terminate({
         reason: TerminationReason.CHECKPOINT_FAILED,
         message: checkpointError.message,
@@ -333,10 +365,24 @@ export class CheckpointManager implements Checkpoint {
         setImmediate(() => {
           this.processQueue();
         });
+      } else {
+        // Queue is empty and processing is done - notify all waiting promises
+        this.notifyQueueCompletion();
       }
     }
   }
 
+  private notifyQueueCompletion(): void {
+    if (this.queueCompletionResolver) {
+      if (this.queueCompletionTimeout) {
+        clearTimeout(this.queueCompletionTimeout);
+        this.queueCompletionTimeout = null;
+      }
+      this.queueCompletionResolver();
+      this.queueCompletionResolver = null;
+    }
+  }
+
   private async processBatch(batch: QueuedCheckpoint[]): Promise<void> {
     const updates: OperationUpdate[] = batch.map((item) => {
       const hashedStepId = hashId(item.stepId);
@@ -368,10 +414,7 @@ export class CheckpointManager implements Checkpoint {
       })),
     });
 
-    const response = await this.storage.checkpoint(
-      checkpointData,
-      this.logger,
-    );
+    const response = await this.storage.checkpoint(checkpointData, this.logger);
 
     if (response.CheckpointToken) {
       this.currentTaskToken = response.CheckpointToken;

packages/aws-durable-execution-sdk-js/src/utils/checkpoint/checkpoint-queue-completion.test.ts

@@ -0,0 +1,157 @@
+import { CheckpointManager } from "./checkpoint-manager";
+import { createMockExecutionContext } from "../../testing/mock-context";
+import { TerminationManager } from "../../termination-manager/termination-manager";
+import { EventEmitter } from "events";
+import { createDefaultLogger } from "../logger/default-logger";
+
+jest.mock("../../utils/logger/logger", () => ({
+  log: jest.fn(),
+}));
+
+describe("CheckpointManager Queue Completion", () => {
+  let checkpointManager: CheckpointManager;
+  let mockContext: any;
+  let mockTerminationManager: TerminationManager;
+  let mockEmitter: EventEmitter;
+  let mockLogger: any;
+
+  beforeEach(() => {
+    mockContext = createMockExecutionContext();
+    mockTerminationManager = new TerminationManager();
+    mockEmitter = new EventEmitter();
+    mockLogger = createDefaultLogger();
+
+    checkpointManager = new CheckpointManager(
+      "test-arn",
+      mockContext._stepData,
+      mockContext.state,
+      mockTerminationManager,
+      undefined,
+      "test-token",
+      mockEmitter,
+      mockLogger,
+      new Set<string>(),
+    );
+  });
+
+  describe("waitForQueueCompletion", () => {
+    it("should resolve immediately when queue is empty", async () => {
+      await expect(
+        checkpointManager.waitForQueueCompletion(),
+      ).resolves.toBeUndefined();
+    });
+
+    it("should wait until queue is empty", async () => {
+      const mockCheckpoint = jest
+        .fn()
+        .mockImplementation(
+          () => new Promise((resolve) => setTimeout(() => resolve({}), 200)),
+        );
+      (checkpointManager as any).storage = { checkpoint: mockCheckpoint };
+
+      // Add item to queue
+      checkpointManager.checkpoint("test-step", {});
+
+      // Should wait for queue to empty
+      await expect(
+        checkpointManager.waitForQueueCompletion(),
+      ).resolves.toBeUndefined();
+      expect(mockCheckpoint).toHaveBeenCalled();
+    });
+
+    it("should handle multiple concurrent waits", async () => {
+      const waits = Promise.all([
+        checkpointManager.waitForQueueCompletion(),
+        checkpointManager.waitForQueueCompletion(),
+        checkpointManager.waitForQueueCompletion(),
+      ]);
+
+      await expect(waits).resolves.toEqual([undefined, undefined, undefined]);
+    });
+
+    it("should timeout after 3 seconds if queue doesn't complete", async () => {
+      jest.useFakeTimers();
+
+      const mockCheckpoint = jest.fn().mockImplementation(
+        () => new Promise(() => {}), // Never resolves
+      );
+      (checkpointManager as any).storage = { checkpoint: mockCheckpoint };
+
+      // Add item to queue
+      checkpointManager.checkpoint("test-step", {});
+
+      const waitPromise = checkpointManager.waitForQueueCompletion();
+
+      // Fast-forward time by 3 seconds
+      jest.advanceTimersByTime(3000);
+
+      await expect(waitPromise).rejects.toThrow(
+        "Timeout waiting for checkpoint queue completion",
+      );
+
+      jest.useRealTimers();
+    });
+
+    it("should clear queue on timeout", async () => {
+      jest.useFakeTimers();
+
+      const mockCheckpoint = jest.fn().mockImplementation(
+        () => new Promise(() => {}), // Never resolves
+      );
+      (checkpointManager as any).storage = { checkpoint: mockCheckpoint };
+
+      // Add items to queue
+      checkpointManager.checkpoint("test-step-1", {});
+      checkpointManager.checkpoint("test-step-2", {});
+
+      expect((checkpointManager as any).queue.length).toBe(2);
+
+      const waitPromise = checkpointManager.waitForQueueCompletion();
+
+      // Fast-forward time by 3 seconds
+      jest.advanceTimersByTime(3000);
+
+      await expect(waitPromise).rejects.toThrow(
+        "Timeout waiting for checkpoint queue completion",
+      );
+
+      // Queue should be cleared
+      expect((checkpointManager as any).queue.length).toBe(0);
+
+      jest.useRealTimers();
+    });
+  });
+
+  describe("clearQueue", () => {
+    it("should clear the queue", () => {
+      const mockCheckpoint = jest.fn().mockResolvedValue({});
+      (checkpointManager as any).storage = { checkpoint: mockCheckpoint };
+
+      // Add items to queue
+      checkpointManager.checkpoint("test-step-1", {});
+      checkpointManager.checkpoint("test-step-2", {});
+
+      expect((checkpointManager as any).queue.length).toBe(2);
+
+      checkpointManager.clearQueue();
+
+      expect((checkpointManager as any).queue.length).toBe(0);
+    });
+
+    it("should allow waitForQueueCompletion to resolve after clearQueue", async () => {
+      const mockCheckpoint = jest.fn().mockResolvedValue({});
+      (checkpointManager as any).storage = { checkpoint: mockCheckpoint };
+
+      // Add item to queue
+      checkpointManager.checkpoint("test-step", {});
+
+      // Clear queue
+      checkpointManager.clearQueue();
+
+      // Wait should resolve immediately
+      await expect(
+        checkpointManager.waitForQueueCompletion(),
+      ).resolves.toBeUndefined();
+    });
+  });
+});