From 4e30dfeddaf176daea087585165854d0c67c543b Mon Sep 17 00:00:00 2001
From: tm-kwon <tm.kwon@naverlabs.com>
Date: Tue, 14 Feb 2023 23:22:42 +0900
Subject: [PATCH] fix: ensure asyncIterator call return() when fail at calling
 next()

---
 src/__tests__/server.ts | 38 ++++++++++++++++++++++++++++++++++++++
 src/server.ts           |  4 ++++
 2 files changed, 42 insertions(+)

diff --git a/src/__tests__/server.ts b/src/__tests__/server.ts
index ea2b1de8..c2b71a27 100644
--- a/src/__tests__/server.ts
+++ b/src/__tests__/server.ts
@@ -1740,6 +1740,44 @@ describe('Subscribe', () => {
       fail("Shouldn't have received a message");
     }, 20);
   });
+
+  it('should execute "return" of the subscription when error occurs in "next" call', async (done) => {
+    const iterator = (async function* () {
+      yield { data: { greetings: 'Hi' } };
+      throw new Error('error');
+    })();
+
+    jest.spyOn(iterator, 'return').mockImplementation(done);
+
+    const { url } = await startTServer({
+      schema,
+      subscribe: () => iterator,
+    });
+
+    const client = await createTClient(url);
+    client.ws.send(
+      stringifyMessage<MessageType.ConnectionInit>({
+        type: MessageType.ConnectionInit,
+      }),
+    );
+
+    await client.waitForMessage(({ data }) => {
+      expect(parseMessage(data).type).toBe(MessageType.ConnectionAck);
+      client.ws.send(
+        stringifyMessage<MessageType.Subscribe>({
+          id: '1',
+          type: MessageType.Subscribe,
+          payload: {
+            operationName: 'Greetings',
+            query: `subscription Greetings {
+                greetings
+              }`,
+            variables: {},
+          },
+        }),
+      );
+    });
+  });
 });
 
 describe('Disconnect/close', () => {
diff --git a/src/server.ts b/src/server.ts
index cc29f323..9e22600e 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -824,6 +824,10 @@ export function makeServer<
               // completed the subscription, he doesnt need to be reminded
               await emit.complete(id in ctx.subscriptions);
             } finally {
+              // before delete the subcription, we should call `return` of asyncIterator to prevent it from hanging
+              const sub = ctx.subscriptions[id];
+              if (isAsyncGenerator(sub)) await sub.return(undefined);
+
               // whatever happens to the subscription, we finally want to get rid of the reservation
               delete ctx.subscriptions[id];
             }