feat(api-rest): support making external promise cancellable (#12167)

packages/api-rest/__tests__/common/internalPost.test.ts

@@ -8,7 +8,11 @@ import {
 	parseJsonError,
 } from '@aws-amplify/core/internals/aws-client-utils';
 
-import { post, cancel } from '../../src/common/internalPost';
+import {
+	post,
+	cancel,
+	updateRequestToBeCancellable,
+} from '../../src/common/internalPost';
 import { RestApiError, isCancelError } from '../../src/errors';
 
 jest.mock('@aws-amplify/core/internals/aws-client-utils');
@@ -207,13 +211,15 @@ describe('internal post', () => {
 				underLyingHandlerReject = reject;
 			})
 		);
+		const abortController = new AbortController();
 		const promise = post(mockAmplifyInstance, {
 			url: apiGatewayUrl,
+			abortController,
 		});
+		updateRequestToBeCancellable(promise, abortController);
 
 		// mock abort behavior
-		const abortSignal = mockUnauthenticatedHandler.mock.calls[0][1]
-			.abortSignal as AbortSignal;
+		const abortSignal = abortController.signal;
 		abortSignal.addEventListener('abort', () => {
 			const mockAbortError = new Error('AbortError');
 			mockAbortError.name = 'AbortError';

packages/api-rest/src/common/handler.ts

@@ -41,23 +41,7 @@ type SigningServiceInfo = {
  *
  * @internal
  */
-export const transferHandler = (
-	amplify: AmplifyClassV6,
-	options: HandlerOptions,
-	signingServiceInfo?: SigningServiceInfo
-) =>
-	createCancellableOperation(abortSignal =>
-		transferHandlerJob(
-			amplify,
-			{
-				...options,
-				abortSignal,
-			},
-			signingServiceInfo
-		)
-	);
-
-const transferHandlerJob = async (
+export const transferHandler = async (
 	amplify: AmplifyClassV6,
 	options: HandlerOptions & { abortSignal: AbortSignal },
 	signingServiceInfo?: SigningServiceInfo

packages/api-rest/src/common/internalPost.ts

@@ -5,47 +5,66 @@ import { AmplifyClassV6 } from '@aws-amplify/core';
 
 import { InternalPostInput, RestApiResponse } from '../types';
 import { transferHandler } from './handler';
+import { createCancellableOperation } from '../utils';
 
-const cancelTokenMap = new WeakMap<
-	Promise<RestApiResponse>,
-	(cancelMessage?: string) => void
->();
+const cancelTokenMap = new WeakMap<Promise<any>, AbortController>();
 
 /**
  * @internal
  */
 export const post = (
 	amplify: AmplifyClassV6,
-	{ url, options }: InternalPostInput
+	{ url, options, abortController }: InternalPostInput
 ): Promise<RestApiResponse> => {
-	const { response, cancel } = transferHandler(
-		amplify,
-		{
-			url,
-			method: 'POST',
-			...options,
-		},
-		options?.signingServiceInfo
-	);
-	const responseWithCleanUp = response.finally(() => {
+	const controller = abortController ?? new AbortController();
+	const responsePromise = createCancellableOperation(async () => {
+		const response = transferHandler(
+			amplify,
+			{
+				url,
+				method: 'POST',
+				...options,
+				abortSignal: controller.signal,
+			},
+			options?.signingServiceInfo
+		);
+		return response;
+	}, controller);
+
+	const responseWithCleanUp = responsePromise.finally(() => {
 		cancelTokenMap.delete(responseWithCleanUp);
 	});
-	cancelTokenMap.set(responseWithCleanUp, cancel);
 	return responseWithCleanUp;
 };
 
 /**
  * Cancels a request given the promise returned by `post`.
  * If the request is already completed, this function does nothing.
+ * It MUST be used after `updateRequestToBeCancellable` is called.
  */
 export const cancel = (
 	promise: Promise<RestApiResponse>,
 	message?: string
 ): boolean => {
-	const cancelFn = cancelTokenMap.get(promise);
-	if (cancelFn) {
-		cancelFn(message);
+	const controller = cancelTokenMap.get(promise);
+	if (controller) {
+		controller.abort(message);
+		if (controller.signal.reason !== message) {
+			// In runtimes where `AbortSignal.reason` is not supported, we track the reason ourselves.
+			// @ts-expect-error reason is read-only property.
+			controller.signal['reason'] = message;
+		}
 		return true;
 	}
 	return false;
 };
+
+/**
+ * MUST be used to make a promise including internal `post` API call cancellable.
+ */
+export const updateRequestToBeCancellable = (
+	promise: Promise<any>,
+	controller: AbortController
+) => {
+	cancelTokenMap.set(promise, controller);
+};

packages/api-rest/src/internals/index.ts

@@ -16,10 +16,13 @@ import { InternalPostInput } from '../types';
  * * If auth mode is 'apiKey', you MUST set 'x-api-key' custom header.
  * * If auth mode is 'oidc' or 'lambda' or 'userPool', you MUST set 'authorization' header.
  *
+ * To make the internal post cancellable, you must also call `updateRequestToBeCancellable()` with the promise from
+ * internal post call and the abort controller supplied to the internal post call.
+ *
  * @internal
  */
 export const post = (input: InternalPostInput) => {
 	return internalPost(Amplify, input);
 };
 
-export { cancel } from '../common/internalPost';
+export { cancel, updateRequestToBeCancellable } from '../common/internalPost';

packages/api-rest/src/internals/server.ts

@@ -19,6 +19,9 @@ import { InternalPostInput } from '../types';
  * * If auth mode is 'apiKey', you MUST set 'x-api-key' custom header.
  * * If auth mode is 'oidc' or 'lambda' or 'userPool', you MUST set 'authorization' header.
  *
+ * To make the internal post cancellable, you must also call `updateRequestToBeCancellable()` with the promise from
+ * internal post call and the abort controller supplied to the internal post call.
+ *
  * @internal
  */
 export const post = (
@@ -28,4 +31,4 @@ export const post = (
 	return internalPost(getAmplifyServerContext(contextSpec).amplify, input);
 };
 
-export { cancel } from '../common/internalPost';
+export { cancel, updateRequestToBeCancellable } from '../common/internalPost';

packages/api-rest/src/types/index.ts

@@ -93,4 +93,11 @@ export type InternalPostInput = {
 			region?: string;
 		};
 	};
+	/**
+	 * The abort controller to cancel the in-flight POST request.
+	 * Required if you want to make the internal post request cancellable. To make the internal post cancellable, you
+	 * must also call `updateRequestToBeCancellable()` with the promise from internal post call and the abort
+	 * controller.
+	 */
+	abortController?: AbortController;
 };

packages/api-rest/src/utils/createCancellableOperation.ts

@@ -0,0 +1,77 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { HttpResponse } from '@aws-amplify/core/internals/aws-client-utils';
+import { CancelledError, RestApiError } from '../errors';
+import { Operation } from '../types';
+import { parseRestApiServiceError } from './serviceError';
+
+/**
+ * Create a cancellable operation conforming to the internal POST API interface.
+ * @internal
+ */
+export function createCancellableOperation(
+	handler: () => Promise<HttpResponse>,
+	abortController: AbortController
+): Promise<HttpResponse>;
+/**
+ * Create a cancellable operation conforming to the external REST API interface.
+ * @internal
+ */
+export function createCancellableOperation(
+	handler: (signal: AbortSignal) => Promise<HttpResponse>
+): Promise<HttpResponse>;
+
+/**
+ * @internal
+ */
+export function createCancellableOperation(
+	handler: (signal?: AbortSignal) => Promise<HttpResponse>,
+	abortController?: AbortController
+): Operation<HttpResponse> | Promise<HttpResponse> {
+	const isInternalPost = (
+		handler: (signal?: AbortSignal) => Promise<HttpResponse>
+	): handler is () => Promise<HttpResponse> => !!abortController;
+	const signal = abortController?.signal;
+	const job = async () => {
+		try {
+			const response = await (isInternalPost(handler)
+				? handler()
+				: handler(signal));
+			if (response.statusCode >= 300) {
+				throw parseRestApiServiceError(response)!;
+			}
+			return response;
+		} catch (error) {
+			if (error.name === 'AbortError' || signal?.aborted === true) {
+				throw new CancelledError({
+					name: error.name,
+					message: signal.reason ?? error.message,
+					underlyingError: error,
+				});
+			}
+			throw new RestApiError({
+				...error,
+				underlyingError: error,
+			});
+		}
+	};
+
+	if (isInternalPost(handler)) {
+		return job();
+	} else {
+		const cancel = (abortMessage?: string) => {
+			if (signal?.aborted === true) {
+				return;
+			}
+			abortController?.abort(abortMessage);
+			// Abort reason is not widely support enough across runtimes and and browsers, so we set it
+			// if it is not already set.
+			if (signal?.reason !== abortMessage) {
+				// @ts-expect-error reason is a readonly property
+				signal['reason'] = abortMessage;
+			}
+		};
+		return { response: job(), cancel };
+	}
+}

packages/api-rest/src/utils/index.ts

@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-export { createCancellableOperation } from './apiOperation';
+export { createCancellableOperation } from './createCancellableOperation';
 export { resolveCredentials } from './resolveCredentials';
 export { parseUrl } from './parseUrl';
 export { parseRestApiServiceError } from './serviceError';