From 77d239ff814e9163b8594a17385606f31a4bb8ee Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Fri, 22 Dec 2023 16:17:59 +0800
Subject: [PATCH] fix(workspace): leave workspace correctly (#5379)

---
 .../workspace/src/engine/sync/engine.ts       |  2 +-
 .../workspace/src/impl/cloud/awareness.ts     | 10 ++++++-
 .../workspace/src/impl/cloud/sync/index.ts    | 27 ++++++++++---------
 3 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/packages/frontend/workspace/src/engine/sync/engine.ts b/packages/frontend/workspace/src/engine/sync/engine.ts
index f1a9ab297c..ba1a734e9f 100644
--- a/packages/frontend/workspace/src/engine/sync/engine.ts
+++ b/packages/frontend/workspace/src/engine/sync/engine.ts
@@ -92,7 +92,7 @@ export class SyncEngine {
   }
 
   canGracefulStop() {
-    return !!this.status.local && this.status.local.pendingPushUpdates > 0;
+    return !!this.status.local && this.status.local.pendingPushUpdates === 0;
   }
 
   async waitForGracefulStop(abort?: AbortSignal) {
diff --git a/packages/frontend/workspace/src/impl/cloud/awareness.ts b/packages/frontend/workspace/src/impl/cloud/awareness.ts
index cfb91ab3a9..7c2be0dc63 100644
--- a/packages/frontend/workspace/src/impl/cloud/awareness.ts
+++ b/packages/frontend/workspace/src/impl/cloud/awareness.ts
@@ -76,6 +76,11 @@ export function createCloudAwarenessProvider(
     removeAwarenessStates(awareness, [awareness.clientID], 'window unload');
   };
 
+  function handleConnect() {
+    socket.emit('client-handshake-awareness', workspaceId);
+    socket.emit('awareness-init', workspaceId);
+  }
+
   return {
     connect: () => {
       socket.on('server-awareness-broadcast', awarenessBroadcast);
@@ -86,16 +91,19 @@ export function createCloudAwarenessProvider(
 
       socket.connect();
 
+      socket.on('connect', handleConnect);
+
       socket.emit('client-handshake-awareness', workspaceId);
       socket.emit('awareness-init', workspaceId);
     },
     disconnect: () => {
+      removeAwarenessStates(awareness, [awareness.clientID], 'disconnect');
       awareness.off('update', awarenessUpdate);
       socket.emit('client-leave-awareness', workspaceId);
       socket.off('server-awareness-broadcast', awarenessBroadcast);
       socket.off('new-client-awareness-init', newClientAwarenessInitHandler);
+      socket.off('connect', handleConnect);
       window.removeEventListener('unload', windowBeforeUnloadHandler);
-      socket.disconnect();
     },
   };
 }
diff --git a/packages/frontend/workspace/src/impl/cloud/sync/index.ts b/packages/frontend/workspace/src/impl/cloud/sync/index.ts
index 187b39d83b..c5a092cce1 100644
--- a/packages/frontend/workspace/src/impl/cloud/sync/index.ts
+++ b/packages/frontend/workspace/src/impl/cloud/sync/index.ts
@@ -51,8 +51,7 @@ export function createAffineStorage(
     });
   });
 
-  // TODO: handle error
-  socket.on('connect', () => {
+  function handleConnect() {
     socket.emit(
       'client-handshake-sync',
       workspaceId,
@@ -62,10 +61,22 @@ export function createAffineStorage(
         }
       }
     );
-  });
+  }
+
+  socket.on('connect', handleConnect);
 
   socket.connect();
 
+  socket.emit(
+    'client-handshake-sync',
+    workspaceId,
+    (response: { error?: any }) => {
+      if (!response.error) {
+        syncSender.start();
+      }
+    }
+  );
+
   return {
     name: 'affine-cloud',
     async pull(docId, state) {
@@ -124,14 +135,6 @@ export function createAffineStorage(
       await syncSender.send(docId, update);
     },
     async subscribe(cb, disconnect) {
-      const response: { error?: any } = await socket
-        .timeout(10000)
-        .emitWithAck('client-handshake-sync', workspaceId);
-
-      if (response.error) {
-        throw new Error('client-handshake error, ' + response.error);
-      }
-
       const handleUpdate = async (message: {
         workspaceId: string;
         guid: string;
@@ -157,7 +160,7 @@ export function createAffineStorage(
     disconnect() {
       syncSender.stop();
       socket.emit('client-leave-sync', workspaceId);
-      socket.disconnect();
+      socket.off('connect', handleConnect);
     },
   };
 }