Implement keepalive on websockets (#2690)

Author: benjie

.changeset/lovely-moments-argue.md

@@ -0,0 +1,6 @@
+---
+"postgraphile": patch
+"grafserv": patch
+---
+
+Add keepalive to grafserv websocket connections.

grafast/grafserv/src/index.ts

@@ -146,6 +146,13 @@ declare global {
        */
       websockets?: boolean;
 
+      /**
+       * Duration (in milliseconds) between pings. Set to `-1` to disable.
+       *
+       * @defaultValue `12_000`
+       */
+      websocketKeepalive?: number;
+
       /**
        * If you would like to customize the way in which errors are masked, you may
        * pass your own error masking function here. You can also import

grafast/grafserv/src/servers/fastify/v4/index.ts

@@ -24,6 +24,7 @@ import {
   normalizeRequest,
   processHeaders,
 } from "../../../utils.js";
+import { DEFAULT_WEBSOCKET_KEEPALIVE } from "../../../websocketKeepalive.js";
 
 declare global {
   namespace Grafast {
@@ -202,8 +203,11 @@ export class FastifyGrafserv extends GrafservBase {
     };
 
     // Build websocket handler.
+    const keepaliveInterval =
+      this.getPreset().grafserv?.websocketKeepalive ??
+      DEFAULT_WEBSOCKET_KEEPALIVE;
     const wsHandler = websockets
-      ? makeHandler(makeGraphQLWSConfig(this))
+      ? makeHandler(makeGraphQLWSConfig(this), keepaliveInterval)
       : undefined;
 
     // Attach HTTP handler for POST requests.

grafast/grafserv/src/servers/node/index.ts

@@ -20,6 +20,7 @@ import {
   makeGraphQLWSConfig,
   processHeaders,
 } from "../../utils.js";
+import { handleWebSocketKeepalive } from "../../websocketKeepalive.js";
 
 declare global {
   namespace Grafast {
@@ -352,6 +353,8 @@ export async function makeNodeUpgradeHandler(instance: GrafservBase) {
     });
   };
   const onConnection = (socket: WebSocket, request: IncomingMessage) => {
+    handleWebSocketKeepalive(socket, instance.getPreset());
+
     // a new socket opened, let graphql-ws take over
     const closed = graphqlWsServer.opened(
       {

grafast/grafserv/src/websocketKeepalive.ts

@@ -0,0 +1,68 @@
+export const DEFAULT_WEBSOCKET_KEEPALIVE = 12_000;
+import type WebSocket from "ws";
+
+export function handleWebSocketKeepalive(
+  socket: WebSocket,
+  preset: GraphileConfig.ResolvedPreset,
+): void {
+  const keepaliveInterval =
+    preset.grafserv?.websocketKeepalive ?? DEFAULT_WEBSOCKET_KEEPALIVE;
+  if (!Number.isFinite(keepaliveInterval) || keepaliveInterval <= 0) {
+    // Keepalive disabled
+    return;
+  }
+
+  /**
+   * Sending a ping and waiting for a pong are mutually exclusive, so this
+   * timer is used for both.
+   */
+  let timer: NodeJS.Timeout | null = null;
+
+  /**
+   * Cleans up the timer, always call this before re-assigning timer (to ensure
+   * that an out-of-order pong doesn't cause two timers to run concurrently).
+   */
+  const stopTimer = () => {
+    if (timer != null) {
+      clearTimeout(timer);
+      timer = null;
+    }
+  };
+
+  /** First half of a heart beat - send ping */
+  const sendPing = () => {
+    stopTimer();
+    // Schedule timeout
+    timer = setTimeout(handleTimeout, keepaliveInterval);
+    socket.ping();
+  };
+  /** Second half of a heart beat - receive pong */
+  const handlePong = () => {
+    stopTimer();
+    // Schedule the next ping
+    timer = setTimeout(sendPing, keepaliveInterval);
+  };
+
+  /** Terminal handler, due to timeout */
+  const handleTimeout = () => {
+    stopTimer();
+    releaseListeners();
+    // Kill the socket (after we've released the listeners)
+    socket.terminate();
+  };
+  /** Terminal handler, due to natural socket close */
+  const handleClose = (_code: number, _reason: Buffer) => {
+    stopTimer();
+    releaseListeners();
+  };
+
+  const releaseListeners = () => {
+    socket.off("pong", handlePong);
+    socket.off("close", handleClose);
+  };
+  socket.on("pong", handlePong);
+  socket.on("close", handleClose);
+
+  // Schedule the first ping
+  timer = setTimeout(sendPing, keepaliveInterval);
+}

postgraphile/postgraphile/graphile.config.ts

@@ -427,6 +427,7 @@ const preset: GraphileConfig.Preset = {
         extend type Subscription {
           sub(topic: String!): Int
           gql: Int
+          slow: String
         }
       `,
       objects: {
@@ -472,6 +473,26 @@ const preset: GraphileConfig.Preset = {
                 [sleep],
               ),
             },
+            slow: {
+              resolve: EXPORTABLE(
+                () =>
+                  function resolve(e) {
+                    return e;
+                  },
+                [],
+              ),
+              subscribe: EXPORTABLE(
+                (sleep) =>
+                  async function* subscribe() {
+                    while (true) {
+                      yield new Date().toISOString();
+                      // Wait two minutes between ticks
+                      await sleep(120000);
+                    }
+                  },
+                [sleep],
+              ),
+            },
           },
         },
       },