setup y-websocket server with y-postgresql persistence layer

Author: TabrizPlv

backend/collaboration/package.json

@@ -21,6 +21,7 @@
     "env-cmd": "^10.1.0",
     "express": "^4.21.1",
     "http-status-codes": "^2.3.0",
+    "pg": "^8.13.0",
     "pino": "^9.4.0",
     "pino-http": "^10.3.0",
     "redis": "^4.7.0",

backend/collaboration/src/config.ts

@@ -3,3 +3,5 @@ import 'dotenv/config';
 export const UI_HOST = process.env.PEERPREP_UI_HOST!;
 
 export const EXPRESS_PORT = process.env.EXPRESS_PORT;
+
+export const WEBSOCKET_PORT = process.env.WEBSOCKET_PORT;

backend/collaboration/src/index.ts

@@ -4,9 +4,11 @@ import { dbHealthCheck } from '@/server';
 import server from '@/websocket';
 
 const port = Number.parseInt(EXPRESS_PORT || '8001');
+const wsPort = Number.parseInt(WEBSOCKET_PORT || '8002');
 
 const listenMessage = `App listening on port: ${port}`;
 server.listen(port, () => {
   void dbHealthCheck();
   logger.info(listenMessage);
 });
+server.listen(wsPort, () => logger.info(`WebSocket listening on port: ${wsPort}`));

backend/collaboration/src/y-postgresql/interfaces.ts

@@ -0,0 +1,14 @@
+import * as awarenessProtocol from 'y-protocols/awareness.js';
+import * as Y from 'yjs';
+
+export interface IWSSharedDoc extends Y.Doc {
+  name: string;
+  conns: Map<object, Set<number>>;
+  awareness: awarenessProtocol.Awareness;
+}
+
+export interface IPersistence {
+  bindState: (arg1: string, arg2: IWSSharedDoc) => void;
+  writeState: (arg1: string, arg2: IWSSharedDoc) => Promise<any>;
+  provider?: any;
+}

backend/collaboration/src/y-postgresql/utils.ts

@@ -0,0 +1,237 @@
+import * as Y from 'yjs';
+import * as syncProtocol from 'y-protocols/sync';
+import * as awarenessProtocol from 'y-protocols/awareness';
+
+import * as encoding from 'lib0/encoding';
+import * as decoding from 'lib0/decoding';
+import * as map from 'lib0/map';
+
+import { IPersistence, IWSSharedDoc } from '@/y-postgresql/interfaces';
+
+const wsReadyStateConnecting = 0;
+const wsReadyStateOpen = 1;
+const wsReadyStateClosing = 2; // eslint-disable-line
+const wsReadyStateClosed = 3; // eslint-disable-line
+
+// disable gc when using snapshots!
+const gcEnabled = process.env.GC !== 'false' && process.env.GC !== '0';
+
+let persistence: IPersistence | null = null;
+
+export const setPersistence = (persistence_: IPersistence) => {
+  persistence = persistence_;
+};
+
+export const getPersistence = () => persistence;
+
+// exporting docs so that others can use it
+export const docs = new Map<string, IWSSharedDoc>();
+
+const messageSync = 0;
+const messageAwareness = 1;
+// const messageAuth = 2
+
+const updateHandler = (update: Uint8Array, origin: any, doc: Y.Doc) => {
+  const sharedDoc = doc as IWSSharedDoc;
+
+  const encoder = encoding.createEncoder();
+  encoding.writeVarUint(encoder, messageSync);
+  syncProtocol.writeUpdate(encoder, update);
+  const message = encoding.toUint8Array(encoder);
+  sharedDoc.conns.forEach((_, conn) => send(sharedDoc, conn, message));
+};
+
+class WSSharedDoc extends Y.Doc implements IWSSharedDoc {
+  name: string;
+  conns: Map<object, Set<number>>;
+  awareness: awarenessProtocol.Awareness;
+
+  constructor(name: string) {
+    super({ gc: gcEnabled });
+    this.name = name;
+    this.conns = new Map();
+    this.awareness = new awarenessProtocol.Awareness(this);
+    this.awareness.setLocalState(null);
+
+    const awarenessChangeHandler = (
+      {
+        added,
+        updated,
+        removed,
+      }: {
+        added: Array<number>;
+        updated: Array<number>;
+        removed: Array<number>;
+      },
+      conn: object | null
+    ) => {
+      const changedClients = added.concat(updated, removed);
+      if (conn !== null) {
+        const connControlledIDs = /** @type {Set<number>} */ this.conns.get(conn);
+        if (connControlledIDs !== undefined) {
+          added.forEach((clientID) => {
+            connControlledIDs.add(clientID);
+          });
+          removed.forEach((clientID) => {
+            connControlledIDs.delete(clientID);
+          });
+        }
+      }
+      // broadcast awareness update
+      const encoder = encoding.createEncoder();
+      encoding.writeVarUint(encoder, messageAwareness);
+      encoding.writeVarUint8Array(
+        encoder,
+        awarenessProtocol.encodeAwarenessUpdate(this.awareness, changedClients)
+      );
+      const buff = encoding.toUint8Array(encoder);
+      this.conns.forEach((_, c) => {
+        send(this, c, buff);
+      });
+    };
+    this.awareness.on('update', awarenessChangeHandler);
+    this.on('update', updateHandler);
+  }
+}
+
+/**
+ * Gets a Y.Doc by name, whether in memory or on disk
+ *
+ * @param {string} docname - the name of the Y.Doc to find or create
+ * @param {boolean} gc - whether to allow gc on the doc (applies only when created)
+ * @return {WSSharedDoc}
+ */
+export const getYDoc = (docname: string, gc = true) =>
+  map.setIfUndefined(docs, docname, () => {
+    const doc = new WSSharedDoc(docname);
+    doc.gc = gc;
+    if (persistence !== null) {
+      persistence.bindState(docname, doc);
+    }
+    docs.set(docname, doc);
+    return doc;
+  });
+
+const messageListener = (conn: any, doc: IWSSharedDoc, message: Uint8Array) => {
+  try {
+    const encoder = encoding.createEncoder();
+    const decoder = decoding.createDecoder(message);
+    const messageType = decoding.readVarUint(decoder);
+    switch (messageType) {
+      case messageSync:
+        encoding.writeVarUint(encoder, messageSync);
+        syncProtocol.readSyncMessage(decoder, encoder, doc, conn);
+
+        // If the `encoder` only contains the type of reply message and no
+        // message, there is no need to send the message. When `encoder` only
+        // contains the type of reply, its length is 1.
+        if (encoding.length(encoder) > 1) {
+          send(doc, conn, encoding.toUint8Array(encoder));
+        }
+        break;
+      case messageAwareness: {
+        awarenessProtocol.applyAwarenessUpdate(
+          doc.awareness,
+          decoding.readVarUint8Array(decoder),
+          conn
+        );
+        break;
+      }
+    }
+  } catch (err) {
+    console.error(err);
+  }
+};
+
+const closeConn = (doc: IWSSharedDoc, conn: any) => {
+  if (doc.conns.has(conn)) {
+    const controlledIds = doc.conns.get(conn);
+    doc.conns.delete(conn);
+    if (controlledIds) {
+      awarenessProtocol.removeAwarenessStates(doc.awareness, Array.from(controlledIds), null);
+    }
+    if (doc.conns.size === 0 && persistence !== null) {
+      // if persisted, we store state and destroy ydocument
+      persistence.writeState(doc.name, doc).then(() => {
+        doc.destroy();
+      });
+      docs.delete(doc.name);
+    }
+  }
+  conn.close();
+};
+
+const send = (doc: IWSSharedDoc, conn: any, m: Uint8Array) => {
+  if (conn.readyState !== wsReadyStateConnecting && conn.readyState !== wsReadyStateOpen) {
+    closeConn(doc, conn);
+  }
+  try {
+    conn.send(m, (err: any) => {
+      if (err != null) {
+        closeConn(doc, conn);
+      }
+    });
+  } catch (e) {
+    closeConn(doc, conn);
+  }
+};
+
+const pingTimeout = 30000;
+
+export const setupWSConnection = (
+  conn: any,
+  req: any,
+  { docName = req.url.slice(1).split('?')[0], gc = true } = {}
+) => {
+  conn.binaryType = 'arraybuffer';
+  // get doc, initialize if it does not exist yet
+  const doc = getYDoc(docName, gc);
+  doc.conns.set(conn, new Set());
+  // listen and reply to events
+  conn.on('message', (message: ArrayBuffer) => messageListener(conn, doc, new Uint8Array(message)));
+
+  // Check if connection is still alive
+  let pongReceived = true;
+  const pingInterval = setInterval(() => {
+    if (!pongReceived) {
+      if (doc.conns.has(conn)) {
+        closeConn(doc, conn);
+      }
+      clearInterval(pingInterval);
+    } else if (doc.conns.has(conn)) {
+      pongReceived = false;
+      try {
+        conn.ping();
+      } catch (e) {
+        closeConn(doc, conn);
+        clearInterval(pingInterval);
+      }
+    }
+  }, pingTimeout);
+  conn.on('close', () => {
+    closeConn(doc, conn);
+    clearInterval(pingInterval);
+  });
+  conn.on('pong', () => {
+    pongReceived = true;
+  });
+  // put the following in a variables in a block so the interval handlers don't keep in in
+  // scope
+  {
+    // send sync step 1
+    const encoder = encoding.createEncoder();
+    encoding.writeVarUint(encoder, messageSync);
+    syncProtocol.writeSyncStep1(encoder, doc);
+    send(doc, conn, encoding.toUint8Array(encoder));
+    const awarenessStates = doc.awareness.getStates();
+    if (awarenessStates.size > 0) {
+      const encoder = encoding.createEncoder();
+      encoding.writeVarUint(encoder, messageAwareness);
+      encoding.writeVarUint8Array(
+        encoder,
+        awarenessProtocol.encodeAwarenessUpdate(doc.awareness, Array.from(awarenessStates.keys()))
+      );
+      send(doc, conn, encoding.toUint8Array(encoder));
+    }
+  }
+};

docker-compose.local.yaml

@@ -30,7 +30,6 @@ services:
       - "${QUESTION_EXPRESS_DB_PORT}:5432"
     restart: unless-stopped
 
-<<<<<<< HEAD
   collab-db:
     hostname: "collab-db"
     image: postgres:16.4
@@ -44,8 +43,6 @@ services:
       - "${COLLAB_EXPRESS_DB_PORT}:5432"
     restart: unless-stopped
 
-=======
->>>>>>> ba41a4a2223602ec532fc1d08a3b632e0ca698cd
   match-db:
     hostname: "match-db"
     image: redis/redis-stack
@@ -58,29 +55,22 @@ services:
       - "${MATCHING_DB_HOST_MGMT_PORT}:8001"
       - "${MATCHING_DB_HOST_PORT}:6379"
     restart: unless-stopped
-<<<<<<< HEAD
-=======
-  
+
   # match-db-ui:
   #   hostname: "match-db-ui"
   #   image: redis/redisinsight
   #   container_name: "match-db-ui"
   #   ports:
   #     - "${MATCHING_DB_HOST_MGMT_PORT}:5540"
   #   restart: unless-stopped
->>>>>>> ba41a4a2223602ec532fc1d08a3b632e0ca698cd
+
 
 volumes:
   user-db-docker:
     external: true
   question-db-docker:
     external: true
-<<<<<<< HEAD
   collab-db-docker:
     external: true
   match-db-docker:
     external: true
-=======
-  match-db-docker:
-    external: true
->>>>>>> ba41a4a2223602ec532fc1d08a3b632e0ca698cd