From 81cceb483a391074d0d63e7d64b6324933a32103 Mon Sep 17 00:00:00 2001
From: Philipinho <16838612+Philipinho@users.noreply.github.com>
Date: Sun, 25 Jan 2026 18:59:14 +0000
Subject: [PATCH] fix: graceful collab shutdown

---
 .../adapter/collab-ws.adapter.ts              | 16 +++++++++---
 .../collaboration/collaboration.gateway.ts    | 26 ++++++++++++++++---
 .../src/collaboration/collaboration.module.ts |  9 +++----
 .../extensions/logger.extension.ts            |  8 +++---
 .../redis-sync/redis-sync.extension.ts        |  8 ++++--
 5 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/apps/server/src/collaboration/adapter/collab-ws.adapter.ts b/apps/server/src/collaboration/adapter/collab-ws.adapter.ts
index 352fe01f..18685bf0 100644
--- a/apps/server/src/collaboration/adapter/collab-ws.adapter.ts
+++ b/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);
+    }
+  }
 }
diff --git a/apps/server/src/collaboration/collaboration.gateway.ts b/apps/server/src/collaboration/collaboration.gateway.ts
index d51f8c2b..2fde2877 100644
--- a/apps/server/src/collaboration/collaboration.gateway.ts
+++ b/apps/server/src/collaboration/collaboration.gateway.ts
@@ -21,10 +21,10 @@ import { pack, unpack } from 'msgpackr';
 
 @Injectable()
 export class CollaborationGateway {
-  private hocuspocus: Hocuspocus;
+  private readonly hocuspocus: Hocuspocus;
   private redisConfig: RedisConfig;
-  private redisSync: RedisSyncExtension<{}> | null = null;
-  private useRedisSync: boolean;
+  private readonly redisSync: RedisSyncExtension<{}> | null = null;
+  private readonly useRedisSync: boolean;
 
   constructor(
     private authenticationExtension: AuthenticationExtension,
@@ -122,6 +122,24 @@ export class CollaborationGateway {
   }
 
   async destroy(): Promise<void> {
-    //await this.hocuspocus.destroy();
+    await new Promise((r) => setTimeout(r, 10000));
+    // 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('');
+          },
+        });
+
+        if (this.hocuspocus.getDocumentsCount() === 0) resolve('');
+        this.hocuspocus.closeConnections();
+      } catch (error) {
+        console.error(error);
+      }
+    });
+
+    await this.hocuspocus.hooks('onDestroy', { instance: this.hocuspocus });
   }
 }
diff --git a/apps/server/src/collaboration/collaboration.module.ts b/apps/server/src/collaboration/collaboration.module.ts
index 7c26f545..14b9d230 100644
--- a/apps/server/src/collaboration/collaboration.module.ts
+++ b/apps/server/src/collaboration/collaboration.module.ts
@@ -51,11 +51,8 @@ export class CollaborationModule implements OnModuleInit, OnModuleDestroy {
   }
 
   async onModuleDestroy(): Promise<void> {
-    if (this.collaborationGateway) {
-      await this.collaborationGateway.destroy();
-    }
-    if (this.collabWsAdapter) {
-      this.collabWsAdapter.destroy();
-    }
+    this.collabWsAdapter?.close();
+    await this.collaborationGateway?.destroy();
+    this.collabWsAdapter?.destroy();
   }
 }
diff --git a/apps/server/src/collaboration/extensions/logger.extension.ts b/apps/server/src/collaboration/extensions/logger.extension.ts
index 969fa712..bbca47bd 100644
--- a/apps/server/src/collaboration/extensions/logger.extension.ts
+++ b/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);
+  }
 }
diff --git a/apps/server/src/collaboration/extensions/redis-sync/redis-sync.extension.ts b/apps/server/src/collaboration/extensions/redis-sync/redis-sync.extension.ts
index 039f66b8..1a218282 100644
--- a/apps/server/src/collaboration/extensions/redis-sync/redis-sync.extension.ts
+++ b/apps/server/src/collaboration/extensions/redis-sync/redis-sync.extension.ts
@@ -275,7 +275,9 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
 
     const proxyTo = await this.getOrClaimLockThrottled(documentName);
     if (proxyTo && proxyTo !== this.serverId) {
-      this.logger.debug(`Doc "${documentName}" owned by server ${proxyTo}, forwarding event "${eventName}"`);
+      this.logger.debug(
+        `Doc "${documentName}" owned by server ${proxyTo}, forwarding event "${eventName}"`,
+      );
       ++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
@@ -347,7 +349,9 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
 
     const proxyTo = await this.getOrClaimLockThrottled(documentName);
     if (proxyTo && proxyTo !== this.serverId) {
-      this.logger.debug(`Doc "${documentName}" owned by server ${proxyTo}, proxying message`);
+      this.logger.debug(
+        `Doc "${documentName}" owned by server ${proxyTo}, proxying message`,
+      );
       // another server owns the doc
       const proxyMessage: RSAMessageProxy = {
         serializedHTTPRequest: serializedHTTPRequest,