feat: add highwatermakr option to subscription queue (#1197)

Author: simone-sanfratello

docs/api/options.md

@@ -76,6 +76,7 @@
   - `subscription.keepAlive`: `Integer` Optional interval in ms to send the `GQL_CONNECTION_KEEP_ALIVE` message.
   - `subscription.fullWsTransport`: `Boolean` Enable sending every operation via WS.
   - `subscription.wsDefaultSubprotocol`: `String` Set the default subprotocol in case the client does not specify one. See [Supported subprotocols](../graphql-over-websocket.md#websocket-subprotocol)
+  - `subscription.queueHighWaterMark`: `Number` Set the high water mark for the internal readable stream queue used by each subscription connection. This controls the buffering behavior of subscription messages. When the queue reaches this limit, backpressure is applied to prevent memory issues. Defaults to Node.js stream default (16 for object mode). See [Node.js Stream Buffering](https://nodejs.org/api/stream.html#buffering) for more details.
 
 - `persistedQueries`: A hash/query map to resolve the full query text using it's unique hash. Overrides `persistedQueryProvider`.
 - `onlyPersisted`: Boolean. Flag to control whether to allow graphql queries other than persisted. When `true`, it'll make the server reject any queries that are not present in the `persistedQueries` option above. It will also disable any ide available (graphiql). Requires `persistedQueries` to be set, and overrides `persistedQueryProvider`.

index.d.ts

@@ -537,6 +537,7 @@ declare namespace mercurius {
         keepAlive?: number,
         fullWsTransport?: boolean,
         wsDefaultSubprotocol?: string,
+        queueHighWaterMark?: number,
       };
     /**
      * Persisted queries, overrides persistedQueryProvider.

index.js

@@ -134,6 +134,7 @@ const mercurius = fp(async function (app, opts) {
   let keepAlive
   let fullWsTransport
   let wsDefaultSubprotocol
+  let queueHighWaterMark
 
   if (typeof subscriptionOpts === 'object') {
     if (subscriptionOpts.pubsub) {
@@ -149,6 +150,7 @@ const mercurius = fp(async function (app, opts) {
     keepAlive = subscriptionOpts.keepAlive
     fullWsTransport = subscriptionOpts.fullWsTransport
     wsDefaultSubprotocol = subscriptionOpts.wsDefaultSubprotocol
+    queueHighWaterMark = subscriptionOpts.queueHighWaterMark
   } else if (subscriptionOpts === true) {
     emitter = mq()
     subscriber = new PubSub(emitter)
@@ -190,6 +192,11 @@ const mercurius = fp(async function (app, opts) {
       throw new MER_ERR_INVALID_OPTS('wsDefaultSubprotocol must be either graphql-ws or graphql-transport-ws')
     }
   }
+  if (queueHighWaterMark) {
+    if (typeof queueHighWaterMark !== 'number' || queueHighWaterMark <= 0) {
+      throw new MER_ERR_INVALID_OPTS('queueHighWaterMark must be a positive number')
+    }
+  }
 
   fastifyGraphQl.schema = schema
 
@@ -226,6 +233,7 @@ const mercurius = fp(async function (app, opts) {
       keepAlive,
       fullWsTransport,
       wsDefaultSubprotocol,
+      queueHighWaterMark,
       additionalRouteOptions: opts.additionalRouteOptions,
       csrfConfig
     })

lib/routes.js

@@ -204,6 +204,7 @@ module.exports = async function (app, opts) {
     keepAlive,
     fullWsTransport,
     wsDefaultSubprotocol,
+    queueHighWaterMark,
     additionalRouteOptions
   } = opts
 
@@ -349,6 +350,7 @@ module.exports = async function (app, opts) {
       keepAlive,
       fullWsTransport,
       wsDefaultSubprotocol,
+      queueHighWaterMark,
       errorFormatter
     })
   } else {

lib/subscriber.js

@@ -37,11 +37,12 @@ class PubSub {
 
 // One context - and queue for each subscription
 class SubscriptionContext {
-  constructor ({ pubsub, fastify }) {
+  constructor ({ pubsub, fastify, queueHighWaterMark }) {
     this.fastify = fastify
     this.pubsub = pubsub
     this.queue = new Readable({
       objectMode: true,
+      highWaterMark: queueHighWaterMark,
       read: () => {}
     })
     this.closed = false

lib/subscription-connection.js

@@ -22,6 +22,7 @@ module.exports = class SubscriptionConnection {
     keepAlive,
     fullWsTransport,
     wsDefaultSubprotocol,
+    queueHighWaterMark,
     errorFormatter
   }) {
     this.fastify = fastify
@@ -38,6 +39,7 @@ module.exports = class SubscriptionConnection {
     this.keepAlive = keepAlive
     this.fullWsTransport = fullWsTransport
     this.wsDefaultSubprotocol = wsDefaultSubprotocol
+    this.queueHighWaterMark = queueHighWaterMark
     this.errorFormatter = errorFormatter
     this.headers = {}
     this.resolverContext = null
@@ -267,7 +269,8 @@ module.exports = class SubscriptionConnection {
 
     const sc = new SubscriptionContext({
       fastify: this.fastify,
-      pubsub: this.subscriber
+      pubsub: this.subscriber,
+      queueHighWaterMark: this.queueHighWaterMark
     })
 
     // Trigger preSubscriptionParsing hook

lib/subscription.js

@@ -6,7 +6,10 @@ const { kHooks } = require('./symbols')
 const SubscriptionConnection = require('./subscription-connection')
 const { isValidClientProtocol } = require('./subscription-protocol')
 
-function createConnectionHandler ({ subscriber, fastify, onConnect, onDisconnect, entityResolversFactory, subscriptionContextFn, keepAlive, fullWsTransport, wsDefaultSubprotocol, errorFormatter }) {
+function createConnectionHandler ({
+  subscriber, fastify, onConnect, onDisconnect, entityResolversFactory, subscriptionContextFn, keepAlive,
+  fullWsTransport, wsDefaultSubprotocol, queueHighWaterMark, errorFormatter
+}) {
   return async (socket, request) => {
     if (!isValidClientProtocol(socket.protocol, wsDefaultSubprotocol)) {
       console.log('wrong websocket protocol: ' + socket.protocol)
@@ -49,13 +52,14 @@ function createConnectionHandler ({ subscriber, fastify, onConnect, onDisconnect
       keepAlive,
       fullWsTransport,
       wsDefaultSubprotocol,
+      queueHighWaterMark,
       errorFormatter
     })
   }
 }
 
 module.exports = async function (fastify, opts) {
-  const { getOptions, subscriber, verifyClient, onConnect, onDisconnect, entityResolversFactory, subscriptionContextFn, keepAlive, fullWsTransport, wsDefaultSubprotocol, errorFormatter } = opts
+  const { getOptions, subscriber, verifyClient, onConnect, onDisconnect, entityResolversFactory, subscriptionContextFn, keepAlive, fullWsTransport, wsDefaultSubprotocol, queueHighWaterMark, errorFormatter } = opts
 
   // If `fastify.websocketServer` exists, it means `@fastify/websocket` already registered.
   // Without this check, @fastify/websocket will be registered multiple times and raises FST_ERR_DEC_ALREADY_PRESENT.
@@ -80,6 +84,7 @@ module.exports = async function (fastify, opts) {
       keepAlive,
       fullWsTransport,
       wsDefaultSubprotocol,
+      queueHighWaterMark,
       errorFormatter
     })
   })

test/options.test.js

@@ -98,3 +98,33 @@ test('invalid wsDefaultSubprotocol', async t => {
     message: 'Invalid options: wsDefaultSubprotocol must be either graphql-ws or graphql-transport-ws'
   })
 })
+
+test('invalid queueHighWaterMark', async t => {
+  const app = Fastify()
+  t.after(() => app.close())
+
+  app.register(mercurius, {
+    subscription: {
+      queueHighWaterMark: 'invalid'
+    }
+  })
+
+  await t.assert.rejects(app.ready(), {
+    message: 'Invalid options: queueHighWaterMark must be a positive number'
+  })
+})
+
+test('invalid queueHighWaterMark', async t => {
+  const app = Fastify()
+  t.after(() => app.close())
+
+  app.register(mercurius, {
+    subscription: {
+      queueHighWaterMark: -1
+    }
+  })
+
+  await t.assert.rejects(app.ready(), {
+    message: 'Invalid options: queueHighWaterMark must be a positive number'
+  })
+})

test/subscription.test.js

@@ -697,6 +697,70 @@ test('subscription with custom pubsub with custom params on subscribe method', (
   })
 })
 
+test('subscription queue has highWaterMark when queueHighWaterMark is provided', async (t) => {
+  const emitter = new EventEmitter()
+  class SpyPubSub {
+    async subscribe (_, queue) {
+      emitter.emit('onsubscription', queue)
+    }
+  }
+
+  const app = Fastify()
+  t.after(() => app.close())
+
+  const pubsub = new SpyPubSub()
+
+  const schema = `
+    type Query {
+      _placeholder: String
+    }
+
+    type Subscription {
+      onSubscription: String
+    }
+  `
+
+  const resolvers = {
+    Subscription: {
+      onSubscription: {
+        subscribe: (root, args, { pubsub }) => pubsub.subscribe('ON_SUBSCRIPTION')
+      }
+    }
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    subscription: {
+      pubsub,
+      queueHighWaterMark: 2
+    }
+  })
+
+  await app.listen({ port: 0 })
+
+  const ws = new WebSocket('ws://localhost:' + (app.server.address()).port + '/graphql', 'graphql-ws')
+  const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+  t.after(() => client.destroy())
+  client.setEncoding('utf8')
+
+  client.write(JSON.stringify({ type: 'connection_init' }))
+  client.write(JSON.stringify({
+    id: 1,
+    type: 'start',
+    payload: {
+      query: `
+        subscription {
+          onSubscription
+        }  
+      `
+    }
+  }))
+
+  const [queue] = await once(emitter, 'onsubscription')
+  t.assert.strictEqual(queue._readableState.highWaterMark, 2)
+})
+
 test('subscription server sends update to subscriptions with custom context', (t, done) => {
   const app = Fastify()
   t.after(() => app.close())

test/types/index.ts

@@ -362,6 +362,14 @@ app.register(mercurius, {
   }
 })
 
+app.register(mercurius, {
+  schema,
+  resolvers,
+  subscription: {
+    queueHighWaterMark: 100
+  }
+})
+
 app.register(mercurius, {
   schema,
   resolvers,