feat(activity): notify Activities when the Worker starts shutting down

packages/activity/src/index.ts

@@ -378,8 +378,40 @@ export class Context {
     /**
      * Holder object for activity cancellation details
      */
-    protected readonly _cancellationDetails: ActivityCancellationDetailsHolder
-  ) {}
+    protected readonly _cancellationDetails: ActivityCancellationDetailsHolder,
+
+    /**
+     * A Promise that fails with a {@link CancelledFailure} when the Worker running this Activity starts shutting down,
+     * i.e. as soon as shutdown is initiated, rather than on expiration of the shutdown grace period. The promise is
+     * guaranteed to never successfully resolve.
+     *
+     * Await this promise in a long running Activity to get an early notice that the Worker is going away, so that the
+     * Activity may checkpoint its progress and return, rather than being abruptly cancelled once the grace period
+     * expires. Unlike {@link Context.cancelled}, an Activity does _not_ need to {@link Context.heartbeat} to get
+     * notified of Worker shutdown.
+     *
+     * @experimental Worker shutdown notification is experimental and may be subject to change.
+     */
+    public readonly workerShuttingDown: Promise<never> = new Promise<never>(() => undefined),
+
+    /**
+     * An {@link https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal | `AbortSignal`} that is aborted when the
+     * Worker running this Activity starts shutting down, i.e. as soon as shutdown is initiated, rather than on
+     * expiration of the shutdown grace period.
+     *
+     * This can be passed in to libraries such as
+     * {@link https://www.npmjs.com/package/node-fetch#request-cancellation-with-abortsignal | fetch}, in the same way
+     * as {@link Context.cancellationSignal}. Unlike cancellation, an Activity does _not_ need to
+     * {@link Context.heartbeat} to get notified of Worker shutdown.
+     *
+     * @experimental Worker shutdown notification is experimental and may be subject to change.
+     */
+    public readonly workerShuttingDownSignal: AbortSignal = new AbortController().signal
+  ) {
+    // The `workerShuttingDown` promise is meant to be raced against by user code; it may legitimately
+    // never be awaited, so make sure it doesn't surface as an unhandled rejection.
+    this.workerShuttingDown.catch(() => undefined);
+  }
 
   /**
    * Send a {@link https://docs.temporal.io/encyclopedia/detecting-activity-failures#activity-heartbeat | heartbeat} from an Activity.
@@ -555,6 +587,41 @@ export function cancellationSignal(): AbortSignal {
   return Context.current().cancellationSignal;
 }
 
+/**
+ * Return a Promise that fails with a {@link CancelledFailure} when the Worker running this Activity starts shutting
+ * down, i.e. as soon as shutdown is initiated, rather than on expiration of the shutdown grace period. The promise is
+ * guaranteed to never successfully resolve.
+ *
+ * Await this promise in a long running Activity to get an early notice that the Worker is going away, so that the
+ * Activity may checkpoint its progress and return, rather than being abruptly cancelled once the grace period expires.
+ * Unlike {@link cancelled}, an Activity does _not_ need to {@link heartbeat} to get notified of Worker shutdown.
+ *
+ * This is a shortcut for `Context.current().workerShuttingDown` (see {@link Context.workerShuttingDown}).
+ *
+ * @experimental Worker shutdown notification is experimental and may be subject to change.
+ */
+export function workerShuttingDown(): Promise<never> {
+  return Context.current().workerShuttingDown;
+}
+
+/**
+ * Return an {@link https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal | `AbortSignal`} that is aborted when
+ * the Worker running this Activity starts shutting down, i.e. as soon as shutdown is initiated, rather than on
+ * expiration of the shutdown grace period.
+ *
+ * This can be passed in to libraries such as
+ * {@link https://www.npmjs.com/package/node-fetch#request-cancellation-with-abortsignal | fetch}, in the same way as
+ * {@link cancellationSignal}. Unlike cancellation, an Activity does _not_ need to {@link heartbeat} to get notified of
+ * Worker shutdown.
+ *
+ * This is a shortcut for `Context.current().workerShuttingDownSignal` (see {@link Context.workerShuttingDownSignal}).
+ *
+ * @experimental Worker shutdown notification is experimental and may be subject to change.
+ */
+export function workerShuttingDownSignal(): AbortSignal {
+  return Context.current().workerShuttingDownSignal;
+}
+
 /**
  * A Temporal Client, bound to the same Temporal Namespace as the Worker executing this Activity.
  *

packages/testing/src/__tests__/test-mockactivityenv.ts

@@ -47,3 +47,58 @@ test('MockActivityEnvironment injects provided info', async (t) => {
   }, 1);
   t.is(res, 4);
 });
+
+test('MockActivityEnvironment notifies activities of worker shutdown', async (t) => {
+  const env = new MockActivityEnvironment();
+  const res = await env.run(async (x: number): Promise<number> => {
+    t.false(activity.workerShuttingDownSignal().aborted);
+    setImmediate(() => env.notifyWorkerShuttingDown());
+    // A worker shutdown notification is not a cancellation, so this activity is free to complete normally.
+    await t.throwsAsync(activity.workerShuttingDown(), {
+      instanceOf: activity.CancelledFailure,
+      message: 'WORKER_SHUTDOWN',
+    });
+    t.true(activity.workerShuttingDownSignal().aborted);
+    return x + 1;
+  }, 3);
+  t.is(res, 4);
+});
+
+test('Worker shutdown notification does not cancel the activity', async (t) => {
+  const env = new MockActivityEnvironment();
+  await env.run(async (): Promise<void> => {
+    env.notifyWorkerShuttingDown();
+    t.true(activity.workerShuttingDownSignal().aborted);
+    // Cancellation is a distinct concern and must remain untouched.
+    t.false(activity.cancellationSignal().aborted);
+    t.is(activity.cancellationDetails(), undefined);
+    await activity.sleep(1);
+    t.pass();
+  });
+});
+
+test('Worker shutdown notification is observable through Promise.race', async (t) => {
+  const env = new MockActivityEnvironment();
+  const res = await env.run(async (): Promise<string> => {
+    setImmediate(() => env.notifyWorkerShuttingDown());
+    // The idiomatic usage: bail out of long running work as soon as the worker starts going away.
+    return await Promise.race([
+      activity.sleep(30_000).then(() => 'completed'),
+      activity.workerShuttingDown().catch(() => 'interrupted'),
+    ]);
+  });
+  t.is(res, 'interrupted');
+});
+
+test('Worker shutdown notification is idempotent and applies to activities started afterwards', async (t) => {
+  const env = new MockActivityEnvironment();
+  env.notifyWorkerShuttingDown();
+  env.notifyWorkerShuttingDown();
+  await env.run(async (): Promise<void> => {
+    t.true(activity.workerShuttingDownSignal().aborted);
+    await t.throwsAsync(activity.workerShuttingDown(), {
+      instanceOf: activity.CancelledFailure,
+      message: 'WORKER_SHUTDOWN',
+    });
+  });
+});

packages/testing/src/mocking-activity-environment.ts

@@ -33,6 +33,14 @@ export interface MockActivityEnvironmentOptions {
  */
 export class MockActivityEnvironment extends events.EventEmitter {
   public cancel: (reason?: CancelReason, details?: ActivityCancellationDetails) => void = () => undefined;
+
+  /**
+   * Simulate initiation of Worker shutdown, resolving {@link activity.Context.workerShuttingDown} and aborting
+   * {@link activity.Context.workerShuttingDownSignal}. This doesn't cancel the Activity.
+   *
+   * @experimental Worker shutdown notification is experimental and may be subject to change.
+   */
+  public notifyWorkerShuttingDown: () => void = () => undefined;
   public readonly context: activity.Context;
   private readonly activity: Activity;
 
@@ -63,6 +71,7 @@ export class MockActivityEnvironment extends events.EventEmitter {
       opts?.interceptors ?? []
     );
     this.context = this.activity.context;
+    this.notifyWorkerShuttingDown = () => this.activity.notifyWorkerShuttingDown();
     this.cancel = (reason?: CancelReason, details?: ActivityCancellationDetails) => {
       // Default to CANCELLED if nothing provided.
       const r = reason ?? 'CANCELLED';

packages/worker/src/activity.ts

@@ -47,6 +47,12 @@ export class Activity {
   public cancel: (reason: CancelReason, details: ActivityCancellationDetails) => void = () => undefined;
   public readonly abortController: AbortController = new AbortController();
 
+  /**
+   * Aborted as soon as the Worker running this Activity initiates shutdown, i.e. before expiration of the shutdown
+   * grace period. Distinct from {@link abortController}, which is only aborted once the Activity itself gets cancelled.
+   */
+  public readonly workerShuttingDownAbortController: AbortController = new AbortController();
+
   /**
    * Logger bound to `sdkComponent: worker`, with metadata from this activity.
    * This is the logger to use for all log messages emitted by the activity
@@ -74,7 +80,8 @@ export class Activity {
     private readonly _client: Client | undefined, // May be undefined in the case of MockActivityEnvironment
     workerLogger: Logger,
     workerMetricMeter: MetricMeter,
-    interceptors: ActivityInterceptorsFactory[]
+    interceptors: ActivityInterceptorsFactory[],
+    workerShuttingDownSignal?: AbortSignal
   ) {
     this.workerLogger = LoggerWithComposedMetadata.compose(workerLogger, this.getLogAttributes.bind(this));
     this.metricMeter = MetricMeterWithComposedTags.compose(workerMetricMeter, this.getMetricTags.bind(this));
@@ -88,6 +95,22 @@ export class Activity {
         reject(err);
       };
     });
+    // The Worker owns a single signal shared by all of its Activities; each Activity derives its own controller from
+    // it, so that user code may add listeners without accumulating them on the Worker's own signal.
+    const workerShuttingDownPromise = new Promise<never>((_, reject) => {
+      this.workerShuttingDownAbortController.signal.addEventListener(
+        'abort',
+        () => reject(this.workerShuttingDownAbortController.signal.reason),
+        { once: true }
+      );
+    });
+    if (workerShuttingDownSignal !== undefined) {
+      if (workerShuttingDownSignal.aborted) {
+        this.notifyWorkerShuttingDown();
+      } else {
+        workerShuttingDownSignal.addEventListener('abort', () => this.notifyWorkerShuttingDown(), { once: true });
+      }
+    }
     this.context = new Context(
       info,
       promise,
@@ -97,10 +120,13 @@ export class Activity {
       // This is the activity context logger, to be used exclusively from user code
       LoggerWithComposedMetadata.compose(this.workerLogger, { sdkComponent: SdkComponent.activity }),
       this.metricMeter,
-      this.cancellationDetails
+      this.cancellationDetails,
+      workerShuttingDownPromise,
+      this.workerShuttingDownAbortController.signal
     );
     // Prevent unhandled rejection
     promise.catch(() => undefined);
+    workerShuttingDownPromise.catch(() => undefined);
     this.interceptors = { inbound: [], outbound: [] };
     interceptors
       .map((factory) => factory(this.context))
@@ -110,6 +136,17 @@ export class Activity {
       });
   }
 
+  /**
+   * Notify this Activity that the Worker running it has started shutting down.
+   *
+   * Contrary to {@link cancel}, this doesn't request cancellation of the Activity; it merely gives the Activity an
+   * early chance to wrap up its work before the shutdown grace period expires, at which point it would get cancelled.
+   */
+  public notifyWorkerShuttingDown(): void {
+    if (this.workerShuttingDownAbortController.signal.aborted) return;
+    this.workerShuttingDownAbortController.abort(new CancelledFailure('WORKER_SHUTDOWN'));
+  }
+
   protected getLogAttributes(): Record<string, unknown> {
     const logAttributes = activityLogAttributes(this.info);
     // In case some interceptor uses the logger while initializing...

packages/worker/src/worker.ts

@@ -478,6 +478,9 @@ export class Worker {
 
   protected readonly numInFlightActivationsSubject = new BehaviorSubject<number>(0);
   protected readonly numInFlightActivitiesSubject = new BehaviorSubject<number>(0);
+  // Aborted as soon as shutdown of this Worker is initiated; shared by all Activities executed by this Worker, each of
+  // which derives its own AbortController from it. See `Context.workerShuttingDown`.
+  protected readonly workerShuttingDownAbortController = new AbortController();
   protected readonly numInFlightNonLocalActivitiesSubject = new BehaviorSubject<number>(0);
   protected readonly numInFlightLocalActivitiesSubject = new BehaviorSubject<number>(0);
   protected readonly numInFlightNexusOperationsSubject = new BehaviorSubject<number>(0);
@@ -959,6 +962,11 @@ export class Worker {
 
   protected set state(state: State) {
     this.logger.info('Worker state changed', { state });
+    // Notify running Activities as soon as the Worker leaves the RUNNING state, whatever the reason, so that they get
+    // a chance to wrap up before the shutdown grace period expires. See `Context.workerShuttingDown`.
+    if (state !== 'INITIALIZED' && state !== 'RUNNING' && !this.workerShuttingDownAbortController.signal.aborted) {
+      this.workerShuttingDownAbortController.abort(new CancelledFailure('WORKER_SHUTDOWN'));
+    }
     this.stateSubject.next(state);
   }
 
@@ -1149,7 +1157,8 @@ export class Worker {
                       this.client,
                       this.logger,
                       this.metricMeter,
-                      this.options.interceptors.activity
+                      this.options.interceptors.activity,
+                      this.workerShuttingDownAbortController.signal
                     );
                     output = { type: 'run', activity, input };
                     break;