feat(server): Add onDisconnect callback (#94)

BREAKING CHANGE: The return function of `server.opened` (`closed`) now requires the close event code and reason for reporting to the `onDisconnect` callback.

Author: enisdenjo

docs/interfaces/_server_.server.md

@@ -35,8 +35,9 @@ of the `Context`. You may pass the initial request or the
 original WebSocket, if you need it down the road.
 
 Returns a function that should be called when the same socket
-has been closed, for whatever reason. The returned promise will
-resolve once the internal cleanup is complete.
+has been closed, for whatever reason. The close code and reason
+must be passed for reporting to the `onDisconnect` callback. Returned
+promise will resolve once the internal cleanup is complete.
 
 #### Parameters:
 

docs/interfaces/_server_.serveroptions.md

@@ -23,6 +23,7 @@ Name | Default |
 * [execute](_server_.serveroptions.md#execute)
 * [onComplete](_server_.serveroptions.md#oncomplete)
 * [onConnect](_server_.serveroptions.md#onconnect)
+* [onDisconnect](_server_.serveroptions.md#ondisconnect)
 * [onError](_server_.serveroptions.md#onerror)
 * [onNext](_server_.serveroptions.md#onnext)
 * [onOperation](_server_.serveroptions.md#onoperation)
@@ -130,6 +131,20 @@ in the close event reason.
 
 ___
 
+### onDisconnect
+
+• `Optional` **onDisconnect**: undefined \| (ctx: [Context](_server_.context.md)<E\>, code: number, reason: string) => Promise<void\> \| void
+
+Called when the socket/client closes/disconnects for
+whatever reason. Provides the close event too. Beware
+that this callback happens AFTER all subscriptions have
+been gracefuly completed.
+
+If you are interested in tracking the subscriptions completions,
+consider using the `onComplete` callback.
+
+___
+
 ### onError
 
 • `Optional` **onError**: undefined \| (ctx: [Context](_server_.context.md)<E\>, message: [ErrorMessage](_message_.errormessage.md), errors: readonly GraphQLError[]) => Promise<readonly GraphQLError[] \| void\> \| readonly GraphQLError[] \| void

docs/interfaces/_server_.websocket.md

@@ -36,7 +36,8 @@ to validate agains the supported ones.
 ▸ **close**(`code`: number, `reason`: string): Promise<void\> \| void
 
 Closes the socket gracefully. Will always provide
-the appropriate code and close reason.
+the appropriate code and close reason. `onDisconnect`
+callback will be called.
 
 The returned promise is used to control the graceful
 closure.

src/server.ts

@@ -163,6 +163,20 @@ export interface ServerOptions<E = unknown> {
     | Record<string, unknown>
     | boolean
     | void;
+  /**
+   * Called when the socket/client closes/disconnects for
+   * whatever reason. Provides the close event too. Beware
+   * that this callback happens AFTER all subscriptions have
+   * been gracefuly completed.
+   *
+   * If you are interested in tracking the subscriptions completions,
+   * consider using the `onComplete` callback.
+   */
+  onDisconnect?: (
+    ctx: Context<E>,
+    code: number,
+    reason: string,
+  ) => Promise<void> | void;
   /**
    * The subscribe callback executed right after
    * acknowledging the request before any payload
@@ -294,10 +308,14 @@ export interface Server<E = undefined> {
    * original WebSocket, if you need it down the road.
    *
    * Returns a function that should be called when the same socket
-   * has been closed, for whatever reason. The returned promise will
-   * resolve once the internal cleanup is complete.
+   * has been closed, for whatever reason. The close code and reason
+   * must be passed for reporting to the `onDisconnect` callback. Returned
+   * promise will resolve once the internal cleanup is complete.
    */
-  opened(socket: WebSocket, ctxExtra: E): () => Promise<void>; // closed
+  opened(
+    socket: WebSocket,
+    ctxExtra: E,
+  ): (code: number, reason: string) => Promise<void>; // closed
 }
 
 export interface WebSocket {
@@ -320,7 +338,8 @@ export interface WebSocket {
   send(data: string): Promise<void> | void;
   /**
    * Closes the socket gracefully. Will always provide
-   * the appropriate code and close reason.
+   * the appropriate code and close reason. `onDisconnect`
+   * callback will be called.
    *
    * The returned promise is used to control the graceful
    * closure.
@@ -392,6 +411,7 @@ export function makeServer<E = unknown>(options: ServerOptions<E>): Server<E> {
     subscribe,
     connectionInitWaitTimeout = 3 * 1000, // 3 seconds
     onConnect,
+    onDisconnect,
     onSubscribe,
     onOperation,
     onNext,
@@ -653,12 +673,13 @@ export function makeServer<E = unknown>(options: ServerOptions<E>): Server<E> {
         }
       });
 
-      // wait for close and cleanup
-      return async () => {
+      // wait for close, cleanup and the disconnect callback
+      return async (code, reason) => {
         if (connectionInitWait) clearTimeout(connectionInitWait);
         for (const sub of Object.values(ctx.subscriptions)) {
           await sub?.return?.();
         }
+        await onDisconnect?.(ctx, code, reason);
       };
     },
   };

src/tests/server.ts

@@ -1383,3 +1383,26 @@ describe('Subscribe', () => {
     client.ws.terminate();
   });
 });
+
+describe('Disconnect', () => {
+  it('should report close code and reason to disconnect callback', async (done) => {
+    const { url, waitForConnect } = await startTServer({
+      onDisconnect: (_ctx, code, reason) => {
+        expect(code).toBe(4321);
+        expect(reason).toBe('Byebye');
+        done();
+      },
+    });
+
+    const client = await createTClient(url);
+
+    client.ws.send(
+      stringifyMessage<MessageType.ConnectionInit>({
+        type: MessageType.ConnectionInit,
+      }),
+    );
+    await waitForConnect();
+
+    client.ws.close(4321, 'Byebye');
+  });
+});

src/use/ws.ts

@@ -105,10 +105,10 @@ export function useServer(
       { socket, request },
     );
 
-    socket.once('close', () => {
+    socket.once('close', (code, reason) => {
       if (pongWait) clearTimeout(pongWait);
       if (pingInterval) clearInterval(pingInterval);
-      closed();
+      closed(code, reason);
     });
   });