fix(server): Dispose of subscriptions on close even if added late to …

…the subscriptions list (#534)

Closes #532

src/__tests__/server.ts

@@ -1841,9 +1841,56 @@ describe('Disconnect/close', () => {
 
     client.ws.close(4321, 'Byebye');
   });
+
+  it('should dispose of subscriptions on close even if added late to the subscriptions list', async () => {
+    let resolveOnOperation: () => void = () => {
+      throw new Error('On operation resolved early');
+    };
+    const waitForOnOperation = new Promise<void>(
+      (resolve) => (resolveOnOperation = resolve),
+    );
+    let resolveOperation: () => void = () => {
+      throw new Error('Operation resolved early');
+    };
+    const { url, waitForConnect, waitForComplete, waitForClientClose } =
+      await startTServer({
+        onOperation: () => {
+          resolveOnOperation();
+          return new Promise((resolve) => (resolveOperation = resolve));
+        },
+      });
+
+    const client = await createTClient(url);
+
+    client.ws.send(
+      stringifyMessage<MessageType.ConnectionInit>({
+        type: MessageType.ConnectionInit,
+      }),
+    );
+    await waitForConnect();
+
+    client.ws.send(
+      stringifyMessage<MessageType.Subscribe>({
+        type: MessageType.Subscribe,
+        id: '1',
+        payload: {
+          query: 'subscription { ping }',
+        },
+      }),
+    );
+
+    await waitForOnOperation;
+
+    client.ws.close(4321, 'Byebye');
+    await waitForClientClose();
+
+    resolveOperation();
+
+    await waitForComplete();
+  });
 });
 
-it('should only accept a Set, Array or string in handleProtocols', () => {
+it('should only accept a Set, Array or string in handleProtocol', () => {
   for (const test of [
     {
       in: new Set(['not', 'me']),

src/server.ts

@@ -848,8 +848,9 @@ export function makeServer<
       // wait for close, cleanup and the disconnect callback
       return async (code, reason) => {
         if (connectionInitWait) clearTimeout(connectionInitWait);
-        for (const sub of Object.values(ctx.subscriptions)) {
+        for (const [id, sub] of Object.entries(ctx.subscriptions)) {
           if (isAsyncGenerator(sub)) await sub.return(undefined);
+          delete ctx.subscriptions[id]; // deleting the subscription means no further activity should take place
         }
         if (ctx.acknowledged) await onDisconnect?.(ctx, code, reason);
         await onClose?.(ctx, code, reason);