feat: allow to use the adapter within a Node.js cluster

Related: socketio/socket.io-cluster-adapter#1

Author: darrachequesne

README.md

@@ -9,10 +9,15 @@ The `@socket.io/postgres-adapter` package allows broadcasting packets between mu
 
 **Table of contents**
 
-- [Supported features](#supported-features)
-- [Installation](#installation)
-- [Usage](#usage)
-- [License](#license)
+<!-- TOC -->
+  * [Supported features](#supported-features)
+  * [Installation](#installation)
+  * [Usage](#usage)
+    * [Standalone](#standalone)
+    * [With Node.js cluster](#with-nodejs-cluster)
+  * [Options](#options)
+  * [License](#license)
+<!-- TOC -->
 
 ## Supported features
 
@@ -31,6 +36,8 @@ npm install @socket.io/postgres-adapter
 
 ## Usage
 
+### Standalone
+
 ```js
 import { Server } from "socket.io";
 import { createAdapter } from "@socket.io/postgres-adapter";
@@ -62,6 +69,60 @@ io.adapter(createAdapter(pool));
 io.listen(3000);
 ```
 
+### With Node.js cluster
+
+```js
+import cluster from "node:cluster";
+import { createServer } from "node:http";
+import { availableParallelism } from "node:os";
+import { Server } from "socket.io";
+import { setupPrimary } from "@socket.io/postgres-adapter"
+import { createAdapter } from "@socket.io/cluster-adapter";
+import pg from "pg";
+
+if (cluster.isPrimary) {
+  const pool = new pg.Pool({
+    user: "postgres",
+    password: "changeit",
+  });
+  
+  await pool.query(`
+    CREATE TABLE IF NOT EXISTS socket_io_attachments (
+        id          bigserial UNIQUE,
+        created_at  timestamptz DEFAULT NOW(),
+        payload     bytea
+    );
+  `);
+
+  setupPrimary(pool);
+
+  for (let i = 0; i < availableParallelism(); i++) {
+    cluster.fork();
+  }
+} else {
+  const io = new Server({
+    adapter: createAdapter(),
+  });
+
+  io.on("connection", (socket) => {
+    /* ... */
+  });
+
+  io.listen(3000);
+}
+```
+
+## Options
+
+| Name                | Description                                                                                                | Default value           |
+|---------------------|------------------------------------------------------------------------------------------------------------|-------------------------|
+| `channelPrefix`     | The prefix of the notification channel.                                                                    | `socket.io`             |
+| `tableName`         | The name of the table for payloads over the 8000 bytes limit or containing binary data.                    | `socket_io_attachments` |
+| `payloadThreshold`  | The threshold for the payload size in bytes (see https://www.postgresql.org/docs/current/sql-notify.html). | `8_000`                 |
+| `cleanupInterval`   | Number of ms between two cleanup queries.                                                                  | `30_000`                |
+| `heartbeatInterval` | The number of ms between two heartbeats.                                                                   | `5_000`                 |
+| `heartbeatTimeout`  | The number of ms without heartbeat before we consider a node down.                                         | `10_000`                |
+
 ## License
 
 [MIT](LICENSE)

lib/adapter.ts

@@ -0,0 +1,133 @@
+import { type Pool } from "pg";
+import { ClusterAdapterWithHeartbeat } from "socket.io-adapter";
+import type {
+  ClusterAdapterOptions,
+  ClusterMessage,
+  ClusterResponse,
+  Offset,
+  ServerId,
+} from "socket.io-adapter";
+import debugModule from "debug";
+import { PubSubClient } from "./util";
+
+const debug = debugModule("socket.io-postgres-adapter");
+
+export interface PostgresAdapterOptions {
+  /**
+   * The prefix of the notification channel
+   * @default "socket.io"
+   */
+  channelPrefix?: string;
+  /**
+   * The name of the table for payloads over the 8000 bytes limit or containing binary data
+   * @default "socket_io_attachments"
+   */
+  tableName?: string;
+  /**
+   * The threshold for the payload size in bytes (see https://www.postgresql.org/docs/current/sql-notify.html)
+   * @default 8000
+   */
+  payloadThreshold?: number;
+  /**
+   * Number of ms between two cleanup queries
+   * @default 30000
+   */
+  cleanupInterval?: number;
+  /**
+   * Handler for errors. If undefined, the errors will be simply logged.
+   *
+   * @default undefined
+   */
+  errorHandler?: (err: Error) => void;
+}
+
+/**
+ * Returns a function that will create a PostgresAdapter instance.
+ *
+ * @param pool - a pg.Pool instance
+ * @param opts - additional options
+ *
+ * @public
+ */
+export function createAdapter(
+  pool: Pool,
+  opts: PostgresAdapterOptions & ClusterAdapterOptions = {}
+) {
+  const options = Object.assign(
+    {
+      channelPrefix: "socket.io",
+      tableName: "socket_io_attachments",
+      payloadThreshold: 8_000,
+      cleanupInterval: 30_000,
+      errorHandler: (err: Error) => debug(err),
+    },
+    opts
+  );
+
+  const namespaces = new Map<string, PostgresAdapter>();
+  const client = new PubSubClient(
+    pool,
+    options,
+    (msg) => {
+      // @ts-expect-error uid is protected
+      return namespaces.get(msg.nsp)?.uid === msg.uid;
+    },
+    (msg) => {
+      namespaces.get(msg.nsp)?.onMessage(msg);
+    }
+  );
+
+  return function (nsp: any) {
+    let adapter = new PostgresAdapter(nsp, opts, client);
+
+    namespaces.set(nsp.name, adapter);
+    client.addNamespace(nsp.name);
+
+    const defaultClose = adapter.close;
+
+    adapter.close = () => {
+      namespaces.delete(nsp.name);
+
+      if (namespaces.size === 0) {
+        client.close();
+      }
+
+      defaultClose.call(adapter);
+    };
+
+    return adapter;
+  };
+}
+
+export class PostgresAdapter extends ClusterAdapterWithHeartbeat {
+  /**
+   * Adapter constructor.
+   *
+   * @param nsp - the namespace
+   * @param opts - additional options
+   * @param client
+   *
+   * @public
+   */
+  constructor(
+    nsp: any,
+    opts: ClusterAdapterOptions,
+    private readonly client: PubSubClient
+  ) {
+    super(nsp, opts);
+  }
+
+  protected override doPublish(message: ClusterMessage): Promise<Offset> {
+    return this.client.publish(message).then(() => {
+      // connection state recovery is not currently supported
+      return "";
+    });
+  }
+
+  protected override doPublishResponse(
+    _requesterUid: ServerId,
+    response: ClusterResponse
+  ) {
+    return this.client.publish(response);
+  }
+}