refactor: decouple TimerTask cancel via injected handler (match durabletask-python)

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: YunchuWang

packages/durabletask-js/src/task/timer-task.ts

@@ -1,22 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import * as pb from "../proto/orchestrator_service_pb";
 import { CompletableTask } from "./completable-task";
 
-/**
- * Minimal structural view of the orchestration context bookkeeping that a
- * {@link TimerTask} needs in order to cancel a pending timer.
- *
- * This is intentionally a narrow interface (rather than importing
- * `RuntimeOrchestrationContext`) so that `TimerTask` has no circular dependency
- * on the worker runtime. The concrete context satisfies it structurally.
- */
-export interface TimerBookkeeping {
-  _pendingActions: Record<number, pb.OrchestratorAction>;
-  _pendingTasks: Record<number, CompletableTask<any>>;
-}
-
 /**
  * A durable timer task returned by `OrchestrationContext.createTimer`.
  *
@@ -25,6 +11,11 @@ export interface TimerBookkeeping {
  * task wins (e.g. via `whenAny`), the losing timeout timer is canceled so the
  * orchestration is no longer waiting on it.
  *
+ * `TimerTask` is decoupled from the orchestration context's internal
+ * bookkeeping: the context injects a cancel handler via {@link setCancelHandler}
+ * and this class knows nothing about pending actions or tasks. This mirrors the
+ * `CancellableTask.set_cancel_handler` pattern in the Python SDK.
+ *
  * @example Cancel the loser of a race
  * ```typescript
  * const timeoutTask = context.createTimer(expirationDate);
@@ -36,45 +27,48 @@ export interface TimerBookkeeping {
  * ```
  */
 export class TimerTask extends CompletableTask<undefined> {
-  private readonly _bookkeeping: TimerBookkeeping;
-  private readonly _timerId: number;
+  private _cancelHandler?: () => void;
   private _isCanceled = false;
 
   /**
-   * Creates a new TimerTask.
-   *
-   * @param bookkeeping - The orchestration context bookkeeping holding the pending
-   *   actions and tasks (the concrete `RuntimeOrchestrationContext` satisfies this).
-   * @param timerId - The sequence id of the timer's `CreateTimer` action / pending task.
+   * Whether this timer has been canceled via {@link cancel}.
    */
-  constructor(bookkeeping: TimerBookkeeping, timerId: number) {
-    super();
-    this._bookkeeping = bookkeeping;
-    this._timerId = timerId;
+  get isCanceled(): boolean {
+    return this._isCanceled;
   }
 
   /**
-   * Whether this timer has been canceled via {@link cancel}.
+   * Registers the handler invoked when this timer is first canceled.
+   *
+   * The orchestration context supplies a closure that removes the timer's
+   * pending `CreateTimer` action and pending-task entry, so this class needs no
+   * knowledge of the context's internals.
+   *
+   * @param handler - Called once, when {@link cancel} first transitions the timer
+   *   to the canceled state.
    */
-  get isCanceled(): boolean {
-    return this._isCanceled;
+  setCancelHandler(handler: () => void): void {
+    this._cancelHandler = handler;
   }
 
   /**
    * Cancels this timer so the orchestration stops waiting on it.
    *
-   * Semantics:
-   * - If the timer's `CreateTimer` action has not yet been dispatched to the
-   *   sidecar (i.e. it is still pending in the current turn), it is removed so
-   *   the timer is never scheduled at all.
-   * - If the timer was already scheduled on a prior turn, it is removed from the
-   *   pending-task set so the orchestrator no longer waits on it; the backend
-   *   timer is reaped when the orchestration completes, and a late `TimerFired`
-   *   event is ignored because no pending task remains for it.
+   * The actual bookkeeping is performed by the cancel handler injected via
+   * {@link setCancelHandler}. The orchestration context's handler:
+   * - Removes the timer's `CreateTimer` action if it has not yet been dispatched
+   *   to the sidecar (i.e. it is still pending in the current turn), so the timer
+   *   is never scheduled at all.
+   * - Otherwise drops the timer from the pending-task set so the orchestrator no
+   *   longer waits on it; the backend timer is reaped when the orchestration
+   *   completes, and a late `TimerFired` event is ignored because no pending task
+   *   remains for it.
    *
    * This is deterministic and replay-safe: it consumes no sequence number and
-   * only deletes bookkeeping entries by key. Calling `cancel()` after the timer
-   * has already fired (completed) or after it was already canceled is a no-op.
+   * only runs the injected handler. Cancel does NOT mark the task complete
+   * (`isCompleted` stays false). Calling `cancel()` after the timer has already
+   * fired (completed) or after it was already canceled is a no-op, so the handler
+   * runs at most once.
    */
   cancel(): void {
     if (this._isComplete || this._isCanceled) {
@@ -83,11 +77,6 @@ export class TimerTask extends CompletableTask<undefined> {
     }
 
     this._isCanceled = true;
-
-    // Drop the pending CreateTimer action (if it has not been dispatched yet)
-    // so it is not scheduled, and stop tracking this timer as outstanding so
-    // the orchestrator is no longer waiting on it.
-    delete this._bookkeeping._pendingActions[this._timerId];
-    delete this._bookkeeping._pendingTasks[this._timerId];
+    this._cancelHandler?.();
   }
 }

packages/durabletask-js/src/worker/runtime-orchestration-context.ts

@@ -335,7 +335,11 @@ export class RuntimeOrchestrationContext extends OrchestrationContext {
     const action = ph.newCreateTimerAction(id, fireAtDate);
     this._pendingActions[action.getId()] = action;
 
-    const timerTask = new TimerTask(this, id);
+    const timerTask = new TimerTask();
+    timerTask.setCancelHandler(() => {
+      delete this._pendingActions[id];
+      delete this._pendingTasks[id];
+    });
     this._pendingTasks[id] = timerTask;
 
     return timerTask;

packages/durabletask-js/test/timer-task.spec.ts

@@ -1,27 +1,16 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import * as pb from "../src/proto/orchestrator_service_pb";
 import { CompletableTask } from "../src/task/completable-task";
-import { TimerTask, TimerBookkeeping } from "../src/task/timer-task";
-
-/**
- * Minimal fake bookkeeping that mirrors the fields TimerTask.cancel() touches on
- * the real RuntimeOrchestrationContext.
- */
-function makeBookkeeping(): TimerBookkeeping {
-  return {
-    _pendingActions: {} as Record<number, pb.OrchestratorAction>,
-    _pendingTasks: {} as Record<number, CompletableTask<any>>,
-  };
-}
+import { TimerTask } from "../src/task/timer-task";
+import { RuntimeOrchestrationContext } from "../src/worker/runtime-orchestration-context";
 
-describe("TimerTask", () => {
-  const TIMER_ID = 1;
+// A far-future fire time so timers created in these tests never fire on their own.
+const FUTURE_FIRE_AT = new Date(Date.now() + 24 * 60 * 60 * 1000);
 
+describe("TimerTask", () => {
   it("should start incomplete, not failed, and not canceled", () => {
-    const bookkeeping = makeBookkeeping();
-    const timer = new TimerTask(bookkeeping, TIMER_ID);
+    const timer = new TimerTask();
 
     expect(timer.isComplete).toBe(false);
     expect(timer.isCompleted).toBe(false);
@@ -30,34 +19,54 @@ describe("TimerTask", () => {
   });
 
   it("should be a Task/CompletableTask instance (identity preserving, no wrapper)", () => {
-    const bookkeeping = makeBookkeeping();
-    const timer = new TimerTask(bookkeeping, TIMER_ID);
+    const timer = new TimerTask();
 
     // TimerTask must extend CompletableTask so whenAny/whenAll return the real
     // instance and `winner === timerTask` identity holds for callers.
     expect(timer).toBeInstanceOf(CompletableTask);
   });
 
+  it("should return the same TimerTask instance that the context stores in _pendingTasks", () => {
+    // Identity: createTimer must hand back the exact instance it tracks so that
+    // `winner === timerTask` holds after whenAny.
+    const ctx = new RuntimeOrchestrationContext("test-instance");
+    const timer = ctx.createTimer(FUTURE_FIRE_AT);
+    const timerId = ctx._sequenceNumber;
+
+    expect(ctx._pendingTasks[timerId]).toBe(timer);
+  });
+
   describe("cancel()", () => {
-    it("should flip isCanceled and drop the pending action and pending task", () => {
-      const bookkeeping = makeBookkeeping();
-      const action = new pb.OrchestratorAction();
-      action.setId(TIMER_ID);
-      const timer = new TimerTask(bookkeeping, TIMER_ID);
-      bookkeeping._pendingActions[TIMER_ID] = action;
-      bookkeeping._pendingTasks[TIMER_ID] = timer;
+    it("should flip isCanceled and run the injected cancel handler once", () => {
+      const timer = new TimerTask();
+      const cancelHandler = jest.fn();
+      timer.setCancelHandler(cancelHandler);
 
       timer.cancel();
 
       expect(timer.isCanceled).toBe(true);
-      expect(bookkeeping._pendingActions[TIMER_ID]).toBeUndefined();
-      expect(bookkeeping._pendingTasks[TIMER_ID]).toBeUndefined();
+      expect(cancelHandler).toHaveBeenCalledTimes(1);
+    });
+
+    it("should drop the pending CreateTimer action and pending task via the context handler", () => {
+      // Drive a real context so the injected closure exercises the actual deletion.
+      const ctx = new RuntimeOrchestrationContext("test-instance");
+      const timer = ctx.createTimer(FUTURE_FIRE_AT);
+      const timerId = ctx._sequenceNumber;
+
+      expect(ctx._pendingActions[timerId]).toBeDefined();
+      expect(ctx._pendingTasks[timerId]).toBe(timer);
+
+      timer.cancel();
+
+      expect(timer.isCanceled).toBe(true);
+      expect(ctx._pendingActions[timerId]).toBeUndefined();
+      expect(ctx._pendingTasks[timerId]).toBeUndefined();
     });
 
     it("should not mark the timer complete (isCompleted stays false, isFaulted false)", () => {
-      const bookkeeping = makeBookkeeping();
-      const timer = new TimerTask(bookkeeping, TIMER_ID);
-      bookkeeping._pendingTasks[TIMER_ID] = timer;
+      const timer = new TimerTask();
+      timer.setCancelHandler(jest.fn());
 
       timer.cancel();
 
@@ -66,45 +75,49 @@ describe("TimerTask", () => {
       expect(timer.isFaulted).toBe(false);
     });
 
-    it("should be idempotent when called multiple times", () => {
-      const bookkeeping = makeBookkeeping();
-      const timer = new TimerTask(bookkeeping, TIMER_ID);
-      bookkeeping._pendingTasks[TIMER_ID] = timer;
+    it("should be idempotent when called multiple times (handler runs only once)", () => {
+      const timer = new TimerTask();
+      const cancelHandler = jest.fn();
+      timer.setCancelHandler(cancelHandler);
 
       timer.cancel();
       expect(() => timer.cancel()).not.toThrow();
+
       expect(timer.isCanceled).toBe(true);
+      expect(cancelHandler).toHaveBeenCalledTimes(1);
     });
 
     it("should not remove an unrelated pending action/task with a different id", () => {
-      const bookkeeping = makeBookkeeping();
-      const otherAction = new pb.OrchestratorAction();
-      otherAction.setId(2);
-      const otherTask = new CompletableTask<number>();
-      bookkeeping._pendingActions[2] = otherAction;
-      bookkeeping._pendingTasks[2] = otherTask;
-
-      const timer = new TimerTask(bookkeeping, TIMER_ID);
-      bookkeeping._pendingTasks[TIMER_ID] = timer;
-
-      timer.cancel();
-
-      // The sibling entries at id 2 must remain untouched.
-      expect(bookkeeping._pendingActions[2]).toBe(otherAction);
-      expect(bookkeeping._pendingTasks[2]).toBe(otherTask);
+      // Two sibling timers on one real context; canceling one must not touch the other.
+      const ctx = new RuntimeOrchestrationContext("test-instance");
+      const firstTimer = ctx.createTimer(FUTURE_FIRE_AT);
+      const firstId = ctx._sequenceNumber;
+      const secondTimer = ctx.createTimer(FUTURE_FIRE_AT);
+      const secondId = ctx._sequenceNumber;
+
+      firstTimer.cancel();
+
+      // The sibling entries at the other id must remain untouched.
+      expect(ctx._pendingActions[secondId]).toBeDefined();
+      expect(ctx._pendingTasks[secondId]).toBe(secondTimer);
+      // The canceled timer's entries are gone.
+      expect(ctx._pendingActions[firstId]).toBeUndefined();
+      expect(ctx._pendingTasks[firstId]).toBeUndefined();
     });
 
     it("should be a no-op after the timer has already fired (completed)", () => {
-      const bookkeeping = makeBookkeeping();
-      const timer = new TimerTask(bookkeeping, TIMER_ID);
+      const timer = new TimerTask();
+      const cancelHandler = jest.fn();
+      timer.setCancelHandler(cancelHandler);
 
       // Simulate the timer firing (handleTimerFired calls complete(undefined)).
       timer.complete(undefined);
 
       expect(() => timer.cancel()).not.toThrow();
-      // Canceling a fired timer must not flip isCanceled.
+      // Canceling a fired timer must not flip isCanceled or run the handler.
       expect(timer.isCanceled).toBe(false);
       expect(timer.isCompleted).toBe(true);
+      expect(cancelHandler).not.toHaveBeenCalled();
     });
   });
 });