Send a close to all websocket connections when server closes (#46)

Author: mcollina

index.js

@@ -66,6 +66,18 @@ function fastifyWebsocket (fastify, opts, next) {
 
   fastify.addHook('onClose', close)
 
+  // Fastify is missing a pre-close event, or the ability to
+  // add a hook before the server.close call. We need to resort
+  // to monkeypatching for now.
+  const oldClose = fastify.server.close
+  fastify.server.close = function (cb) {
+    const server = fastify.websocketServer
+    for (const client of server.clients) {
+      client.close()
+    }
+    oldClose.call(this, cb)
+  }
+
   function handleRouting (connection, request) {
     const response = new ServerResponse(request)
     request[kWs] = WebSocket.createWebSocketStream(connection)
@@ -77,7 +89,8 @@ function fastifyWebsocket (fastify, opts, next) {
 }
 
 function close (fastify, done) {
-  fastify.websocketServer.close(done)
+  const server = fastify.websocketServer
+  server.close(done)
 }
 
 module.exports = fp(fastifyWebsocket, {

test/base.js

@@ -144,3 +144,46 @@ test('Should throw on an invalid handle parameter', (t) => {
     t.equal(err.message, 'invalid handle function')
   })
 })
+
+test('Should gracefully close with a connected client', (t) => {
+  t.plan(6)
+
+  const fastify = Fastify()
+
+  fastify.register(fastifyWebsocket, { handle })
+
+  function handle (connection) {
+    connection.setEncoding('utf8')
+    connection.write('hello client')
+
+    connection.once('data', (chunk) => {
+      t.equal(chunk, 'hello server')
+    })
+
+    connection.on('end', () => {
+      t.pass('end emitted on server side')
+    })
+    // this connection stays alive untile we close the server
+  }
+
+  fastify.listen(0, (err) => {
+    t.error(err)
+
+    const ws = new WebSocket('ws://localhost:' + fastify.server.address().port)
+    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8' })
+
+    client.setEncoding('utf8')
+    client.write('hello server')
+
+    client.on('end', () => {
+      t.pass('end emitted on client side')
+    })
+
+    client.once('data', (chunk) => {
+      t.equal(chunk, 'hello client')
+      fastify.close(function (err) {
+        t.error(err)
+      })
+    })
+  })
+})