feat: WebSocket Ping and Pong as keep-alive (enisdenjo#11)

* feat(server): begin with simple implementation

* fix(server): issue pings only on open sockets

* test(server): check a few keep-alive cases

* docs(protocol): specify keep-alive

PROTOCOL.md

@@ -20,6 +20,14 @@ The server can terminate the socket (kick the client off) at any time. The close
 
 The client terminates the socket and closes the connection by dispatching a `1000: Normal Closure` close event to the server indicating a normal closure.
 
+## Keep-Alive
+
+The server will occasionally check if the client is still "alive", available and listening. In order to perform this check, implementation leverages the standardized [Pings and Pongs: The Heartbeat of WebSockets](https://developer.mozilla.org/en-US/docs/Web/API/WebSockets_API/Writing_WebSocket_servers#Pings_and_Pongs_The_Heartbeat_of_WebSockets).
+
+Keep-Alive interval and the "pong wait" timeout can be tuned by using the accompanying configuration parameter on the server.
+
+Ping and Pong feature is a mandatory requirement by [The WebSocket Protocol](https://tools.ietf.org/html/rfc6455#section-5.5.2). All clients that don't support it are **not** RFC6455 compliant and will simply have their socket terminated after the pong wait has passed.
+
 ## Message types
 
 ### `ConnectionInit`

src/server.ts

@@ -136,6 +136,15 @@ export interface ServerOptions {
    * has been closed.
    */
   onComplete?: (ctx: Context, message: CompleteMessage) => void;
+  /**
+   * The timout between dispatched keep-alive messages. Internally the lib
+   * uses the [WebSocket Ping and Pongs]((https://developer.mozilla.org/en-US/docs/Web/API/WebSockets_API/Writing_WebSocket_servers#Pings_and_Pongs_The_Heartbeat_of_WebSockets)) to check that the link between
+   * the clients and the server is operating and to prevent the link from being broken due to idling.
+   * Set to nullish value to disable.
+   *
+   * @default 12 * 1000 (12 seconds)
+   */
+  keepAlive?: number;
 }
 
 export interface Context {
@@ -197,6 +206,7 @@ export function createServer(
     formatExecutionResult,
     onSubscribe,
     onComplete,
+    keepAlive = 12 * 1000, // 12 seconds
   } = options;
   const webSocketServer = isWebSocketServer(websocketOptionsOrServer)
     ? websocketOptionsOrServer
@@ -238,12 +248,45 @@ export function createServer(
         }
       }, connectionInitWaitTimeout);
 
+    // keep alive through ping-pong messages
+    // read more about the websocket heartbeat here: https://developer.mozilla.org/en-US/docs/Web/API/WebSockets_API/Writing_WebSocket_servers#Pings_and_Pongs_The_Heartbeat_of_WebSockets
+    let pongWait: NodeJS.Timeout | null;
+    const pingInterval =
+      keepAlive && // even 0 disables it
+      keepAlive !== Infinity &&
+      setInterval(() => {
+        // ping pong on open sockets only
+        if (socket.readyState === WebSocket.OPEN) {
+          // terminate the connection after pong wait has passed because the client is idle
+          pongWait = setTimeout(() => {
+            socket.terminate();
+          }, keepAlive);
+
+          // listen for client's pong and stop socket termination
+          socket.once('pong', () => {
+            if (pongWait) {
+              clearTimeout(pongWait);
+              pongWait = null;
+            }
+          });
+
+          // issue a ping to the client
+          socket.ping();
+        }
+      }, keepAlive);
+
     function errorOrCloseHandler(
       errorOrClose: WebSocket.ErrorEvent | WebSocket.CloseEvent,
     ) {
       if (connectionInitWait) {
         clearTimeout(connectionInitWait);
       }
+      if (pongWait) {
+        clearTimeout(pongWait);
+      }
+      if (pingInterval) {
+        clearInterval(pingInterval);
+      }
 
       if (isErrorEvent(errorOrClose)) {
         // TODO-db-200805 leaking sensitive information by sending the error message too?

src/tests/server.ts

@@ -674,3 +674,87 @@ describe('Subscribe', () => {
     await wait(20);
   });
 });
+
+describe('keepAlive', () => {
+  it('should dispatch pings after the timeout has passed', async () => {
+    await makeServer({
+      keepAlive: 50,
+    });
+
+    const client = new WebSocket(url, GRAPHQL_TRANSPORT_WS_PROTOCOL);
+    client.onopen = () => {
+      client.send(
+        stringifyMessage<MessageType.ConnectionInit>({
+          type: MessageType.ConnectionInit,
+        }),
+      );
+    };
+    await wait(10);
+
+    const onPingFn = jest.fn();
+    client.once('ping', onPingFn);
+    await wait(50);
+
+    expect(onPingFn).toBeCalled();
+  });
+
+  it('should not dispatch pings if disabled with nullish timeout', async () => {
+    await makeServer({
+      keepAlive: 0,
+    });
+
+    const client = new WebSocket(url, GRAPHQL_TRANSPORT_WS_PROTOCOL);
+    client.onopen = () => {
+      client.send(
+        stringifyMessage<MessageType.ConnectionInit>({
+          type: MessageType.ConnectionInit,
+        }),
+      );
+    };
+    await wait(10);
+
+    const onPingFn = jest.fn();
+    client.once('ping', onPingFn);
+    await wait(50);
+
+    expect(onPingFn).not.toBeCalled();
+  });
+
+  it('should terminate the socket if no pong is sent in response to a ping', async () => {
+    expect.assertions(4);
+
+    await makeServer({
+      keepAlive: 50,
+    });
+
+    const client = new WebSocket(url, GRAPHQL_TRANSPORT_WS_PROTOCOL);
+    client.onopen = () => {
+      client.send(
+        stringifyMessage<MessageType.ConnectionInit>({
+          type: MessageType.ConnectionInit,
+        }),
+      );
+    };
+    await wait(10);
+
+    // disable pong
+    client.pong = () => {
+      /**/
+    };
+    client.onclose = (event) => {
+      // termination is not graceful or clean
+      expect(event.code).toBe(1006);
+      expect(event.wasClean).toBeFalsy();
+    };
+
+    const onPingFn = jest.fn();
+    client.once('ping', onPingFn);
+    await wait(50);
+
+    expect(onPingFn).toBeCalled(); // ping is received
+
+    await wait(50 + 10); // wait for the timeout to pass and termination to settle
+
+    expect(client.readyState).toBe(WebSocket.CLOSED);
+  });
+});