Abort handler signal on error/return (#786)

connectrpc · Aug 28, 2023 · 9dd8d56 · 9dd8d56
1 parent 3cb108f
commit 9dd8d56
Show file tree

Hide file tree

Showing 6 changed files with 87 additions and 9 deletions.
diff --git a/packages/connect/src/protocol-connect/handler-factory.spec.ts b/packages/connect/src/protocol-connect/handler-factory.spec.ts
@@ -161,11 +161,14 @@ describe("createHandlerFactory()", function () {
     it("should propogate errors back to the handler", async function () {
       let resolve: (e: unknown) => void;
       const catchError = new Promise<unknown>((r) => (resolve = r));
+      let abortResolve: () => void;
+      const abortCalled = new Promise<void>((r) => (abortResolve = r));
       const { handler } = setupTestHandler(
         testService.methods.serverStreaming,
         {},
         // eslint-disable-next-line @typescript-eslint/require-await
-        async function* (req) {
+        async function* (req, { signal }) {
+          signal.addEventListener("abort", abortResolve);
           try {
             yield { value: `${req.value}` };
             fail("expected error");
@@ -192,7 +195,8 @@ describe("createHandlerFactory()", function () {
       await it.next();
       const writeError = new Error("write error");
       await it.throw?.(writeError).catch(() => {});
-      expect(await catchError).toEqual(writeError);
+      await expectAsync(catchError).toBeResolvedTo(writeError);
+      await expectAsync(abortCalled).toBeResolved();
     });
   });
 

diff --git a/packages/connect/src/protocol-connect/handler-factory.ts b/packages/connect/src/protocol-connect/handler-factory.ts
@@ -417,8 +417,30 @@ function createStreamHandler<I extends Message<I>, O extends Message<O>>(
         // raises an error, but we want to be lenient
       ),
     );
+    const it = transformInvokeImplementation<I, O>(
+      spec,
+      context,
+    )(inputIt)[Symbol.asyncIterator]();
     const outputIt = pipe(
-      transformInvokeImplementation<I, O>(spec, context)(inputIt),
+      // We wrap the iterator in an async iterator to ensure that the
+      // abort signal is aborted when the iterator is done.
+      {
+        [Symbol.asyncIterator]() {
+          return {
+            next: () => it.next(),
+            throw: (e: unknown) => {
+              context.abort(e);
+              return it.throw?.(e) ?? Promise.reject({ done: true });
+            },
+            return: (v: unknown) => {
+              context.abort();
+              return (
+                it.return?.(v) ?? Promise.resolve({ done: true, value: v })
+              );
+            },
+          };
+        },
+      },
       transformSerializeEnvelope(serialization.getO(type.binary)),
       transformCatchFinally<EnvelopedMessage>((e) => {
         context.abort();

diff --git a/packages/connect/src/protocol-grpc-web/handler-factory.spec.ts b/packages/connect/src/protocol-grpc-web/handler-factory.spec.ts
@@ -124,11 +124,14 @@ describe("createHandlerFactory()", function () {
     it("should propogate errors back to the handler", async function () {
       let resolve: (e: unknown) => void;
       const catchError = new Promise<unknown>((r) => (resolve = r));
+      let abortResolve: () => void;
+      const abortCalled = new Promise<void>((r) => (abortResolve = r));
       const { handler } = setupTestHandler(
         testService.methods.serverStreaming,
         {},
         // eslint-disable-next-line @typescript-eslint/require-await
-        async function* (req) {
+        async function* (req, { signal }) {
+          signal.addEventListener("abort", abortResolve);
           try {
             yield { value: `${req.value}` };
             fail("expected error");
@@ -155,7 +158,8 @@ describe("createHandlerFactory()", function () {
       await it.next();
       const writeError = new Error("write error");
       await it.throw?.(writeError).catch(() => {});
-      expect(await catchError).toEqual(writeError);
+      await expectAsync(catchError).toBeResolvedTo(writeError);
+      await expectAsync(abortCalled).toBeResolved();
     });
   });
 

diff --git a/packages/connect/src/protocol-grpc-web/handler-factory.ts b/packages/connect/src/protocol-grpc-web/handler-factory.ts
@@ -169,8 +169,30 @@ function createHandler<I extends Message<I>, O extends Message<O>>(
         // raises an error, but we want to be lenient
       ),
     );
+    const it = transformInvokeImplementation<I, O>(
+      spec,
+      context,
+    )(inputIt)[Symbol.asyncIterator]();
     const outputIt = pipe(
-      transformInvokeImplementation<I, O>(spec, context)(inputIt),
+      // We wrap the iterator in an async iterator to ensure that the
+      // abort signal is aborted when the iterator is done.
+      {
+        [Symbol.asyncIterator]() {
+          return {
+            next: () => it.next(),
+            throw: (e: unknown) => {
+              context.abort(e);
+              return it.throw?.(e) ?? Promise.reject({ done: true });
+            },
+            return: (v: unknown) => {
+              context.abort();
+              return (
+                it.return?.(v) ?? Promise.resolve({ done: true, value: v })
+              );
+            },
+          };
+        },
+      },
       transformSerializeEnvelope(serialization.getO(type.binary)),
       transformCatchFinally<EnvelopedMessage>((e) => {
         context.abort();

diff --git a/packages/connect/src/protocol-grpc/handler-factory.spec.ts b/packages/connect/src/protocol-grpc/handler-factory.spec.ts
@@ -124,11 +124,14 @@ describe("createHandlerFactory()", function () {
     it("should propogate errors back to the handler", async function () {
       let resolve: (e: unknown) => void;
       const catchError = new Promise<unknown>((r) => (resolve = r));
+      let abortResolve: () => void;
+      const abortCalled = new Promise<void>((r) => (abortResolve = r));
       const { handler } = setupTestHandler(
         testService.methods.serverStreaming,
         {},
         // eslint-disable-next-line @typescript-eslint/require-await
-        async function* (req) {
+        async function* (req, { signal }) {
+          signal.addEventListener("abort", abortResolve);
           try {
             yield { value: `${req.value}` };
             fail("expected error");
@@ -155,7 +158,8 @@ describe("createHandlerFactory()", function () {
       await it.next();
       const writeError = new Error("write error");
       await it.throw?.(writeError).catch(() => {});
-      expect(await catchError).toEqual(writeError);
+      await expectAsync(catchError).toBeResolvedTo(writeError);
+      await expectAsync(abortCalled).toBeResolved();
     });
   });
 

diff --git a/packages/connect/src/protocol-grpc/handler-factory.ts b/packages/connect/src/protocol-grpc/handler-factory.ts
@@ -155,8 +155,30 @@ function createHandler<I extends Message<I>, O extends Message<O>>(
       transformDecompressEnvelope(compression.request, opt.readMaxBytes),
       transformParseEnvelope(serialization.getI(type.binary)),
     );
+    const it = transformInvokeImplementation<I, O>(
+      spec,
+      context,
+    )(inputIt)[Symbol.asyncIterator]();
     const outputIt = pipe(
-      transformInvokeImplementation<I, O>(spec, context)(inputIt),
+      // We wrap the iterator in an async iterator to ensure that the
+      // abort signal is aborted when the iterator is done.
+      {
+        [Symbol.asyncIterator]() {
+          return {
+            next: () => it.next(),
+            throw: (e: unknown) => {
+              context.abort(e);
+              return it.throw?.(e) ?? Promise.reject({ done: true });
+            },
+            return: (v: unknown) => {
+              context.abort();
+              return (
+                it.return?.(v) ?? Promise.resolve({ done: true, value: v })
+              );
+            },
+          };
+        },
+      },
       transformSerializeEnvelope(serialization.getO(type.binary)),
       transformCompressEnvelope(compression.response, opt.compressMinBytes),
       transformJoinEnvelopes(),