Repository: mercurius-js/mercurius (js)

Issue:
Closing out subscription does not execute return() on async iterator from subscription resolver
Hi

I am working on a new application and looking to use mercurius to serve up graphql subscriptions.  During development, I noticed that if I close the the subscription client either abruptly or gracefully by sending a `GQL_STOP` the underlying async iterator from the subscription resolver will keep running.  In my working program (unable to share this particular code as it is  private code from my work) - I am delegating to another remote graphql subscription running in a different service.  So in this case these delegated subscriptions are staying active within my mercurius application to the remote server even though the client has disconnected.  As I am unable to show that code, I have put together a small project which demonstrates this

https://github.com/moonkev/mercurius-sub-repro

If you run src/mercurius.js it will startup a server on port 4500.  
query - http://localhost:4500
subs - ws://localhost:4500

You can make a subscription using this very simple example (using any client, i.e. Altair, graphql-playground, etc).
```
subscription {
  messages { id }
}
```

If from your client you disconnect or try to unsubscribe gracefully, you will notice from the logging that the iterator continues to run.  I have included a counter-example using express-graphql.  You can start this by running src/express.js.  It will start on the same URLs as the mercurius example.  If you close the client connection here, you will see that the async iterators return method is called and the iterator will stop producing data.

The iterator itself is defined in src/schema.js.  Apologies if there is just something that I have misconfigured that can make this work.  I didn't see anything in the documentation about it.  Please let me know if any further information is needed, Thanks!

Pull request description:
Call return on subscription resolver async iterator
fixes #392 

Patch:
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index 51740bf..09cf1a7 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -37,6 +37,7 @@ module.exports = class SubscriptionConnection {
     this.onConnect = onConnect
     this.onDisconnect = onDisconnect
     this.subscriptionContexts = new Map()
+    this.subscriptionIters = new Map()
     this.context = context
     this.isReady = false
     this.resolveContext = resolveContext
@@ -165,12 +166,12 @@ module.exports = class SubscriptionConnection {
       fastify: this.fastify,
       pubsub: this.subscriber
     })
-    this.subscriptionContexts.set(id, sc)
 
     const document = typeof query !== 'string' ? query : parse(query)
     const schema = this.fastify ? this.fastify.graphql.schema : undefined
+    this.subscriptionContexts.set(id, sc)
 
-    const result = await subscribe(
+    const subIter = await subscribe(
       schema,
       document,
       {}, // rootValue
@@ -186,14 +187,15 @@ module.exports = class SubscriptionConnection {
       variables,
       operationName
     )
+    this.subscriptionIters.set(id, subIter)
 
-    if (result.errors) {
-      this.fastify.log.error(result.errors)
-      throw result.errors[0]
+    if (subIter.errors) {
+      this.fastify.log.error(subIter.errors)
+      throw subIter.errors[0]
     }
 
     // TODO implement backpressure
-    for await (const value of result) {
+    for await (const value of subIter) {
       this.sendMessage(GQL_DATA, data.id, value)
     }
 
@@ -201,20 +203,25 @@ module.exports = class SubscriptionConnection {
   }
 
   handleGQLStop (data) {
-    const subscriptionContextToRemove = this.subscriptionContexts.get(data.id)
-    if (!subscriptionContextToRemove) {
-      return
+    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)
     }
-
-    subscriptionContextToRemove.close()
-    this.subscriptionContexts.delete(data.id)
   }
 
   handleConnectionClose () {
     Array
       .from(this.subscriptionContexts.values())
-      .map(subscriptionContext =>
-        subscriptionContext.close())
+      .map((sc) => sc.close())
+    Array
+      .from(this.subscriptionIters.values())
+      .map((subIter) => subIter.return && subIter.return())
     this.socket.close()
 
     if (typeof this.onDisconnect === 'function') {


Test patch:
diff --git a/test/subscription-connection.js b/test/subscription-connection.js
index 2c493b3..01a5811 100644
--- a/test/subscription-connection.js
+++ b/test/subscription-connection.js
@@ -92,7 +92,7 @@ test('subscription connection handles GQL_CONNECTION_TERMINATE message correctly
   }))
 })
 
-test('subscription connection handles GQL_STOP message correctly', async (t) => {
+test('subscription connection closes context on GQL_STOP message correctly', async (t) => {
   t.plan(2)
   const sc = new SubscriptionConnection({
     on () {},
@@ -115,6 +115,29 @@ test('subscription connection handles GQL_STOP message correctly', async (t) =>
   t.equal(sc.subscriptionContexts.size, 0)
 })
 
+test('subscription connection completes resolver iterator on GQL_STOP message correctly', async (t) => {
+  t.plan(2)
+  const sc = new SubscriptionConnection({
+    on () {},
+    close () {},
+    send (message) {}
+  }, {})
+
+  sc.subscriptionIters = new Map()
+  sc.subscriptionIters.set(1, {
+    return () {
+      t.pass()
+    }
+  })
+
+  await sc.handleMessage(JSON.stringify({
+    id: 1,
+    type: 'stop'
+  }))
+
+  t.equal(sc.subscriptionIters.size, 0)
+})
+
 test('handles error in send and closes connection', async t => {
   t.plan(1)