feat: collab redis extension with server affinity (#1873)

* feat(collab): better redis extension
* move types to own file
* debug logging
* fix: graceful collab shutdown
* rename default prefix
* pass wsAdapter to gateway
* expose event handler
* unique collab serverId generation
* uninstall @hocuspocus/extension-redis package
* expose more functions
* sync with latest
* cleanup
* fastify router options
* cleanup type

apps/server/package.json

@@ -76,8 +76,10 @@
     "kysely-migration-cli": "^0.4.2",
     "kysely-postgres-js": "^3.0.0",
     "ldapts": "^7.4.0",
+    "lib0": "^0.2.117",
     "mammoth": "^1.11.0",
     "mime-types": "^2.1.35",
+    "msgpackr": "^1.11.8",
     "nanoid": "3.3.11",
     "nestjs-kysely": "^1.2.0",
     "nestjs-pino": "^4.5.0",
@@ -102,6 +104,7 @@
     "socket.io": "^4.8.3",
     "stripe": "^17.5.0",
     "tmp-promise": "^3.0.3",
+    "tseep": "^1.3.1",
     "typesense": "^2.1.0",
     "ws": "^8.19.0",
     "yauzl": "^3.2.0"

apps/server/src/collaboration/adapter/collab-ws.adapter.ts

@@ -30,14 +30,22 @@ export class CollabWsAdapter {
     return this.wss;
   }
 
-  public destroy() {
+  public close() {
     try {
-      this.wss.clients.forEach((client) => {
-        client.terminate();
-      });
       this.wss.close();
     } catch (err) {
       console.error(err);
     }
   }
+
+  public destroy() {
+    try {
+      this.wss.close();
+      this.wss.clients.forEach((client) => {
+        client.terminate();
+      });
+    } catch (err) {
+      console.error(err);
+    }
+  }
 }

apps/server/src/collaboration/collaboration.gateway.ts

@@ -1,10 +1,9 @@
-import { Hocuspocus, Server as HocuspocusServer } from '@hocuspocus/server';
+import { Hocuspocus } from '@hocuspocus/server';
 import { IncomingMessage } from 'http';
 import WebSocket from 'ws';
 import { AuthenticationExtension } from './extensions/authentication.extension';
 import { PersistenceExtension } from './extensions/persistence.extension';
 import { Injectable } from '@nestjs/common';
-import { Redis } from '@hocuspocus/extension-redis';
 import { EnvironmentService } from '../integrations/environment/environment.service';
 import {
   createRetryStrategy,
@@ -12,19 +11,39 @@ import {
   RedisConfig,
 } from '../common/helpers';
 import { LoggerExtension } from './extensions/logger.extension';
+import {
+  RedisSyncExtension,
+  SerializedHTTPRequest,
+} from './extensions/redis-sync';
+import { WsSocketWrapper } from './extensions/redis-sync/ws-socket-wrapper';
+import RedisClient from 'ioredis';
+import { pack, unpack } from 'msgpackr';
+import { nanoid } from 'nanoid';
+import * as os from 'node:os';
+import { CollabWsAdapter } from './adapter/collab-ws.adapter';
+import {
+  CollaborationHandler,
+  CollabEventHandlers,
+} from './collaboration.handler';
 
 @Injectable()
 export class CollaborationGateway {
-  private hocuspocus: Hocuspocus;
+  private readonly hocuspocus: Hocuspocus;
   private redisConfig: RedisConfig;
+  // @ts-ignore
+  private readonly redisSync: RedisSyncExtension<CollabEventHandlers> | null =
+    null;
+  private readonly withRedis: boolean;
 
   constructor(
     private authenticationExtension: AuthenticationExtension,
     private persistenceExtension: PersistenceExtension,
     private loggerExtension: LoggerExtension,
     private environmentService: EnvironmentService,
+    private collabEventsService: CollaborationHandler,
   ) {
     this.redisConfig = parseRedisUrl(this.environmentService.getRedisUrl());
+    this.withRedis = !this.environmentService.isCollabDisableRedis();
 
     this.hocuspocus = new Hocuspocus({
       debounce: 10000,
@@ -34,26 +53,80 @@ export class CollaborationGateway {
         this.authenticationExtension,
         this.persistenceExtension,
         this.loggerExtension,
-        ...(this.environmentService.isCollabDisableRedis()
-          ? []
-          : [
-              new Redis({
-                host: this.redisConfig.host,
-                port: this.redisConfig.port,
-                options: {
-                  password: this.redisConfig.password,
-                  db: this.redisConfig.db,
-                  family: this.redisConfig.family,
-                  retryStrategy: createRetryStrategy(),
-                },
-              }),
-            ]),
       ],
     });
+
+    if (this.withRedis) {
+      // @ts-ignore
+      this.redisSync = new RedisSyncExtension({
+        redis: new RedisClient({
+          host: this.redisConfig.host,
+          port: this.redisConfig.port,
+          password: this.redisConfig.password,
+          db: this.redisConfig.db,
+          family: this.redisConfig.family,
+          retryStrategy: createRetryStrategy(),
+        }),
+        serverId: `collab-${os?.hostname()}-${nanoid(10)}`,
+        prefix: 'collab',
+        pack,
+        unpack,
+        // @ts-ignore
+        customEvents: this.collabEventsService.getHandlers(this.hocuspocus),
+      });
+      this.hocuspocus.configuration.extensions.push(this.redisSync);
+      // @ts-ignore
+      this.redisSync.onConfigure({ instance: this.hocuspocus });
+    }
+  }
+
+  private serializeRequest(request: IncomingMessage): SerializedHTTPRequest {
+    return {
+      method: request.method ?? 'GET',
+      url: request.url ?? '/',
+      headers: {
+        'sec-websocket-key': request.headers['sec-websocket-key'] ?? '',
+        'sec-websocket-protocol':
+          request.headers['sec-websocket-protocol'] ?? '',
+      },
+      socket: { remoteAddress: request.socket?.remoteAddress ?? '' },
+    };
   }
 
   handleConnection(client: WebSocket, request: IncomingMessage): any {
-    this.hocuspocus.handleConnection(client, request);
+    if (this.redisSync) {
+      const serializedHTTPRequest = this.serializeRequest(request);
+      const socketId = serializedHTTPRequest.headers['sec-websocket-key'];
+
+      // Create wrapper socket that only receives events via emit()
+      // This prevents double-handling since Hocuspocus won't listen to raw WebSocket events
+      const wrappedSocket = new WsSocketWrapper(client);
+
+      // Route through RedisSync extension (this calls handleConnection internally)
+      this.redisSync.onSocketOpen(wrappedSocket as any, serializedHTTPRequest);
+
+      // Forward raw WebSocket messages to the extension
+      client.on('message', (data: ArrayBuffer) => {
+        this.redisSync!.onSocketMessage(
+          wrappedSocket as any,
+          serializedHTTPRequest,
+          data,
+        );
+      });
+
+      // Forward close events
+      client.on('close', (code: number, reason: Buffer) => {
+        this.redisSync!.onSocketClose(socketId, code, reason);
+      });
+
+      // Forward pong events for keepalive
+      client.on('pong', (data: Buffer) => {
+        wrappedSocket.emit('pong', data);
+      });
+    } else {
+      // Fallback to direct Hocuspocus connection
+      this.hocuspocus.handleConnection(client, request);
+    }
   }
 
   getConnectionCount() {
@@ -64,7 +137,52 @@ export class CollaborationGateway {
     return this.hocuspocus.getDocumentsCount();
   }
 
-  async destroy(): Promise<void> {
-    //await this.hocuspocus.destroy();
+  handleYjsEvent<TName extends keyof CollabEventHandlers>(
+    eventName: TName,
+    documentName: string,
+    payload: Parameters<CollabEventHandlers[TName]>[1],
+  ) {
+    return this.redisSync?.handleEvent(eventName, documentName, payload);
+  }
+
+  openDirectConnection(documentName: string, context?: any) {
+    return this.hocuspocus.openDirectConnection(documentName, context);
+  }
+
+  /*
+   *Can be used before calling openDirectConnection directly
+   */
+  async lockDocument(documentName: string) {
+    return this.redisSync.lockDocument(documentName);
+  }
+
+  /*
+   *Releases a document lock and stops the interval that maintains it.
+   */
+  async releaseLock(documentName: string) {
+    return this.redisSync.releaseLock(documentName);
+  }
+
+  async destroy(collabWsAdapter: CollabWsAdapter): Promise<void> {
+    // eslint-disable-next-line no-async-promise-executor
+    await new Promise(async (resolve) => {
+      try {
+        // Wait for all documents to unload
+        this.hocuspocus.configuration.extensions.push({
+          async afterUnloadDocument({ instance }) {
+            if (instance.getDocumentsCount() === 0) resolve('');
+          },
+        });
+
+        collabWsAdapter?.close();
+
+        if (this.hocuspocus.getDocumentsCount() === 0) resolve('');
+        this.hocuspocus.closeConnections();
+      } catch (error) {
+        console.error(error);
+      }
+    });
+
+    await this.hocuspocus.hooks('onDestroy', { instance: this.hocuspocus });
   }
 }

apps/server/src/collaboration/collaboration.handler.ts

@@ -0,0 +1,42 @@
+import { Injectable, Logger } from '@nestjs/common';
+import { Hocuspocus, Document } from '@hocuspocus/server';
+
+export type CollabEventHandlers = ReturnType<
+  CollaborationHandler['getHandlers']
+>;
+
+@Injectable()
+export class CollaborationHandler {
+  private readonly logger = new Logger(CollaborationHandler.name);
+
+  constructor() {}
+
+  getHandlers(hocuspocus: Hocuspocus) {
+    return {
+      alterState: async (documentName: string, payload: { pageId: string }) => {
+        // dummy
+        // this.logger.log('Processing', documentName, payload);
+        // await this.withYdocConnection(hocuspocus, documentName, {}, (doc) => {
+        //   const fragment = doc.getXmlFragment('default');
+        //});
+      },
+    };
+  }
+
+  async withYdocConnection(
+    hocuspocus: Hocuspocus,
+    documentName: string,
+    context: any = {},
+    fn: (doc: Document) => void,
+  ): Promise<void> {
+    const connection = await hocuspocus.openDirectConnection(
+      documentName,
+      context,
+    );
+    try {
+      await connection.transact(fn);
+    } finally {
+      await connection.disconnect();
+    }
+  }
+}

apps/server/src/collaboration/collaboration.module.ts

@@ -9,6 +9,7 @@ import { WebSocket } from 'ws';
 import { TokenModule } from '../core/auth/token.module';
 import { HistoryListener } from './listeners/history.listener';
 import { LoggerExtension } from './extensions/logger.extension';
+import { CollaborationHandler } from './collaboration.handler';
 
 @Module({
   providers: [
@@ -17,6 +18,7 @@ import { LoggerExtension } from './extensions/logger.extension';
     PersistenceExtension,
     LoggerExtension,
     HistoryListener,
+    CollaborationHandler,
   ],
   exports: [CollaborationGateway],
   imports: [TokenModule],
@@ -46,16 +48,12 @@ export class CollaborationModule implements OnModuleInit, OnModuleDestroy {
     });
 
     wss.on('error', (error) =>
-      this.logger.log('WebSocket server error:', error),
+      this.logger.error('WebSocket server error:', error),
     );
   }
 
   async onModuleDestroy(): Promise<void> {
-    if (this.collaborationGateway) {
-      await this.collaborationGateway.destroy();
-    }
-    if (this.collabWsAdapter) {
-      this.collabWsAdapter.destroy();
-    }
+    await this.collaborationGateway?.destroy(this.collabWsAdapter);
+    this.collabWsAdapter?.destroy();
   }
 }

apps/server/src/collaboration/extensions/logger.extension.ts

@@ -9,11 +9,11 @@ import { Injectable, Logger } from '@nestjs/common';
 export class LoggerExtension implements Extension {
   private readonly logger = new Logger('Collab' + LoggerExtension.name);
 
-  async onDisconnect(data: onDisconnectPayload) {
-    this.logger.debug(`User disconnected from "${data.documentName}".`);
-  }
-
   async afterUnloadDocument(data: onLoadDocumentPayload) {
     this.logger.debug('Unloaded ' + data.documentName + ' from memory');
   }
+
+  async onDisconnect(data: onDisconnectPayload) {
+    this.logger.debug('User disconnected from ' + data.documentName);
+  }
 }

apps/server/src/collaboration/extensions/redis-sync/collab-proxy-socket.ts

@@ -0,0 +1,70 @@
+import type RedisClient from 'ioredis';
+import { EventEmitter } from 'tseep';
+import type {
+  Pack,
+  RSAMessageClose,
+  RSAMessagePing,
+  RSAMessageSend,
+} from './redis-sync.types';
+
+export class CollabProxySocket extends EventEmitter {
+  private readonly replyTo: string;
+  private readonly serverChannel: string;
+  private readonly socketId: string;
+  private pub: RedisClient;
+  private readonly pack: Pack;
+  readyState = 1;
+
+  constructor(
+    pub: RedisClient,
+    pack: Pack,
+    replyTo: string,
+    serverChannel: string,
+    socketId: string,
+  ) {
+    super();
+    this.replyTo = replyTo;
+    this.socketId = socketId;
+    this.serverChannel = serverChannel;
+    this.pub = pub;
+    this.pack = pack;
+    this.once('close', () => {
+      this.readyState = 3;
+    });
+  }
+
+  private publish(msg: RSAMessageClose | RSAMessagePing | RSAMessageSend) {
+    this.pub.publish(this.replyTo, this.pack(msg));
+  }
+
+  close(code?: number, reason?: string) {
+    if (this.readyState !== 1) return;
+    const msg: RSAMessageClose = {
+      type: 'close',
+      code,
+      reason,
+      socketId: this.socketId,
+    };
+    this.publish(msg);
+  }
+
+  ping() {
+    if (this.readyState !== 1) return;
+    const msg: RSAMessagePing = {
+      type: 'ping',
+      socketId: this.socketId,
+      replyTo: this.serverChannel,
+    };
+    this.publish(msg);
+  }
+
+  send(message: Uint8Array) {
+    if (this.readyState !== 1) return;
+    const msg: RSAMessageSend = {
+      type: 'send',
+      socketId: this.socketId,
+      message,
+    };
+    this.publish(msg);
+  }
+}

apps/server/src/collaboration/extensions/redis-sync/index.ts

@@ -0,0 +1,2 @@
+export * from './redis-sync.extension';
+export type { SerializedHTTPRequest } from './redis-sync.extension';

apps/server/src/collaboration/extensions/redis-sync/redis-sync.extension.ts

@@ -0,0 +1,376 @@
+// Source https://github.com/ueberdosis/hocuspocus/pull/1008 - MIT
+import {
+  Extension,
+  Hocuspocus,
+  IncomingMessage,
+  afterUnloadDocumentPayload,
+  onConfigurePayload,
+  onLoadDocumentPayload,
+} from '@hocuspocus/server';
+import RedisClient from 'ioredis';
+import { readVarString } from 'lib0/decoding.js';
+import { CollabProxySocket } from './collab-proxy-socket';
+import {
+  BaseWebSocket,
+  Configuration,
+  CustomEvents,
+  Pack,
+  RSAMessage,
+  RSAMessageCloseProxy,
+  RSAMessageCustomEventComplete,
+  RSAMessageCustomEventStart,
+  RSAMessagePong,
+  RSAMessageProxy,
+  RSAMessageUnload,
+  SerializedHTTPRequest,
+  Unpack,
+} from './redis-sync.types';
+
+export type { Pack, SerializedHTTPRequest } from './redis-sync.types';
+
+type ServerId = string;
+type DocumentName = string;
+type SocketId = string;
+
+export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
+  priority = 1000;
+  private readonly pub: RedisClient;
+  private sub: RedisClient;
+  private readonly pack: Pack;
+  private readonly unpack: Unpack;
+  private originSockets: Record<SocketId, BaseWebSocket> = {};
+  private locks: Record<DocumentName, NodeJS.Timeout> = {};
+  private lockPromises: Record<DocumentName, Promise<ServerId | null>> = {};
+  private proxySockets: Record<SocketId, CollabProxySocket> = {};
+  private readonly prefix: string;
+  private readonly lockPrefix: string;
+  private readonly msgChannel: string;
+  private readonly serverId: ServerId;
+  private readonly customEventTTL: number;
+  private readonly lockTTL: number;
+  private instance!: Hocuspocus;
+  private readonly customEvents: TCE;
+  private replyIdCounter: number = 0;
+  // @ts-ignore
+  private pendingReplies: Record<number, PromiseWithResolvers<any>['resolve']> =
+    {};
+
+  constructor(configuration: Configuration<TCE>) {
+    const {
+      redis,
+      pack,
+      unpack,
+      serverId,
+      lockTTL,
+      prefix,
+      customEvents,
+      customEventTTL,
+    } = configuration;
+    this.pub = redis.duplicate();
+    this.sub = redis.duplicate();
+    this.pack = pack;
+    this.unpack = unpack;
+    this.serverId = serverId;
+    this.lockTTL = lockTTL ?? 10_000;
+    this.customEventTTL = customEventTTL ?? 30_000;
+    this.prefix = prefix ?? 'collab';
+    this.lockPrefix = `${this.prefix}Lock`;
+    this.msgChannel = `${this.prefix}Msg`;
+    this.customEvents = (customEvents as any) ?? ({} as any as CustomEvents);
+    this.sub.subscribe(this.msgChannel, `${this.msgChannel}:${this.serverId}`);
+    this.sub.on('messageBuffer', this.handleRedisMessage);
+  }
+  private getKey(documentName: string) {
+    return `${this.lockPrefix}:${documentName}`;
+  }
+
+  private closeProxy(socketId: string) {
+    const proxySocket = this.proxySockets[socketId];
+    if (proxySocket) {
+      proxySocket.emit(
+        'close',
+        1000,
+        Buffer.from('provider_initiated', 'utf-8'),
+      );
+      delete this.proxySockets[socketId];
+    }
+  }
+
+  private pongProxy(socketId: string) {
+    this.proxySockets[socketId]?.emit('pong');
+  }
+
+  private handleProxyMessage(
+    msg: Pick<RSAMessageProxy, 'replyTo' | 'message' | 'serializedHTTPRequest'>,
+  ) {
+    const { replyTo, message, serializedHTTPRequest } = msg;
+    const { headers } = serializedHTTPRequest;
+    const socketId = headers['sec-websocket-key']!;
+    let socket = this.proxySockets[socketId];
+    if (!socket) {
+      socket = new CollabProxySocket(
+        this.pub,
+        this.pack,
+        replyTo,
+        `${this.msgChannel}:${this.serverId}`,
+        socketId,
+      );
+      this.proxySockets[socketId] = socket;
+      this.instance.handleConnection(
+        socket as any,
+        serializedHTTPRequest as any,
+        {},
+      );
+    }
+    socket.emit('message', message);
+  }
+
+  private getOrClaimLock(documentName: string) {
+    const lockPromise = this.pub.set(
+      this.getKey(documentName),
+      this.serverId,
+      'PX',
+      this.lockTTL,
+      'NX',
+      'GET',
+    );
+    this.lockPromises[documentName] = lockPromise;
+    // Briefly cache the serverId that claimed the doc to reduce load on redis
+    // When the claimant unloads the doc, it will send an unload message to immediately clear this
+    // a lockTTL / 2 guarantees stale reads < lockTTL upon server crash
+    setTimeout(() => {
+      delete this.lockPromises[documentName];
+    }, this.lockTTL / 2);
+    return lockPromise;
+  }
+
+  private getOrClaimLockThrottled(documentName: string) {
+    const existingWorkerIdPromise = this.lockPromises[documentName];
+    if (existingWorkerIdPromise) return existingWorkerIdPromise;
+    return this.getOrClaimLock(documentName);
+  }
+
+  private handleRedisMessage = async (
+    _channel: Buffer,
+    packedMessage: Buffer,
+  ) => {
+    const msg = this.unpack(packedMessage) as RSAMessage;
+    const { type } = msg;
+    if (type === 'proxy') {
+      this.handleProxyMessage(msg);
+      return;
+    }
+    if (type === 'closeProxy') {
+      this.closeProxy(msg.socketId);
+      return;
+    }
+    if (type === 'pong') {
+      this.pongProxy(msg.socketId);
+      return;
+    }
+    if (type === 'unload') {
+      delete this.lockPromises[msg.documentName];
+      return;
+    }
+    if (type === 'customEventStart') {
+      const { documentName, eventName, payload, replyTo, replyId } = msg;
+      const res = await this.handleEventLocally(
+        eventName as Extract<keyof TCE, string>,
+        documentName,
+        payload,
+      );
+      const reply: RSAMessageCustomEventComplete = {
+        type: 'customEventComplete',
+        replyId,
+        payload: res,
+      };
+      this.pub.publish(`${replyTo}`, this.pack(reply));
+      return;
+    }
+    if (type === 'customEventComplete') {
+      const { replyId, payload } = msg;
+      const resolveFn = this.pendingReplies[replyId];
+      if (!resolveFn) return;
+      delete this.pendingReplies[replyId];
+      resolveFn(payload);
+      return;
+    }
+    const { socketId } = msg;
+    const socket = this.originSockets[socketId];
+    if (!socket) {
+      // origin socket already cleaned up
+      return;
+    }
+    if (type === 'close') {
+      socket.close(msg.code, msg.reason);
+    } else if (type === 'ping') {
+      // Reply instantly to the proxy socket, without forwarding to client
+      // The origin socket handles heartbeat for itself
+      const { replyTo, socketId } = msg;
+      const reply: RSAMessagePong = {
+        type: 'pong',
+        socketId,
+      };
+      this.pub.publish(`${replyTo}`, this.pack(reply));
+    } else if (type === 'send') {
+      socket.send(msg.message);
+    }
+  };
+
+  async maintainLock(documentName: string) {
+    this.locks[documentName] = setInterval(() => {
+      this.pub.set(
+        this.getKey(documentName),
+        this.serverId,
+        'PX',
+        this.lockTTL,
+      );
+    }, this.lockTTL / 2);
+  }
+
+  async releaseLock(documentName: string) {
+    clearInterval(this.locks[documentName]);
+    delete this.locks[documentName];
+    return this.pub.del(this.getKey(documentName));
+  }
+
+  private async handleEventLocally<TName extends Extract<keyof TCE, string>>(
+    eventName: TName,
+    documentName: string,
+    payload: any,
+  ) {
+    const handler = this.customEvents[eventName];
+    if (!handler) throw new Error(`Invalid eventName: ${eventName}`);
+    const result = await handler(documentName, payload);
+    return result as Promise<ReturnType<TCE[TName]>>;
+  }
+
+  async handleEvent<TName extends Extract<keyof TCE, string>>(
+    eventName: TName,
+    documentName: string,
+    payload: any,
+  ) {
+    const isDocLoadedOnInstance = this.instance.documents.has(documentName);
+
+    if (isDocLoadedOnInstance) {
+      return this.handleEventLocally(eventName, documentName, payload);
+    }
+
+    const proxyTo = await this.getOrClaimLockThrottled(documentName);
+    if (proxyTo && proxyTo !== this.serverId) {
+      ++this.replyIdCounter; // bug in biome thinks this.replyIdCounter is not used if written on the line below
+      const replyId = this.replyIdCounter;
+      // another server owns the doc
+      const proxyMessage: RSAMessageCustomEventStart = {
+        eventName,
+        documentName,
+        payload,
+        replyTo: `${this.msgChannel}:${this.serverId}`,
+        replyId,
+        type: 'customEventStart',
+      };
+      const msg = this.pack(proxyMessage);
+      this.pub.publish(`${this.msgChannel}:${proxyTo}`, msg);
+      // @ts-ignore
+      const { promise, resolve, reject } = Promise.withResolvers();
+      this.pendingReplies[replyId] = resolve;
+      setTimeout(() => {
+        reject('TIMEOUT');
+      }, this.customEventTTL);
+      return promise as Promise<ReturnType<TCE[TName]>>;
+    }
+    // This server owns the document, but hocuspocus hasn't loaded it yet
+    return this.handleEventLocally(eventName, documentName, payload);
+  }
+
+  async lockDocument(documentName: string) {
+    const proxyTo = await this.getOrClaimLockThrottled(documentName);
+    if (proxyTo && proxyTo !== this.serverId) {
+      throw new Error(`Could not lock document: ${documentName}`);
+    }
+    this.maintainLock(documentName);
+    return () => this.releaseLock(documentName);
+  }
+
+  /* WebSocket Server Hooks */
+  onSocketOpen(
+    ws: BaseWebSocket,
+    serializedHTTPRequest: SerializedHTTPRequest,
+    context = {},
+  ) {
+    const socketId = serializedHTTPRequest.headers['sec-websocket-key']!;
+    this.originSockets[socketId] = ws;
+    this.instance.handleConnection(
+      ws as any,
+      serializedHTTPRequest as any,
+      context,
+    );
+  }
+
+  async onSocketMessage(
+    ws: BaseWebSocket,
+    serializedHTTPRequest: SerializedHTTPRequest,
+    detachableMsg: ArrayBuffer,
+  ) {
+    const message = new Uint8Array(detachableMsg.slice());
+    const tmpMsg = new IncomingMessage(detachableMsg);
+    const documentName = readVarString(tmpMsg.decoder);
+    const isDocLoadedOnInstance = this.instance.documents.has(documentName);
+
+    if (isDocLoadedOnInstance) {
+      ws.emit('message', message);
+      return;
+    }
+
+    const proxyTo = await this.getOrClaimLockThrottled(documentName);
+    if (proxyTo && proxyTo !== this.serverId) {
+      // another server owns the doc
+      const proxyMessage: RSAMessageProxy = {
+        serializedHTTPRequest: serializedHTTPRequest,
+        replyTo: `${this.msgChannel}:${this.serverId}`,
+        message,
+        type: 'proxy',
+      };
+      const msg = this.pack(proxyMessage);
+      this.pub.publish(`${this.msgChannel}:${proxyTo}`, msg);
+      return;
+    }
+    // This server owns the document, but hocuspocus hasn't loaded it yet
+    ws.emit('message', message);
+  }
+
+  onSocketClose(socketId: string, code?: number, reason?: ArrayBuffer) {
+    const socket = this.originSockets[socketId];
+    if (!socket) return;
+    // at this point the socket is considered GC'd and we cannot call close
+    // The origin socket did not set up any connections for the proxy, so none of the hooks will work if we just emit
+    socket?.emit('close', code, reason);
+    delete this.originSockets[socketId];
+    const msg: RSAMessageCloseProxy = { type: 'closeProxy', socketId };
+    this.pub.publish(this.msgChannel, this.pack(msg)).catch(() => {});
+  }
+
+  /* Hocuspocus hooks */
+  async onConfigure({ instance }: onConfigurePayload) {
+    this.instance = instance;
+  }
+
+  async onLoadDocument(data: onLoadDocumentPayload) {
+    const { documentName } = data;
+    // Refresh the lock TTL
+    this.maintainLock(documentName);
+  }
+
+  async afterUnloadDocument(data: afterUnloadDocumentPayload) {
+    const { documentName } = data;
+    this.releaseLock(documentName);
+    // Broadcast to cluster to immediately remove the cached redis value
+    const msg: RSAMessageUnload = { type: 'unload', documentName };
+    this.pub.publish(this.msgChannel, this.pack(msg));
+  }
+
+  async onDestroy() {
+    this.pub.disconnect(false);
+    this.sub.disconnect(false);
+  }
+}

apps/server/src/collaboration/extensions/redis-sync/redis-sync.types.ts

@@ -0,0 +1,121 @@
+import EventEmitter from 'node:events';
+import { IncomingHttpHeaders } from 'node:http2';
+import RedisClient from 'ioredis';
+
+export type SecondParam<T> = T extends (
+  arg1: unknown,
+  arg2: infer A,
+  ...args: unknown[]
+) => unknown
+  ? A
+  : never;
+
+export type SerializedHTTPRequest = {
+  method: string;
+  url: string;
+  headers: IncomingHttpHeaders;
+  socket: { remoteAddress: string };
+};
+
+export type RSAMessageProxy = {
+  type: 'proxy';
+  replyTo: string;
+  message: Uint8Array<ArrayBufferLike>;
+  serializedHTTPRequest: SerializedHTTPRequest;
+};
+
+export type RSAMessageCloseProxy = {
+  type: 'closeProxy';
+  socketId: string;
+};
+
+export type RSAMessageUnload = {
+  type: 'unload';
+  documentName: string;
+};
+
+export type RSAMessageClose = {
+  type: 'close';
+  code?: number;
+  reason?: string;
+  socketId: string;
+};
+
+export type RSAMessagePing = {
+  type: 'ping';
+  socketId: string;
+  replyTo: string;
+};
+
+export type RSAMessagePong = {
+  type: 'pong';
+  socketId: string;
+};
+
+export type RSAMessageSend = {
+  type: 'send';
+  // @ts-ignore
+  message: Uint8Array<ArrayBufferLike>;
+  socketId: string;
+};
+
+export type RSAMessageCustomEventStart<TName = string, TPayload = unknown> = {
+  type: 'customEventStart';
+  documentName: string;
+  eventName: TName;
+  payload: TPayload;
+  replyTo: string;
+  replyId: number;
+};
+
+export type RSAMessageCustomEventComplete = {
+  type: 'customEventComplete';
+  replyId: number;
+  payload: unknown;
+};
+
+export type RSAMessage =
+  | RSAMessageProxy
+  | RSAMessageCloseProxy
+  | RSAMessageUnload
+  | RSAMessageClose
+  | RSAMessagePing
+  | RSAMessagePong
+  | RSAMessageSend
+  | RSAMessageCustomEventStart
+  | RSAMessageCustomEventComplete;
+
+// @ts-ignore
+export type Pack = (msg: RSAMessage) => string | Buffer<ArrayBufferLike>;
+
+export type Unpack = (
+  // @ts-ignore
+  packedMessage: Uint8Array | Buffer<ArrayBufferLike>,
+) => RSAMessage;
+
+type ServerId = string;
+type DocumentName = string;
+type CustomEventName = string;
+
+export type CustomEvents = Record<
+  CustomEventName,
+  (documentName: string, payload: unknown) => Promise<unknown>
+>;
+
+export interface Configuration<TCE> {
+  redis: RedisClient;
+  pack: Pack;
+  unpack: Unpack;
+  serverId: ServerId;
+  lockTTL?: number;
+  customEventTTL?: number;
+  prefix?: string;
+  customEvents?: TCE;
+}
+
+export type BaseWebSocket = EventEmitter & {
+  readyState: number;
+  close(code?: number, reason?: string): void;
+  ping(): void;
+  send(message: Uint8Array): void;
+};

apps/server/src/collaboration/extensions/redis-sync/ws-socket-wrapper.ts

@@ -0,0 +1,47 @@
+import { EventEmitter } from 'events';
+import type WebSocket from 'ws';
+
+/**
+ * Wrapper around ws WebSocket that only receives events via emit().
+ * This prevents double-handling when used with RedisSyncExtension.
+ */
+export class WsSocketWrapper extends EventEmitter {
+  private ws: WebSocket;
+  readyState = 1;
+
+  constructor(ws: WebSocket) {
+    super();
+    this.ws = ws;
+    this.once('close', () => {
+      this.readyState = 3;
+    });
+  }
+
+  close(code?: number, reason?: string) {
+    if (this.readyState !== 1) return;
+    this.readyState = 3;
+    try {
+      this.ws.close(code, reason);
+    } catch (e) {
+      // Socket already closed
+    }
+  }
+
+  ping() {
+    if (this.readyState !== 1) return;
+    try {
+      this.ws.ping();
+    } catch (e) {
+      // Socket already closed
+    }
+  }
+
+  send(message: Uint8Array) {
+    if (this.readyState !== 1) return;
+    try {
+      this.ws.send(message);
+    } catch (e) {
+      // Socket already closed
+    }
+  }
+}

apps/server/src/collaboration/server/collab-main.ts

@@ -12,9 +12,11 @@ async function bootstrap() {
   const app = await NestFactory.create<NestFastifyApplication>(
     CollabAppModule,
     new FastifyAdapter({
-      ignoreTrailingSlash: true,
-      ignoreDuplicateSlashes: true,
-      maxParamLength: 500,
+      routerOptions: {
+        maxParamLength: 1000,
+        ignoreTrailingSlash: true,
+        ignoreDuplicateSlashes: true,
+      },
     }),
     {
       bufferLogs: true,