feat(client): Add withAbortSignal context manager (#1272)

## What was changed

Clients and the Connection class get a `withAbortSignal` context manager
that can be used to cancel all ongoing connections started within a
function's scope.

```ts
const ctrl = new AbortController();
setTimeout(() => ctrl.abort(), 10_000);
// 👇 throws if incomplete by the timeout.
await conn.withAbortSignal(ctrl.signal, () => { /* make client or direct grpc calls */);
```

## Why?

Provide a way to abort client grpc calls.

Author: bergundy

packages/activity/src/index.ts

@@ -69,6 +69,7 @@
  * @module
  */
 
+// Keep this around until we drop support for Node 14.
 import 'abort-controller/polyfill'; // eslint-disable-line import/no-unassigned-import
 import { AsyncLocalStorage } from 'node:async_hooks';
 import { Logger, Duration, LogLevel, LogMetadata } from '@temporalio/common';

packages/client/src/base-client.ts

@@ -1,3 +1,5 @@
+// Keep this around until we drop support for Node 14.
+import 'abort-controller/polyfill'; // eslint-disable-line import/no-unassigned-import
 import os from 'node:os';
 import { DataConverter, LoadedDataConverter } from '@temporalio/common';
 import { isLoadedDataConverter, loadDataConverter } from '@temporalio/common/lib/internal-non-workflow';
@@ -67,6 +69,18 @@ export class BaseClient {
     return await this.connection.withDeadline(deadline, fn);
   }
 
+  /**
+   * Set an {@link https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal | `AbortSignal`} that, when aborted,
+   * cancels any ongoing service requests executed in `fn`'s scope.
+   *
+   * @returns value returned from `fn`
+   *
+   * @see {@link Connection.withAbortSignal}
+   */
+  async withAbortSignal<R>(abortSignal: AbortSignal, fn: () => Promise<R>): Promise<R> {
+    return await this.connection.withAbortSignal(abortSignal, fn);
+  }
+
   /**
    * Set metadata for any service requests executed in `fn`'s scope.
    *

packages/client/src/connection.ts

@@ -1,3 +1,5 @@
+// Keep this around until we drop support for Node 14.
+import 'abort-controller/polyfill'; // eslint-disable-line import/no-unassigned-import
 import { AsyncLocalStorage } from 'node:async_hooks';
 import * as grpc from '@grpc/grpc-js';
 import type { RPCImpl } from 'protobufjs';
@@ -375,7 +377,7 @@ export class Connection {
   }: RPCImplOptions): RPCImpl {
     return (method: { name: string }, requestData: any, callback: grpc.requestCallback<any>) => {
       const metadataContainer = new grpc.Metadata();
-      const { metadata, deadline } = callContextStorage.getStore() ?? {};
+      const { metadata, deadline, abortSignal } = callContextStorage.getStore() ?? {};
       for (const [k, v] of Object.entries(staticMetadata)) {
         metadataContainer.set(k, v);
       }
@@ -384,7 +386,7 @@ export class Connection {
           metadataContainer.set(k, v);
         }
       }
-      return client.makeUnaryRequest(
+      const call = client.makeUnaryRequest(
         `/${serviceName}/${method.name}`,
         (arg: any) => arg,
         (arg: any) => arg,
@@ -393,6 +395,11 @@ export class Connection {
         { interceptors, deadline },
         callback
       );
+      if (abortSignal != null) {
+        abortSignal.addEventListener('abort', () => call.cancel());
+      }
+
+      return call;
     };
   }
 
@@ -403,7 +410,27 @@ export class Connection {
    */
   async withDeadline<ReturnType>(deadline: number | Date, fn: () => Promise<ReturnType>): Promise<ReturnType> {
     const cc = this.callContextStorage.getStore();
-    return await this.callContextStorage.run({ deadline, metadata: cc?.metadata }, fn);
+    return await this.callContextStorage.run({ ...cc, deadline }, fn);
+  }
+
+  /**
+   * Set an {@link https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal | `AbortSignal`} that, when aborted,
+   * cancels any ongoing requests executed in `fn`'s scope.
+   *
+   * @returns value returned from `fn`
+   *
+   * @example
+   *
+   * ```ts
+   * const ctrl = new AbortController();
+   * setTimeout(() => ctrl.abort(), 10_000);
+   * // 👇 throws if incomplete by the timeout.
+   * await conn.withAbortSignal(ctrl.signal, () => client.workflow.execute(myWorkflow, options));
+   * ```
+   */
+  async withAbortSignal<ReturnType>(abortSignal: AbortSignal, fn: () => Promise<ReturnType>): Promise<ReturnType> {
+    const cc = this.callContextStorage.getStore();
+    return await this.callContextStorage.run({ ...cc, abortSignal }, fn);
   }
 
   /**
@@ -416,16 +443,15 @@ export class Connection {
    *
    * @example
    *
-   *```ts
-   *const workflowHandle = await conn.withMetadata({ apiKey: 'secret' }, () =>
-   *  conn.withMetadata({ otherKey: 'set' }, () => client.start(options)))
-   *);
-   *```
+   * ```ts
+   * const workflowHandle = await conn.withMetadata({ apiKey: 'secret' }, () =>
+   *   conn.withMetadata({ otherKey: 'set' }, () => client.start(options)))
+   * );
+   * ```
    */
   async withMetadata<ReturnType>(metadata: Metadata, fn: () => Promise<ReturnType>): Promise<ReturnType> {
     const cc = this.callContextStorage.getStore();
-    metadata = { ...cc?.metadata, ...metadata };
-    return await this.callContextStorage.run({ metadata, deadline: cc?.deadline }, fn);
+    return await this.callContextStorage.run({ ...cc, metadata: { ...cc?.metadata, ...metadata } }, fn);
   }
 
   /**

packages/client/src/types.ts

@@ -82,6 +82,8 @@ export interface CallContext {
    * Metadata to set on gRPC requests
    */
   metadata?: Metadata;
+
+  abortSignal?: AbortSignal;
 }
 
 /**
@@ -105,4 +107,12 @@ export interface ConnectionLike {
    * @returns returned value of `fn`
    */
   withMetadata<R>(metadata: Metadata, fn: () => Promise<R>): Promise<R>;
+
+  /**
+   * Set an {@link https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal | `AbortSignal`} that, when aborted,
+   * cancels any ongoing requests executed in `fn`'s scope.
+   *
+   * @returns value returned from `fn`
+   */
+  withAbortSignal<R>(abortSignal: AbortSignal, fn: () => Promise<R>): Promise<R>;
 }

packages/test/src/test-client-connection.ts

@@ -4,7 +4,7 @@ import fs from 'node:fs/promises';
 import test from 'ava';
 import * as grpc from '@grpc/grpc-js';
 import * as protoLoader from '@grpc/proto-loader';
-import { Connection, defaultGrpcRetryOptions, makeGrpcRetryInterceptor } from '@temporalio/client';
+import { Connection, defaultGrpcRetryOptions, isGrpcServiceError, makeGrpcRetryInterceptor } from '@temporalio/client';
 import pkg from '@temporalio/client/lib/pkg';
 import { temporal, grpc as grpcProto } from '@temporalio/proto';
 
@@ -38,7 +38,7 @@ async function bindLocalhostTls(server: grpc.Server): Promise<number> {
   return await util.promisify(server.bindAsync.bind(server))('localhost:0', credentials);
 }
 
-test('withMetadata / withDeadline set the CallContext for RPC call', async (t) => {
+test('withMetadata / withDeadline / withAbortSignal set the CallContext for RPC call', async (t) => {
   const server = new grpc.Server();
   let gotTestHeaders = false;
   let gotDeadline = false;
@@ -73,6 +73,9 @@ test('withMetadata / withDeadline set the CallContext for RPC call', async (t) =
       }
       callback(null, {});
     },
+    updateNamespace() {
+      // Simulate a hanging call to test abort signal support.
+    },
   });
   const port = await bindLocalhost(server);
   server.start();
@@ -84,6 +87,10 @@ test('withMetadata / withDeadline set the CallContext for RPC call', async (t) =
   );
   t.true(gotTestHeaders);
   t.true(gotDeadline);
+  const ctrl = new AbortController();
+  setTimeout(() => ctrl.abort(), 10);
+  const err = await t.throwsAsync(conn.withAbortSignal(ctrl.signal, () => conn.workflowService.updateNamespace({})));
+  t.true(isGrpcServiceError(err) && err.code === grpc.status.CANCELLED);
 });
 
 test('healthService works', async (t) => {