Repository: mcollina/fastify-gql
Pull request #137 (2020-04-16 16:32:21)

Subscription support in gateway
This PR enables subscriptions when the server runs in gateway mode.

Todo:
- [x] Enable WebSocket connection to the gateway
- [x] Forward subscription to correct service
- [x] Forward received data from service to all clients
- [x] Generate subscription query
- [x] Support full query execution after data received from service and pushed to clients
- [x] Websocket client updates in gateway

mcollina: Good work!

Patch:
diff --git a/README.md b/README.md
index de0d422..f8b3fc6 100644
--- a/README.md
+++ b/README.md
@@ -431,7 +431,6 @@ In Gateway mode the following options are not allowed (the plugin will throw an
 * `schema`
 * `resolvers`
 * `loaders`
-* `subscription`
 
 Also, using the following decorator methods will throw:
 * `app.graphql.defineResolvers`
@@ -503,6 +502,8 @@ __fastify-gql__ supports the following options:
     * `service.name`: A unique name for the service. Required.
     * `service.url`: The url of the service endpoint. Required
     * `service.rewriteHeaders`: `Function` A function that gets the original headers as a parameter and returns an object containing values that should be added to the headers
+    * `service.wsUrl`: The url of the websocket endpoint
+    * `service.wsConnectionParams`: `Function` or `Object`
 
 #### queryDepth example
 ```
diff --git a/examples/gateway-subscription.js b/examples/gateway-subscription.js
new file mode 100644
index 0000000..dcb7e84
--- /dev/null
+++ b/examples/gateway-subscription.js
@@ -0,0 +1,278 @@
+'use strict'
+const Fastify = require('fastify')
+const GQL = require('..')
+
+async function createService (port, schema, resolvers = {}) {
+  const service = Fastify()
+
+  service.register(GQL, {
+    schema,
+    resolvers,
+    federationMetadata: true,
+    ide: 'playground',
+    routes: true,
+    jit: 1,
+    subscription: true
+  })
+  await service.listen(port)
+}
+
+const users = {
+  u1: {
+    id: 'u1',
+    name: 'John'
+  },
+  u2: {
+    id: 'u2',
+    name: 'Jane'
+  },
+  u3: {
+    id: 'u3',
+    name: 'Jack'
+  }
+}
+
+const posts = {
+  p1: {
+    pid: 'p1',
+    title: 'Post 1',
+    content: 'Content 1',
+    authorId: 'u1'
+  },
+  p2: {
+    pid: 'p2',
+    title: 'Post 2',
+    content: 'Content 2',
+    authorId: 'u2'
+  },
+  p3: {
+    pid: 'p3',
+    title: 'Post 3',
+    content: 'Content 3',
+    authorId: 'u1'
+  },
+  p4: {
+    pid: 'p4',
+    title: 'Post 4',
+    content: 'Content 4',
+    authorId: 'u2'
+  }
+}
+
+const comments = {}
+
+async function start () {
+  await createService(4001, `
+    extend type Query {
+      me: User
+    }
+
+    type User @key(fields: "id") {
+      id: ID!
+      name: String!
+      fullName: String
+      avatar(size: AvatarSize): String
+      friends: [User]
+    }
+
+    enum AvatarSize {
+      small
+      medium
+      large
+    }
+  `, {
+    Query: {
+      me: (root, args, context, info) => {
+        return users.u1
+      }
+    },
+    User: {
+      __resolveReference: (user, args, context, info) => {
+        return users[user.id]
+      },
+      avatar: (user, { size }) => `avatar-${size}.jpg`,
+      friends: (user) => Object.values(users).filter(u => u.id !== user.id),
+      fullName: (user) => user.name + ' Doe'
+    }
+  })
+
+  await createService(4002, `
+    type Post @key(fields: "pid") {
+      pid: ID!
+      title: String
+      content: String
+      author: User @requires(fields: "pid title")
+    }
+
+    extend type Query {
+      topPosts(count: Int): [Post]
+    }
+
+    extend type User @key(fields: "id") {
+      id: ID! @external
+      name: String @external
+      posts: [Post]
+      numberOfPosts: Int @requires(fields: "id name")
+    }
+
+    extend type Mutation {
+      createPost(post: PostInput!): Post
+    }
+
+    input PostInput {
+      title: String!
+      content: String!
+      authorId: String!
+    }
+  `, {
+    Post: {
+      __resolveReference: (post, args, context, info) => {
+        return posts[post.pid]
+      },
+      author: (post, args, context, info) => {
+        return {
+          __typename: 'User',
+          id: post.authorId
+        }
+      }
+    },
+    User: {
+      posts: (user, args, context, info) => {
+        return Object.values(posts).filter(p => p.authorId === user.id)
+      },
+      numberOfPosts: (user) => {
+        return Object.values(posts).filter(p => p.authorId === user.id).length
+      }
+    },
+    Query: {
+      topPosts: (root, { count = 2 }) => Object.values(posts).slice(0, count)
+    },
+    Mutation: {
+      createPost: (root, { post }) => {
+        const pid = `p${Object.values(posts).length + 1}`
+
+        const result = {
+          pid,
+          ...post
+        }
+        posts[pid] = result
+
+        return result
+      }
+    }
+  })
+
+  await createService(4003, `
+    type Comment @key(fields: "cid") {
+      cid: ID!
+      text: String!
+      author: User
+      post: Post
+    }
+
+    extend type User @key(fields: "id") {
+      id: ID! @external
+      comments: [Comment]
+    }
+
+    extend type Post @key(fields: "pid") {
+      pid: ID! @external
+      comments: [Comment]
+    }
+
+    input CommentInput {
+      text: String!
+      authorId: ID!
+      postId: ID!
+    }
+
+    extend type Mutation {
+      addComment(comment: CommentInput): Comment
+    }
+
+    extend type Subscription {
+      commentAdded(postId: ID!): Comment
+    }
+  `, {
+    Comment: {
+      __resolveReference: (comment) => {
+        return comments[comment.id]
+      },
+      author: (comment) => {
+        return {
+          __typename: 'User',
+          id: comment.authorId
+        }
+      },
+      post: (comment) => {
+        return {
+          __typename: 'Post',
+          pid: comment.postId
+        }
+      }
+    },
+    Post: {
+      comments: (post) => {
+        return Object.values(comments).filter(c => post.pid === c.postId)
+      }
+    },
+    User: {
+      comments: (user) => {
+        return Object.values(comments).filter(c => user.id === c.authorId)
+      }
+    },
+    Mutation: {
+      async addComment (parent, { comment }, { pubsub }) {
+        const cid = `c${Object.values(comments).length + 1}`
+
+        const result = {
+          cid,
+          ...comment
+        }
+        comments[cid] = result
+
+        await pubsub.publish({
+          topic: `COMMENT_ADDED_${comment.postId}`,
+          payload: {
+            commentAdded: result
+          }
+        })
+        return result
+      }
+    },
+    Subscription: {
+      commentAdded: {
+        subscribe: async (root, { postId }, { pubsub }) => {
+          const subscription = await pubsub.subscribe(`COMMENT_ADDED_${postId}`)
+
+          return subscription
+        }
+      }
+    }
+  })
+
+  const gateway = Fastify()
+  gateway.register(GQL, {
+    routes: true,
+    ide: 'playground',
+    subscription: true,
+    jit: 1,
+    gateway: {
+      services: [{
+        name: 'user',
+        url: 'http://localhost:4001/graphql'
+      }, {
+        name: 'post',
+        url: 'http://localhost:4002/graphql'
+      }, {
+        name: 'comment',
+        url: 'http://localhost:4003/graphql',
+        wsUrl: 'ws://localhost:4003/graphql'
+      }]
+    }
+  })
+
+  await gateway.listen(4000)
+}
+
+start()
diff --git a/index.js b/index.js
index 65ac9b0..69dd77c 100644
--- a/index.js
+++ b/index.js
@@ -24,6 +24,8 @@ const {
 const queryDepth = require('./lib/queryDepth')
 const buildFederationSchema = require('./lib/federation')
 const buildGateway = require('./lib/gateway')
+const mq = require('mqemitter')
+const { PubSub } = require('./lib/subscriber')
 
 const kLoaders = Symbol('fastify-gql.loaders')
 
@@ -62,9 +64,25 @@ module.exports = fp(async function (app, opts) {
   const root = {}
   let schema = opts.schema
   let gateway = opts.gateway
+  const subscriptionOpts = opts.subscription
+  let emitter
 
-  if (gateway && (schema || opts.resolvers || opts.loaders || opts.subscription)) {
-    throw new Error('Adding "schema", "resolvers", "loaders" or "subscription" to plugin options when plugin is running in gateway mode is not allowed')
+  let subscriber
+  let verifyClient
+
+  if (typeof subscriptionOpts === 'object') {
+    emitter = subscriptionOpts.emitter || mq()
+    verifyClient = subscriptionOpts.verifyClient
+  } else if (subscriptionOpts === true) {
+    emitter = mq()
+  }
+
+  if (subscriptionOpts) {
+    subscriber = new PubSub(emitter)
+  }
+
+  if (gateway && (schema || opts.resolvers || opts.loaders)) {
+    throw new Error('Adding "schema", "resolvers", "loaders" or to plugin options when plugin is running in gateway mode is not allowed')
   }
 
   if (typeof schema === 'string') {
@@ -86,10 +104,12 @@ module.exports = fp(async function (app, opts) {
     })
   }
 
+  let entityResolversFactory
   if (gateway) {
     gateway = await buildGateway(gateway, app)
 
     schema = gateway.schema
+    entityResolversFactory = gateway.entityResolversFactory
 
     app.onClose((fastify, next) => {
       gateway.close()
@@ -117,9 +137,10 @@ module.exports = fp(async function (app, opts) {
       path: opts.path,
       context: opts.context,
       schema,
-      subscription: opts.subscription,
-      federationMetadata: opts.federationMetadata,
-      gateway
+      subscriber,
+      verifyClient,
+      lruGatewayResolvers,
+      entityResolversFactory
     })
   }
 
diff --git a/lib/gateway.js b/lib/gateway.js
index 922ea04..94ce8dd 100644
--- a/lib/gateway.js
+++ b/lib/gateway.js
@@ -13,7 +13,7 @@ const {
   createQueryOperation,
   createFieldResolverOperation,
   createEntityReferenceResolverOperation,
-  kEntityResolvels
+  kEntityResolvers
 } = require('./gateway/make-resolver')
 
 function isDefaultType (type) {
@@ -85,12 +85,22 @@ function defineResolvers (schema, typeToServiceMap, serviceMap, typeFieldsToServ
 
           if (serviceForFieldType !== serviceForType) {
             if (serviceForType === null) {
-              field.resolve = makeResolver({
-                service: serviceMap[serviceForFieldType],
-                createOperation: createQueryOperation,
-                transformData: response => response.json.data[fieldName],
-                isQuery: true
-              })
+              // Either query, mutation, subscription
+              if (type.name === 'Subscription') {
+                field.subscribe = makeResolver({
+                  service: serviceMap[serviceForFieldType],
+                  createOperation: createQueryOperation,
+                  isQuery: true,
+                  isSubscription: true
+                })
+              } else {
+                field.resolve = makeResolver({
+                  service: serviceMap[serviceForFieldType],
+                  createOperation: createQueryOperation,
+                  transformData: response => response.json.data[fieldName],
+                  isQuery: true
+                })
+              }
             } else {
               // check if the field is default field of the type
               if (serviceMap[serviceForType].typeMap[type.name].has(fieldName)) {
@@ -133,9 +143,9 @@ async function buildGateway (gatewayOpts, app) {
   const typeToServiceMap = {}
   const typeFieldsToService = {}
   const factory = new Factory()
-  app.decorateReply(kEntityResolvels)
+  app.decorateReply(kEntityResolvers)
   app.addHook('onRequest', async function (req, reply) {
-    reply[kEntityResolvels] = factory.create()
+    reply[kEntityResolvers] = factory.create()
   })
 
   for (const [service, serviceDefinition] of Object.entries(serviceMap)) {
@@ -253,12 +263,14 @@ async function buildGateway (gatewayOpts, app) {
 
   typeToServiceMap.Query = null
   typeToServiceMap.Mutation = null
+  typeToServiceMap.Subscription = null
 
   defineResolvers(schema, typeToServiceMap, serviceMap, typeFieldsToService)
 
   return {
     schema,
     serviceMap,
+    entityResolversFactory: factory,
     close () {
       for (const service of Object.values(serviceMap)) {
         service.close()
diff --git a/lib/gateway/make-resolver.js b/lib/gateway/make-resolver.js
index 867781b..155c72a 100644
--- a/lib/gateway/make-resolver.js
+++ b/lib/gateway/make-resolver.js
@@ -7,7 +7,7 @@ const {
   Kind
 } = require('graphql')
 
-const kEntityResolvels = Symbol('fastify-gql.entity-resolvers')
+const kEntityResolvers = Symbol('fastify-gql.entity-resolvers')
 
 function getFieldType (schema, type, fieldName) {
   return getNamedType(schema.getType(type).getFields()[fieldName].type)
@@ -304,7 +304,7 @@ function generatePathKey (path) {
  *  - Field entity resolver: when the field was added through type extension in the service of the field's type
  *
  */
-function makeResolver ({ service, createOperation, transformData, isQuery, isReference }) {
+function makeResolver ({ service, createOperation, transformData, isQuery, isReference, isSubscription }) {
   return function (parent, args, context, info) {
     const {
       fieldNodes,
@@ -317,7 +317,7 @@ function makeResolver ({ service, createOperation, transformData, isQuery, isRef
       schema
     } = info
     const resolverKey = generatePathKey(info.path).join('.').replace(/\d/g, '_IDX_')
-    const { reply, __currentQuery, lruGatewayResolvers } = context
+    const { reply, __currentQuery, lruGatewayResolvers, pubsub } = context
 
     const cached = lruGatewayResolvers.get(`${__currentQuery}_${resolverKey}`)
     let variableNamesToDefine
@@ -382,13 +382,18 @@ function makeResolver ({ service, createOperation, transformData, isQuery, isRef
 
     if (isReference) {
       variables.representations = [removeNonIdProperties(parent[fieldName], type)]
-    } else if (!isQuery) {
+    } else if (!isQuery && !isSubscription) {
       variables.representations = [{
         ...removeNonIdProperties(parent, parentType),
         ...getRequiredFields(parent, schema.getType(parentType).getFields()[fieldName])
       }]
     }
 
+    if (isSubscription) {
+      const subscriptionId = service.createSubscription(query, variables, pubsub.publish.bind(pubsub))
+      return pubsub.subscribe(`${service.name}_${subscriptionId}`)
+    }
+
     if (isQuery) {
       return service.sendRequest({
         method: 'POST',
@@ -398,13 +403,13 @@ function makeResolver ({ service, createOperation, transformData, isQuery, isRef
         }),
         originalRequestHeaders: reply.request.headers
       }).then(transformData)
-    } else {
-      return reply[kEntityResolvels][`${service.name}Entity`]({
-        query,
-        variables,
-        originalRequestHeaders: reply.request.headers
-      }).then(transformData)
     }
+
+    return reply[kEntityResolvers][`${service.name}Entity`]({
+      query,
+      variables,
+      originalRequestHeaders: reply.request.headers
+    }).then(transformData)
   }
 }
 
@@ -446,5 +451,5 @@ module.exports = {
   createQueryOperation,
   createFieldResolverOperation,
   createEntityReferenceResolverOperation,
-  kEntityResolvels
+  kEntityResolvers
 }
diff --git a/lib/gateway/service-map.js b/lib/gateway/service-map.js
index f3d8d23..7ac3d9a 100644
--- a/lib/gateway/service-map.js
+++ b/lib/gateway/service-map.js
@@ -8,6 +8,7 @@ const {
 } = require('graphql')
 
 const { buildRequest, sendRequest } = require('./request')
+const SubscriptionClient = require('../subscription-client')
 
 function hasExternalDirective (field) {
   return field.directives.some(directive => directive.name.value === 'external')
@@ -84,6 +85,14 @@ async function buildServiceMap (services) {
       }
     }
 
+    if (service.wsUrl) {
+      const client = new SubscriptionClient(service.wsUrl, {
+        serviceName: service.name
+      })
+
+      serviceConfig.createSubscription = client.createSubscription.bind(client)
+    }
+
     serviceMap[service.name] = serviceConfig
   }
 
diff --git a/lib/routes.js b/lib/routes.js
index b197647..2782167 100644
--- a/lib/routes.js
+++ b/lib/routes.js
@@ -4,8 +4,6 @@ const { join } = require('path')
 const Static = require('fastify-static')
 const { BadRequest } = require('http-errors')
 const { formatError, GraphQLError } = require('graphql')
-const mq = require('mqemitter')
-const { PubSub } = require('./subscriber')
 const subscription = require('./subscription')
 
 const responseSchema = {
@@ -90,23 +88,13 @@ module.exports = async function (app, opts) {
   }
   const contextFn = opts.context
 
-  const subscriptionOpts = opts.subscription
-  let emitter
-  let subscriber
-  let verifyClient
-
-  if (typeof subscriptionOpts === 'object') {
-    emitter = subscriptionOpts.emitter || mq()
-    verifyClient = subscriptionOpts.verifyClient
-  } else if (subscriptionOpts === true) {
-    emitter = mq()
-  }
-
-  if (subscriptionOpts) {
-    subscriber = new PubSub(emitter)
-  }
-
-  const { path: graphqlPath = '/graphql' } = opts
+  const {
+    path: graphqlPath = '/graphql',
+    subscriber,
+    verifyClient,
+    lruGatewayResolvers,
+    entityResolversFactory
+  } = opts
 
   const getOptions = {
     url: graphqlPath,
@@ -157,12 +145,14 @@ module.exports = async function (app, opts) {
     }
   }
 
-  if (subscriptionOpts) {
+  if (subscriber) {
     app.register(subscription, {
       getOptions,
       schema: opts.schema,
       subscriber,
-      verifyClient
+      verifyClient,
+      lruGatewayResolvers,
+      entityResolversFactory
     })
   } else {
     app.route(getOptions)
diff --git a/lib/subscription-client.js b/lib/subscription-client.js
new file mode 100644
index 0000000..63b91ce
--- /dev/null
+++ b/lib/subscription-client.js
@@ -0,0 +1,240 @@
+'use strict'
+
+const WebSocket = require('ws')
+const {
+  GQL_CONNECTION_INIT,
+  GQL_CONNECTION_ACK,
+  GQL_START,
+  GQL_DATA,
+  GQL_COMPLETE,
+  GQL_STOP,
+  GRAPHQL_WS
+} = require('./subscription-protocol')
+
+class SubscriptionClient {
+  constructor (uri, config) {
+    this.uri = uri
+    this.socket = null
+    this.operationId = 0
+    this.ready = false
+    this.operations = new Map()
+    this.operationsCount = {}
+    this.subscriptionQueryMap = {}
+    const {
+      protocols = [],
+      reconnect,
+      maxReconnectAttempts = Infinity,
+      serviceName,
+      connectionCallback,
+      failedReconnectCallback
+    } = config
+
+    this.protocols = [GRAPHQL_WS, ...protocols]
+    this.tryReconnect = reconnect
+    this.maxReconnectAttempts = maxReconnectAttempts
+    this.serviceName = serviceName
+    this.reconnectAttempts = 0
+    this.connectionCallback = connectionCallback
+    this.failedReconnectCallback = failedReconnectCallback
+
+    this.connect()
+  }
+
+  connect () {
+    this.socket = new WebSocket(this.uri, this.protocols)
+
+    this.socket.onopen = () => {
+      /* istanbul ignore else */
+      if (this.socket && this.socket.readyState === WebSocket.OPEN) {
+        this.sendMessage(null, GQL_CONNECTION_INIT)
+      }
+    }
+
+    this.socket.onclose = () => {
+      if (!this.closedByUser) {
+        this.close(this.tryReconnect, false)
+      }
+    }
+
+    this.socket.onerror = () => {}
+
+    this.socket.onmessage = ({ data }) => {
+      this.handleMessage(data)
+    }
+  }
+
+  close (tryReconnect, closedByUser = true) {
+    this.closedByUser = closedByUser
+    this.isReady = false
+
+    if (this.socket !== null) {
+      if (closedByUser) {
+        this.unsubscribeAll()
+      }
+
+      this.socket.close()
+      this.socket = null
+      this.reconnecting = false
+
+      if (tryReconnect) {
+        for (const operationId of this.operations.keys()) {
+          const { options, handler } = this.operations.get(operationId)
+
+          this.operations.set(operationId, {
+            options,
+            handler,
+            started: false
+          })
+        }
+
+        this.reconnect()
+      }
+    }
+  }
+
+  getReconnectDelay () {
+    const delayMs = 100 * Math.pow(2, this.reconnectAttempts)
+
+    return Math.min(delayMs, 10000)
+  }
+
+  reconnect () {
+    if (
+      this.reconnecting ||
+      this.reconnectAttempts > this.maxReconnectAttempts
+    ) {
+      return this.failedReconnectCallback && this.failedReconnectCallback()
+    }
+
+    this.reconnectAttempts++
+    this.reconnecting = true
+
+    const delay = this.getReconnectDelay()
+
+    this.reconnectTimeoutId = setTimeout(() => {
+      this.connect()
+    }, delay)
+  }
+
+  unsubscribe (operationId, forceUnsubscribe) {
+    let count = this.operationsCount[operationId]
+    count--
+
+    if (count === 0 || forceUnsubscribe) {
+      this.sendMessage(operationId, GQL_STOP, null)
+      this.operationsCount[operationId] = 0
+    } else {
+      this.operationsCount[operationId] = count
+    }
+  }
+
+  unsubscribeAll () {
+    for (const operationId of this.operations.keys()) {
+      this.unsubscribe(operationId, true)
+    }
+  }
+
+  sendMessage (operationId, type, payload) {
+    this.socket.send(
+      JSON.stringify({
+        id: operationId,
+        type,
+        payload
+      })
+    )
+  }
+
+  handleMessage (message) {
+    let data
+    let operationId
+    let operation
+
+    try {
+      data = JSON.parse(message)
+      operationId = data.id
+    } catch (e) {
+      /* istanbul ignore next */
+      throw new Error(
+        `Invalid message received: "${message}" Message must be JSON parsable.`
+      )
+    }
+
+    if (operationId) {
+      operation = this.operations.get(operationId)
+    }
+
+    switch (data.type) {
+      case GQL_CONNECTION_ACK:
+        this.reconnecting = false
+        this.ready = true
+        this.reconnectAttempts = 0
+
+        for (const operationId of this.operations.keys()) {
+          this.startOperation(operationId)
+        }
+
+        if (this.connectionCallback) {
+          this.connectionCallback()
+        }
+
+        break
+      case GQL_DATA:
+        /* istanbul ignore else */
+        if (operation) {
+          operation.handler(data.payload.data)
+        }
+        break
+      case GQL_COMPLETE:
+        /* istanbul ignore else */
+        if (operation) {
+          operation.handler(null)
+          this.operations.delete(operationId)
+        }
+
+        break
+      /* istanbul ignore next */
+      default:
+        /* istanbul ignore next */
+        throw new Error(`Invalid message type: "${data.type}"`)
+    }
+  }
+
+  startOperation (operationId) {
+    const { started, options, handler } = this.operations.get(operationId)
+    if (this.ready && !started) {
+      this.operations.set(operationId, { started: true, options, handler })
+      this.sendMessage(operationId, GQL_START, options)
+    }
+  }
+
+  createSubscription (query, variables, publish) {
+    const subscriptionString = JSON.stringify({ query, variables })
+    let operationId = this.subscriptionQueryMap[subscriptionString]
+
+    if (operationId && this.operations.get(operationId)) {
+      this.operationsCount[operationId] = this.operationsCount[operationId] + 1
+      return operationId
+    }
+
+    operationId = String(++this.operationId)
+
+    this.operations.set(operationId, {
+      started: false,
+      options: { query, variables },
+      handler: async (data) => {
+        await publish({
+          topic: `${this.serviceName}_${operationId}`,
+          payload: data
+        })
+      }
+    })
+    this.startOperation(operationId)
+    this.operationsCount[operationId] = 1
+
+    this.subscriptionQueryMap[subscriptionString] = operationId
+
+    return operationId
+  }
+}
+
+module.exports = SubscriptionClient
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index 8cfbce3..cdebbfe 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -1,25 +1,32 @@
+'use strict'
+
 const { subscribe, parse } = require('graphql')
 const { SubscriptionContext } = require('./subscriber')
-
-const GQL_CONNECTION_INIT = 'connection_init' // Client -> Server
-const GQL_CONNECTION_ACK = 'connection_ack' // Server -> Client
-
-const GQL_CONNECTION_TERMINATE = 'connection_terminate' // Client -> Server
-const GQL_START = 'start' // Client -> Server
-const GQL_DATA = 'data' // Server -> Client
-const GQL_ERROR = 'error' // Server -> Client
-const GQL_COMPLETE = 'complete' // Server -> Client
-const GQL_STOP = 'stop' // Client -> Server
+const { kEntityResolvers } = require('./gateway/make-resolver')
+const {
+  GQL_CONNECTION_INIT,
+  GQL_CONNECTION_ACK,
+  GQL_CONNECTION_TERMINATE,
+  GQL_START,
+  GQL_DATA,
+  GQL_ERROR,
+  GQL_COMPLETE,
+  GQL_STOP
+} = require('./subscription-protocol')
 
 module.exports = class SubscriptionConnection {
   constructor (socket, {
     schema,
     subscriber,
-    fastify
+    fastify,
+    lruGatewayResolvers,
+    entityResolvers
   }) {
     this.fastify = fastify
     this.socket = socket
     this.schema = schema
+    this.lruGatewayResolvers = lruGatewayResolvers
+    this.entityResolvers = entityResolvers
     this.subscriber = subscriber
     this.subscriptionContexts = new Map()
 
@@ -83,7 +90,12 @@ module.exports = class SubscriptionConnection {
       document,
       {}, // rootValue
       {
-        pubsub: sc
+        pubsub: sc,
+        lruGatewayResolvers: this.lruGatewayResolvers,
+        reply: {
+          [kEntityResolvers]: this.entityResolvers,
+          request: { headers: {} }
+        }
       },
       variables,
       operationName
diff --git a/lib/subscription-protocol.js b/lib/subscription-protocol.js
new file mode 100644
index 0000000..54318d8
--- /dev/null
+++ b/lib/subscription-protocol.js
@@ -0,0 +1,13 @@
+'use strict'
+
+module.exports.GQL_CONNECTION_INIT = 'connection_init' // Client -> Server
+module.exports.GQL_CONNECTION_ACK = 'connection_ack' // Server -> Client
+
+module.exports.GQL_CONNECTION_TERMINATE = 'connection_terminate' // Client -> Server
+module.exports.GQL_START = 'start' // Client -> Server
+module.exports.GQL_DATA = 'data' // Server -> Client
+module.exports.GQL_ERROR = 'error' // Server -> Client
+module.exports.GQL_COMPLETE = 'complete' // Server -> Client
+module.exports.GQL_STOP = 'stop' // Client -> Server
+
+module.exports.GRAPHQL_WS = 'graphql-ws'
diff --git a/lib/subscription.js b/lib/subscription.js
index 3b8146b..50d6870 100644
--- a/lib/subscription.js
+++ b/lib/subscription.js
@@ -8,7 +8,7 @@ function handle (conn) {
   conn.end(JSON.stringify({ error: 'unknown route' }))
 }
 
-function createConnectionHandler (schema, subscriber, fastify) {
+function createConnectionHandler ({ schema, subscriber, fastify, lruGatewayResolvers, entityResolversFactory }) {
   return (connection, request) => {
     const { socket } = connection
     if (socket.protocol === undefined ||
@@ -24,7 +24,9 @@ function createConnectionHandler (schema, subscriber, fastify) {
     const subscriptionConnection = new SubscriptionConnection(socket, {
       schema,
       subscriber,
-      fastify
+      fastify,
+      lruGatewayResolvers,
+      entityResolvers: entityResolversFactory && entityResolversFactory.create()
     })
 
     /* istanbul ignore next */
@@ -37,7 +39,8 @@ function createConnectionHandler (schema, subscriber, fastify) {
   }
 }
 
-module.exports = function (fastify, { getOptions, schema, subscriber, verifyClient }, next) {
+module.exports = function (fastify, opts, next) {
+  const { getOptions, schema, subscriber, verifyClient, lruGatewayResolvers, entityResolversFactory } = opts
   fastify.register(fastifyWebsocket, {
     handle,
     options: {
@@ -48,7 +51,13 @@ module.exports = function (fastify, { getOptions, schema, subscriber, verifyClie
 
   fastify.route({
     ...getOptions,
-    wsHandler: createConnectionHandler(schema, subscriber, fastify)
+    wsHandler: createConnectionHandler({
+      schema,
+      subscriber,
+      fastify,
+      lruGatewayResolvers,
+      entityResolversFactory
+    })
   })
 
   next()


Test patch:
diff --git a/test/gateway/config.js b/test/gateway/config.js
index 712dfb7..dfbe607 100644
--- a/test/gateway/config.js
+++ b/test/gateway/config.js
@@ -22,7 +22,7 @@ test('"schema" option not allowed in gateway moode', async (t) => {
   try {
     await app.ready()
   } catch (err) {
-    t.is(err.message, 'Adding "schema", "resolvers", "loaders" or "subscription" to plugin options when plugin is running in gateway mode is not allowed')
+    t.is(err.message, 'Adding "schema", "resolvers", "loaders" or to plugin options when plugin is running in gateway mode is not allowed')
   }
 })
 
@@ -39,7 +39,7 @@ test('"resolvers" option not allowed in gateway moode', async (t) => {
   try {
     await app.ready()
   } catch (err) {
-    t.is(err.message, 'Adding "schema", "resolvers", "loaders" or "subscription" to plugin options when plugin is running in gateway mode is not allowed')
+    t.is(err.message, 'Adding "schema", "resolvers", "loaders" or to plugin options when plugin is running in gateway mode is not allowed')
   }
 })
 
@@ -56,23 +56,6 @@ test('"loaders" option not allowed in gateway moode', async (t) => {
   try {
     await app.ready()
   } catch (err) {
-    t.is(err.message, 'Adding "schema", "resolvers", "loaders" or "subscription" to plugin options when plugin is running in gateway mode is not allowed')
-  }
-})
-
-test('"subscription" option not allowed in gateway moode', async (t) => {
-  const app = Fastify()
-
-  app.register(GQL, {
-    subscription: true,
-    gateway: {
-      services: []
-    }
-  })
-
-  try {
-    await app.ready()
-  } catch (err) {
-    t.is(err.message, 'Adding "schema", "resolvers", "loaders" or "subscription" to plugin options when plugin is running in gateway mode is not allowed')
+    t.is(err.message, 'Adding "schema", "resolvers", "loaders" or to plugin options when plugin is running in gateway mode is not allowed')
   }
 })
diff --git a/test/gateway/schema.js b/test/gateway/schema.js
index 3d09c1c..a622b5e 100644
--- a/test/gateway/schema.js
+++ b/test/gateway/schema.js
@@ -143,8 +143,6 @@ test('It builds the gateway schema correctly', async (t) => {
         name: 'user',
         url: `http://localhost:${userServicePort}/graphql`,
         rewriteHeaders: (headers) => {
-          // console.log('rewrite headers', headers)
-
           if (headers.authorization) {
             return {
               authorization: headers.authorization
diff --git a/test/gateway/subscription.js b/test/gateway/subscription.js
new file mode 100644
index 0000000..9043827
--- /dev/null
+++ b/test/gateway/subscription.js
@@ -0,0 +1,287 @@
+'use strict'
+
+const { test } = require('tap')
+const Fastify = require('fastify')
+const WebSocket = require('ws')
+const GQL = require('../..')
+
+const users = {
+  u1: {
+    id: 'u1',
+    name: 'John'
+  },
+  u2: {
+    id: 'u2',
+    name: 'Jane'
+  }
+}
+
+const messages = {}
+
+const userSchema = `
+extend type Query {
+  me: User
+}
+
+type User @key(fields: "id") {
+  id: ID!
+  name: String!
+}
+`
+
+const messageSchema = `
+extend type Mutation {
+  sendMessage(message: MessageInput!): Message
+}
+
+extend type Subscription {
+  newMessage(toUser: ID!): Message
+}
+
+type Message @key(fields: "id") {
+  id: ID!
+  text: String!
+  from: User
+  to: User
+}
+
+extend type User @key(fields: "id") {
+  id: ID! @external
+  messages: [Message]
+}
+
+input MessageInput {
+  fromUserId: ID!
+  toUserId: ID!
+  text: String!
+}
+`
+
+const userResolvers = {
+  Query: {
+    me: (root, args, context, info) => {
+      return users.u2
+    }
+  },
+  User: {
+    __resolveReference: (user, args, context, info) => {
+      return users[user.id]
+    }
+  }
+}
+
+const messageResolvers = {
+  Mutation: {
+    async sendMessage (root, { message }, { pubsub }) {
+      const id = Object.values(messages).length + 1
+
+      const result = {
+        id,
+        ...message
+      }
+
+      messages[id] = result
+
+      await pubsub.publish({
+        topic: `NEW_MESSAGE_${message.toUserId}`,
+        payload: {
+          newMessage: result
+        }
+      })
+
+      return result
+    }
+  },
+  Subscription: {
+    newMessage: {
+      subscribe: async (root, { toUser }, { pubsub }) => {
+        const subscription = await pubsub.subscribe(`NEW_MESSAGE_${toUser}`)
+
+        return subscription
+      }
+    }
+  },
+  Message: {
+    __resolveReference: (message) => messages[message.id],
+    from: (message) => {
+      return {
+        __typename: 'User',
+        id: message.fromUserId
+      }
+    },
+    to: (message) => {
+      return {
+        __typename: 'User',
+        id: message.toUserId
+      }
+    }
+  }
+}
+
+test('gateway subscription handling works correctly', t => {
+  t.plan(1)
+  let userService
+  let messageService
+  let gateway
+
+  function createUserService (callback) {
+    userService = Fastify()
+    userService.register(GQL, {
+      schema: userSchema,
+      resolvers: userResolvers,
+      federationMetadata: true,
+      subscription: true,
+      ide: 'playground'
+    })
+    userService.listen(0, callback)
+  }
+
+  function createMessageService (callback) {
+    messageService = Fastify()
+    messageService.register(GQL, {
+      schema: messageSchema,
+      resolvers: messageResolvers,
+      federationMetadata: true,
+      subscription: true,
+      ide: 'playground'
+    })
+    messageService.listen(0, callback)
+  }
+
+  function createGateway (callback) {
+    const userServicePort = userService.server.address().port
+    const messageServicePort = messageService.server.address().port
+
+    gateway = Fastify()
+    gateway.register(GQL, {
+      subscription: true,
+      ide: 'playground',
+      jit: 1,
+      gateway: {
+        services: [{
+          name: 'user',
+          url: `http://localhost:${userServicePort}/graphql`,
+          wsUrl: `ws://localhost:${userServicePort}/graphql`
+        }, {
+          name: 'message',
+          url: `http://localhost:${messageServicePort}/graphql`,
+          wsUrl: `ws://localhost:${messageServicePort}/graphql`
+        }]
+      }
+    })
+
+    gateway.listen(0, callback)
+  }
+
+  function runSubscription () {
+    const ws = new WebSocket(`ws://localhost:${(gateway.server.address()).port}/graphql`, 'graphql-ws')
+    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+    t.tearDown(() => {
+      client.destroy()
+      messageService.close()
+      userService.close()
+      gateway.close()
+      process.exit(0)
+    })
+    client.setEncoding('utf8')
+
+    client.write(JSON.stringify({
+      type: 'connection_init'
+    }))
+
+    client.write(JSON.stringify({
+      id: 1,
+      type: 'start',
+      payload: {
+        query: `
+          subscription {
+            newMessage(toUser: "u1") {
+              id
+              text
+              from {
+                id
+                name
+              }
+              to {
+                id
+                name
+              }
+            }
+          }
+        `
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'start',
+      payload: {
+        query: `
+          subscription {
+            newMessage(toUser: "u2") {
+              id
+            }
+          }
+        `
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'stop'
+    }))
+
+    client.on('data', (chunk) => {
+      const data = JSON.parse(chunk)
+
+      if (data.id === 1 && data.type === 'data') {
+        t.equal(chunk, JSON.stringify({
+          type: 'data',
+          id: 1,
+          payload: {
+            data: {
+              newMessage: {
+                id: '1',
+                text: 'Hi there',
+                from: {
+                  id: 'u2',
+                  name: 'Jane'
+                },
+                to: {
+                  id: 'u1',
+                  name: 'John'
+                }
+              }
+            }
+          }
+        }))
+
+        client.end()
+        t.end()
+      } else if (data.id === 2 && data.type === 'complete') {
+        gateway.inject({
+          method: 'POST',
+          url: '/graphql',
+          body: {
+            query: `
+              mutation {
+                sendMessage(message: {
+                  text: "Hi there",
+                  fromUserId: "u2",
+                  toUserId: "u1"
+                }) {
+                  id
+                }
+              }
+            `
+          }
+        })
+      }
+    })
+  }
+
+  const startGateway = createGateway.bind(null, runSubscription)
+  const startMessageService = createMessageService.bind(null, startGateway)
+
+  createUserService(startMessageService)
+})
diff --git a/test/subscription-client.js b/test/subscription-client.js
new file mode 100644
index 0000000..92a21c4
--- /dev/null
+++ b/test/subscription-client.js
@@ -0,0 +1,217 @@
+'use strict'
+const { test } = require('tap')
+const SubscriptionClient = require('../lib/subscription-client')
+const WS = require('ws')
+
+test('subscription client calls the publish method with the correct payload', (t) => {
+  const server = new WS.Server({ port: 0 })
+  const port = server.address().port
+
+  server.on('connection', function connection (ws) {
+    ws.on('message', function incoming (message) {
+      const data = JSON.parse(message)
+      if (data.type === 'start') {
+        ws.send(JSON.stringify({ id: '1', type: 'data', payload: { data: { foo: 'bar' } } }))
+      }
+    })
+
+    ws.send(JSON.stringify({ id: undefined, type: 'connection_ack' }))
+  })
+
+  const client = new SubscriptionClient(`ws://localhost:${port}`, {
+    reconnect: true,
+    maxReconnectAttempts: 10,
+    serviceName: 'test-service'
+  })
+
+  client.createSubscription('query', {}, (data) => {
+    t.deepEqual(data, {
+      topic: 'test-service_1',
+      payload: {
+        foo: 'bar'
+      }
+    })
+    client.close()
+    server.close()
+    t.end()
+  })
+})
+
+test('subscription client calls the publish method with null after GQL_COMPLETE type payload received', (t) => {
+  const server = new WS.Server({ port: 0 })
+  const port = server.address().port
+
+  server.on('connection', function connection (ws) {
+    ws.on('message', function incoming (message) {
+      const data = JSON.parse(message)
+      if (data.type === 'start') {
+        ws.send(JSON.stringify({ id: '1', type: 'complete' }))
+      }
+    })
+
+    ws.send(JSON.stringify({ id: undefined, type: 'connection_ack' }))
+  })
+
+  const client = new SubscriptionClient(`ws://localhost:${port}`, {
+    reconnect: true,
+    maxReconnectAttempts: 10,
+    serviceName: 'test-service'
+  })
+
+  client.createSubscription('query', {}, (data) => {
+    t.deepEqual(data, {
+      topic: 'test-service_1',
+      payload: null
+    })
+    client.close()
+    server.close()
+    t.end()
+  })
+})
+
+test('subscription client tries to reconnect when server closes', (t) => {
+  let server = new WS.Server({ port: 0 })
+  const port = server.address().port
+
+  server.on('connection', function connection (ws) {
+    ws.send(JSON.stringify({ id: undefined, type: 'connection_ack' }))
+  })
+
+  function createSubscription () {
+    client.createSubscription('query', {}, (data) => {
+      t.deepEqual(data, {
+        topic: 'test-service_1',
+        payload: null
+      })
+      client.close()
+      server.close()
+      t.end()
+    })
+  }
+
+  let shouldCloseServer = true
+
+  function connectionCallback () {
+    if (shouldCloseServer) {
+      server.close()
+      shouldCloseServer = false
+      server = new WS.Server({ port }, () => {
+        createSubscription()
+      })
+      server.on('connection', function connection (ws) {
+        ws.on('message', (message) => {
+          const data = JSON.parse(message)
+          if (data.type === 'connection_init') {
+            ws.send(JSON.stringify({ id: undefined, type: 'connection_ack' }))
+          } else if (data.type === 'start') {
+            ws.send(JSON.stringify({ id: '1', type: 'complete' }))
+          }
+        })
+      })
+    }
+  }
+
+  const client = new SubscriptionClient(`ws://localhost:${port}`, {
+    reconnect: true,
+    maxReconnectAttempts: 10,
+    serviceName: 'test-service',
+    connectionCallback
+  })
+})
+
+test('subscription client stops trying reconnecting after maxReconnectAttempts', (t) => {
+  const server = new WS.Server({ port: 0 })
+  const port = server.address().port
+
+  server.on('connection', function connection (ws) {
+    ws.send(JSON.stringify({ id: undefined, type: 'connection_ack' }))
+  })
+
+  const client = new SubscriptionClient(`ws://localhost:${port}`, {
+    reconnect: true,
+    maxReconnectAttempts: 1,
+    serviceName: 'test-service',
+    failedReconnectCallback: () => {
+      client.close()
+      server.close()
+      t.end()
+    }
+  })
+  server.close()
+
+  client.createSubscription('query', {}, (data) => {
+    t.deepEqual(data, {
+      topic: 'test-service_1',
+      payload: null
+    })
+  })
+})
+
+test('subscription client multiple subscriptions is handled by one operation', t => {
+  const server = new WS.Server({ port: 0 })
+  const port = server.address().port
+
+  server.on('connection', function connection (ws) {
+    ws.on('message', function incoming (message) {
+      const data = JSON.parse(message)
+      if (data.type === 'start') {
+        ws.send(JSON.stringify({ id: '1', type: 'complete' }))
+      }
+    })
+
+    ws.send(JSON.stringify({ id: undefined, type: 'connection_ack' }))
+  })
+
+  const client = new SubscriptionClient(`ws://localhost:${port}`, {
+    reconnect: true,
+    maxReconnectAttempts: 10,
+    serviceName: 'test-service'
+  })
+
+  function publish (data) {
+    client.close()
+    server.close()
+    t.end()
+  }
+
+  client.createSubscription('query', {}, publish)
+  client.createSubscription('query', {}, publish)
+})
+
+test('subscription client multiple subscriptions unsubscribe removes only one subscription', t => {
+  const server = new WS.Server({ port: 0 })
+  const port = server.address().port
+
+  server.on('connection', function connection (ws) {
+    ws.on('message', function incoming (message) {
+      const data = JSON.parse(message)
+      if (data.type === 'stop') {
+        ws.send(JSON.stringify({ id: '1', type: 'complete' }))
+      }
+    })
+
+    ws.send(JSON.stringify({ id: undefined, type: 'connection_ack' }))
+  })
+
+  const client = new SubscriptionClient(`ws://localhost:${port}`, {
+    reconnect: true,
+    maxReconnectAttempts: 10,
+    serviceName: 'test-service',
+    connectionCallback: () => {
+      function publish (data) {
+        client.close()
+        server.close()
+        t.end()
+      }
+
+      const operationId1 = client.createSubscription('query', {}, publish)
+      const operationId2 = client.createSubscription('query', {}, publish)
+      t.equal(operationId1, operationId2)
+
+      client.unsubscribe(operationId1)
+      t.equal(client.operationsCount[operationId1], 1)
+
+      client.unsubscribe(operationId1)
+    }
+  })
+})
diff --git a/test/subscription-connection.js b/test/subscription-connection.js
index c42409b..ee1ff6b 100644
--- a/test/subscription-connection.js
+++ b/test/subscription-connection.js
@@ -1,3 +1,4 @@
+'use strict'
 const { test } = require('tap')
 const proxyquire = require('proxyquire')
 const WebSocket = require('ws')