Always teardown subscriptions on completion (#846)

* Always teardown GQL subscriptions on completion * format + unit test
mercurius-js · Aug 24, 2022 · e7baf2b · e7baf2b
1 parent 728beed
commit e7baf2b
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 10 deletions.
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
@@ -271,6 +271,7 @@ module.exports = class SubscriptionConnection {
     }
 
     this.sendMessage(this.protocolMessageTypes.GQL_COMPLETE, id, null)
+    this.handleGQLComplete(id)
   }
 
   async _executeQueryOrMutation ({ query, context, variables, operationName, id }) {
@@ -284,6 +285,19 @@ module.exports = class SubscriptionConnection {
     this.sendMessage(this.protocolMessageTypes.GQL_COMPLETE, id, null)
   }
 
+  handleGQLComplete (id) {
+    const sc = this.subscriptionContexts.get(id)
+    if (sc) {
+      sc.close && sc.close()
+      this.subscriptionContexts.delete(id)
+    }
+    const subIter = this.subscriptionIters.get(id)
+    if (subIter) {
+      subIter.return && subIter.return()
+      this.subscriptionIters.delete(id)
+    }
+  }
+
   async handleGQLStop (data) {
     if (this.context.onSubscriptionEnd) {
       try {
@@ -293,16 +307,8 @@ module.exports = class SubscriptionConnection {
         return this.handleConnectionClose()
       }
     }
-    const sc = this.subscriptionContexts.get(data.id)
-    if (sc) {
-      sc.close && sc.close()
-      this.subscriptionContexts.delete(data.id)
-    }
-    const subIter = this.subscriptionIters.get(data.id)
-    if (subIter) {
-      subIter.return && subIter.return()
-      this.subscriptionIters.delete(data.id)
-    }
+
+    this.handleGQLComplete(data.id)
   }
 
   handleConnectionClose () {

diff --git a/test/subscription-connection.js b/test/subscription-connection.js
@@ -5,6 +5,7 @@ const sinon = require('sinon')
 const WebSocket = require('ws')
 const fastify = require('fastify')
 const mq = require('mqemitter')
+const { makeExecutableSchema } = require('@graphql-tools/schema')
 const SubscriptionConnection = require('../lib/subscription-connection')
 const { PubSub } = require('../lib/subscriber')
 const { GRAPHQL_WS, GRAPHQL_TRANSPORT_WS } = require('../lib/subscription-protocol')
@@ -643,3 +644,60 @@ test('subscription connection handles mutation when fullWsTransport: true', asyn
     ).calledOnce
   )
 })
+
+test('subscription data is released right after it ends', async (t) => {
+  const sc = new SubscriptionConnection({
+    on () {},
+    close () {},
+    send () {},
+    protocol: GRAPHQL_TRANSPORT_WS
+  }, {
+    context: { preSubscriptionParsing: null, preSubscriptionExecution: null },
+    fastify: {
+      graphql: {
+        schema: makeExecutableSchema({
+          typeDefs: ['type Query { blah: String! }', 'type Subscription { onMessage: String! }'],
+          resolvers: {
+            Query: {},
+            Subscription: {
+              onMessage: {
+                async * subscribe () {
+                  return 'blah'
+                }
+              }
+            }
+          }
+        })
+      }
+    }
+  })
+
+  sc.isReady = true
+
+  t.equal(sc.subscriptionIters.size, 0)
+  t.equal(sc.subscriptionContexts.size, 0)
+
+  await sc.handleMessage(JSON.stringify({
+    id: 1,
+    type: 'subscribe',
+    payload: {
+      query: 'subscription { onMessage } '
+    }
+  }))
+
+  await new Promise(resolve => {
+    sc.sendMessage = (type, id, payload) => {
+      t.equal(id, 1)
+      t.equal(type, 'complete')
+      t.equal(payload, null)
+
+      t.equal(sc.subscriptionIters.size, 1)
+      t.equal(sc.subscriptionContexts.size, 1)
+
+      resolve()
+    }
+  })
+
+  t.equal(sc.subscriptionIters.size, 0)
+  t.equal(sc.subscriptionContexts.size, 0)
+})