From cb6974f263e93dd6b08a1da2e35169fbc00429a8 Mon Sep 17 00:00:00 2001
From: liuyi <forehalo@gmail.com>
Date: Mon, 6 Nov 2023 18:21:19 +0800
Subject: [PATCH] fix(server): avoid server overloading by too many updates
 (#4846)

---
 .../backend/server/src/modules/doc/manager.ts | 59 +++++++++++++------
 .../server/src/modules/doc/redis-manager.ts   |  4 +-
 2 files changed, 43 insertions(+), 20 deletions(-)

diff --git a/packages/backend/server/src/modules/doc/manager.ts b/packages/backend/server/src/modules/doc/manager.ts
index e27ec546b..fde40fa8e 100644
--- a/packages/backend/server/src/modules/doc/manager.ts
+++ b/packages/backend/server/src/modules/doc/manager.ts
@@ -68,31 +68,50 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
     this.destroy();
   }
 
-  protected recoverDoc(...updates: Buffer[]): Doc {
+  protected recoverDoc(...updates: Buffer[]): Promise<Doc> {
     const doc = new Doc();
+    const chunks = chunk(updates, 100);
 
-    updates.forEach((update, i) => {
-      try {
-        if (update.length) {
-          applyUpdate(doc, update);
+    return new Promise(resolve => {
+      const next = () => {
+        const updates = chunks.shift();
+        if (updates?.length) {
+          updates.forEach(u => {
+            try {
+              applyUpdate(doc, u);
+            } catch (e) {
+              this.logger.error(
+                `Failed to apply update: ${updates
+                  .map(u => u.toString('hex'))
+                  .join('\n')}`
+              );
+            }
+          });
+
+          // avoid applying too many updates in single round which will take the whole cpu time like dead lock
+          setImmediate(() => {
+            next();
+          });
+        } else {
+          resolve(doc);
         }
-      } catch (e) {
-        this.logger.error(
-          `Failed to apply updates, index: ${i}\nUpdate: ${updates
-            .map(u => u.toString('hex'))
-            .join('\n')}`
-        );
-      }
-    });
+      };
 
-    return doc;
+      next();
+    });
   }
 
-  protected applyUpdates(guid: string, ...updates: Buffer[]): Doc {
-    const doc = this.recoverDoc(...updates);
+  protected async applyUpdates(
+    guid: string,
+    ...updates: Buffer[]
+  ): Promise<Doc> {
+    const doc = await this.recoverDoc(...updates);
 
     // test jwst codec
-    if (this.config.doc.manager.experimentalMergeWithJwstCodec) {
+    if (
+      this.config.doc.manager.experimentalMergeWithJwstCodec &&
+      updates.length < 100 /* avoid overloading */
+    ) {
       this.metrics.jwstCodecMerge(1, {});
       const yjsResult = Buffer.from(encodeStateAsUpdate(doc));
       let log = false;
@@ -312,6 +331,10 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
         workspaceId,
         id: guid,
       },
+      // take it ease, we don't want to overload db and or cpu
+      // if we limit the taken number here,
+      // user will never see the latest doc if there are too many updates pending to be merged.
+      take: 100,
     });
 
     // perf(memory): avoid sorting in db
@@ -402,7 +425,7 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
     const first = updates[0];
     const last = updates[updates.length - 1];
 
-    const doc = this.applyUpdates(
+    const doc = await this.applyUpdates(
       first.id,
       snapshot ? snapshot.blob : Buffer.from([0, 0]),
       ...updates.map(u => u.blob)
diff --git a/packages/backend/server/src/modules/doc/redis-manager.ts b/packages/backend/server/src/modules/doc/redis-manager.ts
index 4550bdd90..3ae69c3b2 100644
--- a/packages/backend/server/src/modules/doc/redis-manager.ts
+++ b/packages/backend/server/src/modules/doc/redis-manager.ts
@@ -103,9 +103,9 @@ export class RedisDocManager extends DocManager {
       const snapshot = await this.getSnapshot(docId.workspace, docId.guid);
 
       // merge
-      const doc = snapshot
+      const doc = await (snapshot
         ? this.applyUpdates(docId.full, snapshot.blob, ...updates)
-        : this.applyUpdates(docId.full, ...updates);
+        : this.applyUpdates(docId.full, ...updates));
 
       // update snapshot
       await this.upsert(docId.workspace, docId.guid, doc, snapshot?.seq);