diff --git a/packages/durabletask-js/src/testing/in-memory-backend.ts b/packages/durabletask-js/src/testing/in-memory-backend.ts index 2c96acb1..38eac32a 100644 --- a/packages/durabletask-js/src/testing/in-memory-backend.ts +++ b/packages/durabletask-js/src/testing/in-memory-backend.ts @@ -337,29 +337,42 @@ export class InMemoryOrchestrationBackend { } return new Promise((resolve, reject) => { - const timer = setTimeout(() => { - const waiters = this.stateWaiters.get(instanceId); - if (waiters) { - const index = waiters.findIndex((w) => w.resolve === resolve); - if (index >= 0) { - waiters.splice(index, 1); - } - } - reject(new Error(`Timeout waiting for orchestration '${instanceId}'`)); - }, timeoutMs); + // Declare timer with let so we can reference it inside the waiter + // and reference the waiter inside the timer callback. + // eslint-disable-next-line prefer-const + let timer: ReturnType; const waiter: StateWaiter = { resolve: (result) => { clearTimeout(timer); + this.pendingTimers.delete(timer); resolve(result); }, reject: (error) => { clearTimeout(timer); + this.pendingTimers.delete(timer); reject(error); }, predicate, }; + timer = setTimeout(() => { + this.pendingTimers.delete(timer); + const waiters = this.stateWaiters.get(instanceId); + if (waiters) { + const index = waiters.indexOf(waiter); + if (index >= 0) { + waiters.splice(index, 1); + } + if (waiters.length === 0) { + this.stateWaiters.delete(instanceId); + } + } + reject(new Error(`Timeout waiting for orchestration '${instanceId}'`)); + }, timeoutMs); + + this.pendingTimers.add(timer); + let waiters = this.stateWaiters.get(instanceId); if (!waiters) { waiters = []; diff --git a/packages/durabletask-js/test/in-memory-backend.spec.ts b/packages/durabletask-js/test/in-memory-backend.spec.ts index 66fe6a96..220ce6fc 100644 --- a/packages/durabletask-js/test/in-memory-backend.spec.ts +++ b/packages/durabletask-js/test/in-memory-backend.spec.ts @@ -377,4 +377,58 @@ describe("In-Memory Backend", () => { expect(state?.runtimeStatus).toEqual(OrchestrationStatus.COMPLETED); expect(state?.serializedOutput).toEqual(JSON.stringify(42)); }); + + it("should clean up stale waiters after waitForState timeout", async () => { + const instanceId = "timeout-cleanup-test"; + backend.createInstance(instanceId, "testOrch"); + + // Call waitForState with a predicate that will never match and a short timeout + await expect( + backend.waitForState( + instanceId, + () => false, // Will never match + 50, // 50ms timeout + ), + ).rejects.toThrow("Timeout waiting for orchestration"); + + // After the timeout, the stale waiter should be cleaned up. + // Access internal stateWaiters to verify the waiter was removed. + const stateWaitersMap = (backend as any).stateWaiters as Map; + + // The timed-out waiter should have been removed, and since it was the only + // waiter, the instance entry should be removed from the map entirely. + expect(stateWaitersMap.has(instanceId)).toBe(false); + }); + + it("should remove only the timed-out waiter when multiple waiters exist", async () => { + const instanceId = "multi-waiter-timeout-test"; + backend.createInstance(instanceId, "testOrch"); + + // Start a waiter with a long timeout (won't time out during the test) + const longWaitPromise = backend.waitForState( + instanceId, + () => false, // Never matches + 60000, // 60 second timeout — won't fire + ); + + // Start a waiter with a very short timeout + await expect( + backend.waitForState( + instanceId, + () => false, // Never matches + 50, // 50ms timeout + ), + ).rejects.toThrow("Timeout waiting for orchestration"); + + // After the short timeout, only the long-lived waiter should remain + const stateWaitersMap = (backend as any).stateWaiters as Map; + const waiters = stateWaitersMap.get(instanceId); + + expect(waiters).toBeDefined(); + expect(waiters!.length).toBe(1); + + // Clean up: reset to clear the remaining waiter and its timer + backend.reset(); + await longWaitPromise.catch(() => {}); // Ignore the reset rejection + }); });