Use the host AbortSignal if provided (#509)

For runtimes that provide a fetch compatible server api,
we should thread trough the abort signal that is provided by
them. For node and lambda we will provide our own.

Author: igalshilman

packages/restate-sdk/src/endpoint/handlers/fetch.ts

@@ -25,6 +25,7 @@ export function fetcher(handler: GenericHandler) {
         headers,
         body: event.body,
         extraArgs,
+        abortSignal: event.signal,
       };
 
       const resp = await handler.handle(request);

packages/restate-sdk/src/endpoint/handlers/generic.ts

@@ -55,6 +55,7 @@ export interface RestateRequest {
   readonly headers: Headers;
   readonly body: ReadableStream<Uint8Array> | null;
   readonly extraArgs: unknown[];
+  readonly abortSignal: AbortSignal;
 }
 
 export interface RestateResponse {
@@ -178,11 +179,13 @@ export class GenericHandler implements RestateHandler {
       this.endpoint.rlog.error(msg);
       return this.toErrorResponse(400, msg);
     }
+
     return this.handleInvoke(
       handler,
       request.body,
       request.headers,
       request.extraArgs,
+      request.abortSignal,
       context ?? {}
     );
   }
@@ -220,6 +223,7 @@ export class GenericHandler implements RestateHandler {
     body: ReadableStream<Uint8Array>,
     headers: Headers,
     extraArgs: unknown[],
+    abortSignal: AbortSignal,
     additionalContext: AdditionalContext
   ): Promise<RestateResponse> {
     const loggerId = Math.floor(Math.random() * 4_294_967_295 /* u32::MAX */);
@@ -277,8 +281,6 @@ export class GenericHandler implements RestateHandler {
     // Get input
     const input = coreVm.sys_input();
 
-    const abortController = new AbortController();
-
     const invocationRequest: Request = {
       id: input.invocation_id,
       headers: input.headers.reduce((headers, { key, value }) => {
@@ -296,7 +298,7 @@ export class GenericHandler implements RestateHandler {
       ),
       body: input.input,
       extraArgs,
-      attemptCompletedSignal: abortController.signal,
+      attemptCompletedSignal: abortSignal,
     };
 
     // Prepare logger
@@ -398,11 +400,6 @@ export class GenericHandler implements RestateHandler {
         inputReader.cancel().catch(() => {});
       })
       .finally(() => {
-        try {
-          abortController.abort();
-        } catch (e) {
-          // suppressed
-        }
         invocationLoggers.delete(loggerId);
       })
       .catch(() => {});

packages/restate-sdk/src/endpoint/handlers/lambda.ts

@@ -17,7 +17,11 @@ import type {
   Context,
 } from "aws-lambda";
 import { Buffer } from "node:buffer";
-import type { GenericHandler, RestateRequest } from "./generic.js";
+import type {
+  GenericHandler,
+  RestateRequest,
+  RestateResponse,
+} from "./generic.js";
 import { WritableStream, type ReadableStream } from "node:stream/web";
 import { OnceStream } from "../../utils/streams.js";
 import { X_RESTATE_SERVER } from "../../user_agent.js";
@@ -48,16 +52,26 @@ export class LambdaHandler {
       body = OnceStream(new TextEncoder().encode(event.body));
     }
 
+    const abortController = new AbortController();
+
     const request: RestateRequest = {
       body,
       headers: event.headers,
       url: path,
       extraArgs: [context],
+      abortSignal: abortController.signal,
     };
 
-    const resp = await this.handler.handle(request, {
-      AWSRequestId: context.awsRequestId,
-    });
+    let resp: RestateResponse;
+
+    try {
+      resp = await this.handler.handle(request, {
+        AWSRequestId: context.awsRequestId,
+      });
+    } catch (e) {
+      abortController.abort();
+      throw e;
+    }
 
     const chunks: Uint8Array[] = [];
 
@@ -85,6 +99,8 @@ export class LambdaHandler {
         isBase64Encoded: false,
         body: JSON.stringify({ message: error.message }),
       };
+    } finally {
+      abortController.abort();
     }
 
     return {

packages/restate-sdk/src/endpoint/node_endpoint.ts

@@ -74,15 +74,25 @@ export class NodeEndpoint implements RestateEndpoint {
 
     return (request, response) => {
       (async () => {
+        const abortController = new AbortController();
+        request.once("aborted", () => {
+          abortController.abort();
+        });
+        request.once("close", () => {
+          abortController.abort();
+        });
+        request.once("error", () => {
+          abortController.abort();
+        });
         try {
           const url = request.url;
           const resp = await handler.handle({
             url,
             headers: request.headers,
             body: Readable.toWeb(request),
             extraArgs: [],
+            abortSignal: abortController.signal,
           });
-
           response.writeHead(resp.statusCode, resp.headers);
           const responseWeb = Writable.toWeb(
             response
@@ -95,6 +105,7 @@ export class NodeEndpoint implements RestateEndpoint {
             "Error while handling connection: " + (error.stack ?? error.message)
           );
           response.destroy(error);
+          abortController.abort();
         }
       })().catch(() => {});
     };