From dd0d4fa7828974b27876e138c0d09f78b2721d2d Mon Sep 17 00:00:00 2001
From: enisdenjo <badurinadenis@gmail.com>
Date: Wed, 13 Jan 2021 11:23:15 +0100
Subject: [PATCH] feat(server): Add `onClose` callback for closures at _any_
 point in time

Note that `onDisconnect` is called only if the connection is acknowledged.
---
 docs/interfaces/server.serveroptions.md | 20 +++++++++++++++++++
 src/server.ts                           | 26 +++++++++++++++++++++++--
 src/tests/server.ts                     | 21 ++++++++++++++------
 3 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/docs/interfaces/server.serveroptions.md b/docs/interfaces/server.serveroptions.md
index a82251f5..f7705fe5 100644
--- a/docs/interfaces/server.serveroptions.md
+++ b/docs/interfaces/server.serveroptions.md
@@ -21,6 +21,7 @@ Name | Default |
 - [connectionInitWaitTimeout](server.serveroptions.md#connectioninitwaittimeout)
 - [context](server.serveroptions.md#context)
 - [execute](server.serveroptions.md#execute)
+- [onClose](server.serveroptions.md#onclose)
 - [onComplete](server.serveroptions.md#oncomplete)
 - [onConnect](server.serveroptions.md#onconnect)
 - [onDisconnect](server.serveroptions.md#ondisconnect)
@@ -84,6 +85,25 @@ in the close event reason.
 
 ___
 
+### onClose
+
+• `Optional` **onClose**: *undefined* \| (`ctx`: [*Context*](server.context.md)<E\>, `code`: *number*, `reason`: *string*) => *void* \| *Promise*<*void*\>
+
+Called when the socket closes for whatever reason, at any
+point in time. Provides the close event too. Beware
+that this callback happens AFTER all subscriptions have
+been gracefully completed and AFTER the `onDisconnect` callback.
+
+If you are interested in tracking the subscriptions completions,
+consider using the `onComplete` callback.
+
+In comparison to `onDisconnect`, this callback will ALWAYS
+be called, regardless if the user succesfully went through
+the connection initialisation or not. `onConnect` might not
+called before the `onClose`.
+
+___
+
 ### onComplete
 
 • `Optional` **onComplete**: *undefined* \| (`ctx`: [*Context*](server.context.md)<E\>, `message`: [*CompleteMessage*](message.completemessage.md)) => *void* \| *Promise*<*void*\>
diff --git a/src/server.ts b/src/server.ts
index 3ba896d5..beb8d33a 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -184,6 +184,25 @@ export interface ServerOptions<E = unknown> {
     code: number,
     reason: string,
   ) => Promise<void> | void;
+  /**
+   * Called when the socket closes for whatever reason, at any
+   * point in time. Provides the close event too. Beware
+   * that this callback happens AFTER all subscriptions have
+   * been gracefully completed and AFTER the `onDisconnect` callback.
+   *
+   * If you are interested in tracking the subscriptions completions,
+   * consider using the `onComplete` callback.
+   *
+   * In comparison to `onDisconnect`, this callback will ALWAYS
+   * be called, regardless if the user succesfully went through
+   * the connection initialisation or not. `onConnect` might not
+   * called before the `onClose`.
+   */
+  onClose?: (
+    ctx: Context<E>,
+    code: number,
+    reason: string,
+  ) => Promise<void> | void;
   /**
    * The subscribe callback executed right after
    * acknowledging the request before any payload
@@ -419,6 +438,7 @@ export function makeServer<E = unknown>(options: ServerOptions<E>): Server<E> {
     connectionInitWaitTimeout = 3 * 1000, // 3 seconds
     onConnect,
     onDisconnect,
+    onClose,
     onSubscribe,
     onOperation,
     onNext,
@@ -430,8 +450,9 @@ export function makeServer<E = unknown>(options: ServerOptions<E>): Server<E> {
     opened(socket, extra) {
       if (socket.protocol !== GRAPHQL_TRANSPORT_WS_PROTOCOL) {
         socket.close(1002, 'Protocol Error');
-        return async () => {
-          /* nothing was set up */
+        return async (code, reason) => {
+          /* nothing was set up, just notify the closure */
+          await onClose?.(ctx, code, reason);
         };
       }
 
@@ -687,6 +708,7 @@ export function makeServer<E = unknown>(options: ServerOptions<E>): Server<E> {
           await sub?.return?.();
         }
         if (ctx.acknowledged) await onDisconnect?.(ctx, code, reason);
+        await onClose?.(ctx, code, reason);
       };
     },
   };
diff --git a/src/tests/server.ts b/src/tests/server.ts
index 341fb1c8..119741f5 100644
--- a/src/tests/server.ts
+++ b/src/tests/server.ts
@@ -1391,12 +1391,18 @@ describe('Subscribe', () => {
   });
 });
 
-describe('Disconnect', () => {
-  it('should report close code and reason to disconnect callback after connection acknowledgement', async (done) => {
+describe('Disconnect/close', () => {
+  it('should report close code and reason to disconnect and close callback after connection acknowledgement', async (done) => {
     const { url, waitForConnect } = await startTServer({
+      // 1st
       onDisconnect: (_ctx, code, reason) => {
         expect(code).toBe(4321);
         expect(reason).toBe('Byebye');
+      },
+      // 2nd
+      onClose: (_ctx, code, reason) => {
+        expect(code).toBe(4321);
+        expect(reason).toBe('Byebye');
         done();
       },
     });
@@ -1413,17 +1419,20 @@ describe('Disconnect', () => {
     client.ws.close(4321, 'Byebye');
   });
 
-  it('should not trigger the disconnect callback if connection is not acknowledged', async () => {
-    const { url, waitForClientClose } = await startTServer({
+  it('should trigger the close callback instead of disconnect if connection is not acknowledged', async (done) => {
+    const { url } = await startTServer({
       onDisconnect: () => {
         fail("Disconnect callback shouldn't be triggered");
       },
+      onClose: (_ctx, code, reason) => {
+        expect(code).toBe(4321);
+        expect(reason).toBe('Byebye');
+        done();
+      },
     });
 
     const client = await createTClient(url);
 
     client.ws.close(4321, 'Byebye');
-
-    await waitForClientClose();
   });
 });