From 39476d16ef124b58cd47f842c83fb8cba08ec286 Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Thu, 14 Mar 2024 11:29:50 +0800
Subject: [PATCH] fix(workspace): data loss (hot-fix) (#6110)

---
 .../common/workspace/src/engine/sync/peer.ts  | 11 ++++++
 .../core/src/pages/workspace/index.tsx        | 36 +++++++++++--------
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/packages/common/workspace/src/engine/sync/peer.ts b/packages/common/workspace/src/engine/sync/peer.ts
index 3c58a3f218..4f1e81329f 100644
--- a/packages/common/workspace/src/engine/sync/peer.ts
+++ b/packages/common/workspace/src/engine/sync/peer.ts
@@ -19,6 +19,7 @@ export interface SyncPeerStatus {
   loadedDocs: number;
   pendingPullUpdates: number;
   pendingPushUpdates: number;
+  rootDocLoaded: boolean;
 }
 
 /**
@@ -54,6 +55,7 @@ export class SyncPeer {
     loadedDocs: 0,
     pendingPullUpdates: 0,
     pendingPushUpdates: 0,
+    rootDocLoaded: false,
   };
   onStatusChange = new Slot<SyncPeerStatus>();
   readonly abort = new AbortController();
@@ -119,6 +121,7 @@ export class SyncPeer {
           loadedDocs: 0,
           pendingPullUpdates: 0,
           pendingPushUpdates: 0,
+          rootDocLoaded: this.status.rootDocLoaded,
         };
         await Promise.race([
           new Promise<void>(resolve => {
@@ -291,6 +294,13 @@ export class SyncPeer {
       (await this.storage.pull(doc.guid, encodeStateVector(doc))) ?? {};
     throwIfAborted(abort);
 
+    if (docData !== undefined && doc.guid === this.rootDoc.guid) {
+      this.status = {
+        ...this.status,
+        rootDocLoaded: true,
+      };
+    }
+
     if (docData) {
       applyUpdate(doc, docData, 'load');
     }
@@ -391,6 +401,7 @@ export class SyncPeer {
         this.state.pullUpdatesQueue.length + (this.state.subdocLoading ? 1 : 0),
       pendingPushUpdates:
         this.state.pushUpdatesQueue.length + (this.state.pushingUpdate ? 1 : 0),
+      rootDocLoaded: this.status.rootDocLoaded,
     };
   }
 
diff --git a/packages/frontend/core/src/pages/workspace/index.tsx b/packages/frontend/core/src/pages/workspace/index.tsx
index 1b58c24560..46f06d60f9 100644
--- a/packages/frontend/core/src/pages/workspace/index.tsx
+++ b/packages/frontend/core/src/pages/workspace/index.tsx
@@ -6,7 +6,8 @@ import {
   workspaceListLoadingStatusAtom,
   workspaceManagerAtom,
 } from '@affine/core/modules/workspace';
-import { type Workspace } from '@affine/workspace';
+import type { Workspace } from '@affine/workspace';
+import { type SyncEngineStatus } from '@affine/workspace';
 import { useAtom, useAtomValue } from 'jotai';
 import {
   type ReactElement,
@@ -75,25 +76,32 @@ export const Component = (): ReactElement => {
 
   const [workspaceIsLoading, setWorkspaceIsLoading] = useState(true);
 
+  const [syncEngineStatus, setSyncEngineStatus] = useState<SyncEngineStatus>();
+
+  useEffect(() => {
+    if (!workspace) {
+      setSyncEngineStatus(undefined);
+    }
+    return workspace?.engine.sync.onStatusChange.on(setSyncEngineStatus)
+      .dispose;
+  }, [workspace]);
+
   // hotfix: avoid doing operation, before workspace is loaded
   useEffect(() => {
     if (!workspace) {
       setWorkspaceIsLoading(true);
       return;
     }
-    const metaYMap = workspace.blockSuiteWorkspace.doc.getMap('meta');
-
-    const handleYMapChanged = () => {
-      setWorkspaceIsLoading(metaYMap.size === 0);
-    };
-
-    handleYMapChanged();
-
-    metaYMap.observe(handleYMapChanged);
-    return () => {
-      metaYMap.unobserve(handleYMapChanged);
-    };
-  }, [workspace]);
+    if (
+      [syncEngineStatus?.local, ...(syncEngineStatus?.remotes ?? [])].some(
+        p => p?.rootDocLoaded === true
+      )
+    ) {
+      setWorkspaceIsLoading(false);
+    } else {
+      setWorkspaceIsLoading(true);
+    }
+  }, [syncEngineStatus, workspace]);
 
   // if listLoading is false, we can show 404 page, otherwise we should show loading page.
   if (listLoading === false && meta === undefined) {