diff --git a/core/packages/gax/src/createApiCall.ts b/core/packages/gax/src/createApiCall.ts index 40b271a441c..f2f58190f37 100644 --- a/core/packages/gax/src/createApiCall.ts +++ b/core/packages/gax/src/createApiCall.ts @@ -135,6 +135,24 @@ export function createApiCall( } const ongoingCall = currentApiCaller.init(callback); + const signal = thisSettings.signal; + if (signal && !(ongoingCall instanceof StreamProxy)) { + if (signal.aborted) { + process.nextTick(() => { + ongoingCall.cancel(); + }); + return currentApiCaller.result(ongoingCall); + } + const abortListener = () => { + ongoingCall.cancel(); + }; + signal.addEventListener('abort', abortListener, {once: true}); + const originalCallback = ongoingCall.callback; + ongoingCall.callback = (...args: Parameters) => { + signal.removeEventListener('abort', abortListener); + originalCallback?.(...args); + }; + } // Track resends per logical request so auto-paginated subsequent page // requests start at resendCount = 0 instead of being counted as retries. diff --git a/core/packages/gax/src/gax.ts b/core/packages/gax/src/gax.ts index 98420f82179..bdb721f9211 100644 --- a/core/packages/gax/src/gax.ts +++ b/core/packages/gax/src/gax.ts @@ -178,6 +178,10 @@ export interface CallOptions { * GAPIC-generated client libraries; do not set manually. */ resumableUpload?: ResumableUploadContext; + /** + * An AbortSignal to cancel the ongoing non-streaming API call. + */ + signal?: AbortSignal; } export class CallSettings { @@ -196,6 +200,7 @@ export class CallSettings { retryRequestOptions?: RetryRequestOptions; enableTelemetryTracing?: boolean; resumableUpload?: ResumableUploadContext; + signal?: AbortSignal; /** * @param {Object} settings - An object containing parameters of this settings. @@ -232,6 +237,7 @@ export class CallSettings { this.enableTelemetryTracing = settings.enableTelemetryTracing; this.resumableUpload = 'resumableUpload' in settings ? settings.resumableUpload : undefined; + this.signal = 'signal' in settings ? settings.signal : undefined; } /** @@ -257,6 +263,7 @@ export class CallSettings { let retryRequestOptions = this.retryRequestOptions; let enableTelemetryTracing = this.enableTelemetryTracing; let resumableUpload = this.resumableUpload; + let signal = this.signal; // If the user provides a timeout to the method, that timeout value will be used // to override the backoff settings. @@ -318,6 +325,9 @@ export class CallSettings { if ('resumableUpload' in options) { resumableUpload = options.resumableUpload; } + if ('signal' in options) { + signal = options.signal; + } return new CallSettings({ timeout, @@ -332,6 +342,7 @@ export class CallSettings { retryRequestOptions, enableTelemetryTracing, resumableUpload, + signal, }); } } diff --git a/core/packages/gax/test/unit/apiCallable.ts b/core/packages/gax/test/unit/apiCallable.ts index 1c40c31d14f..15fb53de373 100644 --- a/core/packages/gax/test/unit/apiCallable.ts +++ b/core/packages/gax/test/unit/apiCallable.ts @@ -1769,8 +1769,8 @@ describe('createApiCall', () => { }); describe('Promise', () => { - it('calls api call', done => { - let deadlineArg: string; + it('calls api call', async () => { + let deadlineArg = ''; function func( argument: {}, metadata: {}, @@ -1782,14 +1782,10 @@ describe('Promise', () => { } const apiCall = createApiCall(func); // eslint-disable-next-line @typescript-eslint/no-explicit-any - (apiCall as any)(null) - .then((response: number[]) => { - assert.ok(Array.isArray(response)); - assert.strictEqual(response[0], 42); - assert.ok(deadlineArg); - return done(); - }) - .catch(done); + const response = await (apiCall as any)(null); + assert.ok(Array.isArray(response)); + assert.strictEqual(response[0], 42); + assert.ok(deadlineArg); }); it('emits error on rejected promise', async () => { @@ -1806,28 +1802,24 @@ describe('Promise', () => { await assert.rejects(apiCall({}, undefined)); }); - it('has cancel method', done => { + it('has cancel method', async () => { function func(argument: {}, metadata: {}, options: {}, callback: Function) { setTimeout(() => { callback(null, 42); }, 0); } - const apiCall = createApiCall(func, {cancel: done}); + const apiCall = createApiCall(func); // eslint-disable-next-line @typescript-eslint/no-explicit-any const promise = (apiCall as any)(null); - promise - .then(() => { - return done(new Error('should not reach')); - }) - .catch((err: {code: number}) => { - assert(err instanceof GoogleError); - assert.strictEqual(err.code, status.CANCELLED); - done(); - }); promise.cancel(); + await assert.rejects(promise, (err: unknown) => { + assert(err instanceof GoogleError); + assert.strictEqual(err.code, status.CANCELLED); + return true; + }); }); - it('cancels retrying call', done => { + it('cancels retrying call', async () => { const retryOptions = utils.createRetryOptions(0, 0, 0, 0, 0, 0, 100); let callCount = 0; @@ -1855,18 +1847,11 @@ describe('Promise', () => { }); // eslint-disable-next-line @typescript-eslint/no-explicit-any const promise = (apiCall as any)(null); - promise - .then(() => { - return done(new Error('should not reach')); - }) - .catch(() => { - assert(callCount < 4); - done(); - }) - .catch(done); setTimeout(() => { promise.cancel(); }, 15); + await assert.rejects(promise); + assert(callCount < 4); }); it('does not return promise when callback is supplied', done => { @@ -1883,6 +1868,91 @@ describe('Promise', () => { undefined, ); }); + + it('cancels promise-based call via AbortSignal', async () => { + function func(argument: {}, metadata: {}, options: {}, callback: Function) { + setTimeout(() => { + callback(null, 42); + }, 10); + } + const controller = new AbortController(); + const apiCall = createApiCall(func); + const promise = apiCall({}, {signal: controller.signal}); + controller.abort(); + await assert.rejects(promise, (err: unknown) => { + assert(err instanceof GoogleError); + assert.strictEqual(err.code, status.CANCELLED); + return true; + }); + }); + + it('cancels callback-based call via AbortSignal', done => { + let cancelCalled = false; + function func(argument: {}, metadata: {}, options: {}, callback: Function) { + const timer = setTimeout(() => { + callback(null, 42); + }, 50); + return () => { + cancelCalled = true; + clearTimeout(timer); + const err = new GoogleError('cancelled'); + err.code = status.CANCELLED; + callback(err); + }; + } + const controller = new AbortController(); + const apiCall = createApiCall(func, {returnCancelFunc: true}); + void apiCall({}, {signal: controller.signal}, err => { + try { + assert.strictEqual(cancelCalled, true); + assert(err instanceof GoogleError); + assert.strictEqual(err.code, status.CANCELLED); + done(); + } catch (assertErr) { + done(assertErr); + } + }); + setTimeout(() => { + controller.abort(); + }, 5); + }); + + it('immediately cancels call when AbortSignal is already aborted', done => { + let funcCalled = false; + let returnedSynchronously = false; + function func(argument: {}, metadata: {}, options: {}, callback: Function) { + funcCalled = true; + callback(null, 42); + } + const controller = new AbortController(); + controller.abort(); + const apiCall = createApiCall(func); + void apiCall({}, {signal: controller.signal}, err => { + try { + assert.strictEqual(returnedSynchronously, true); + assert.strictEqual(funcCalled, false); + assert(err instanceof GoogleError); + assert.strictEqual(err.code, status.CANCELLED); + done(); + } catch (assertErr) { + done(assertErr); + } + }); + returnedSynchronously = true; + }); + + it('removes abort listener from AbortSignal when call completes', async () => { + function func(argument: {}, metadata: {}, options: {}, callback: Function) { + callback(null, 42); + } + const controller = new AbortController(); + const removeSpy = sinon.spy(controller.signal, 'removeEventListener'); + const apiCall = createApiCall(func); + const [response] = await apiCall({}, {signal: controller.signal}); + assert.strictEqual(response, 42); + assert.strictEqual(removeSpy.calledWith('abort'), true); + removeSpy.restore(); + }); }); describe('retryable', () => { @@ -1922,9 +1992,9 @@ describe('retryable', () => { }); }); - it('retries the API call with promise', done => { + it('retries the API call with promise', async () => { let toAttempt = 3; - let deadlineArg: string; + let deadlineArg = ''; function func( argument: {}, metadata: {}, @@ -1940,18 +2010,14 @@ describe('retryable', () => { callback(null, 1729); } const apiCall = createApiCall(func, settings); - apiCall({}, undefined) - .then(resp => { - assert.ok(Array.isArray(resp)); - assert.strictEqual(resp[0], 1729); - assert.strictEqual(toAttempt, 0); - assert.ok(deadlineArg); - return done(); - }) - .catch(done); + const resp = await apiCall({}, undefined); + assert.ok(Array.isArray(resp)); + assert.strictEqual(resp[0], 1729); + assert.strictEqual(toAttempt, 0); + assert.ok(deadlineArg); }); - it('cancels in the middle of retries', done => { + it('cancels in the middle of retries', async () => { let callCount = 0; // eslint-disable-next-line @typescript-eslint/no-explicit-any function func(argument: {}, metadata: {}, options: {}, callback: Function) { @@ -1969,14 +2035,10 @@ describe('retryable', () => { } const apiCall = createApiCall(func, settings); const promise = apiCall({}, undefined); - promise - .then(() => { - return done(new Error('should not reach')); - }) - .catch((err: Error) => { - assert(err instanceof Error); - done(); - }); + await assert.rejects(promise, (err: unknown) => { + assert(err instanceof Error); + return true; + }); }); it("doesn't retry if no codes", done => { @@ -2188,7 +2250,7 @@ describe('retryable', () => { assert(spy.callCount > callsLowerBound); assert(spy.callCount < callsUpperBound); done(); - }).catch(done); + }); }); it.skip('reports A/B testing', () => { @@ -2236,12 +2298,12 @@ describe('retryable', () => { }); }); - it('forwards metadata to builder', done => { + it('forwards metadata to builder', async () => { function func(argument: {}, metadata: {}, options: {}, callback: Function) { callback(null, {}); } - let gotHeaders: {h1?: string; h2?: string}; + let gotHeaders: {h1?: string; h2?: string} = {}; const mockBuilder = (abTest: {}, headers: {}) => { gotHeaders = headers; }; @@ -2255,14 +2317,8 @@ describe('retryable', () => { h1: 'val1', h2: 'val2', }; - void apiCall({}, {otherArgs: {headers}}).then(() => { - try { - assert.strictEqual(gotHeaders.h1, 'val1'); - assert.strictEqual(gotHeaders.h2, 'val2'); - return done(); - } catch (err) { - return done(err); - } - }); + await apiCall({}, {otherArgs: {headers}}); + assert.strictEqual(gotHeaders.h1, 'val1'); + assert.strictEqual(gotHeaders.h2, 'val2'); }); });