fix(workspace): fix sync stuck (#5762) (#5772)

packages/frontend/workspace-impl/src/cloud/awareness.ts

@@ -18,26 +18,70 @@ export function createCloudAwarenessProvider(
   workspaceId: string,
   awareness: Awareness
 ): AwarenessProvider {
-  const socket = getIoManager().socket('/');
+  return new AffineCloudAwarenessProvider(workspaceId, awareness);
+}
 
-  const awarenessBroadcast = ({
+class AffineCloudAwarenessProvider implements AwarenessProvider {
+  socket = getIoManager().socket('/');
+
+  constructor(
+    private readonly workspaceId: string,
+    private readonly awareness: Awareness
+  ) {}
+
+  connect(): void {
+    this.socket.on('server-awareness-broadcast', this.awarenessBroadcast);
+    this.socket.on(
+      'new-client-awareness-init',
+      this.newClientAwarenessInitHandler
+    );
+    this.awareness.on('update', this.awarenessUpdate);
+
+    window.addEventListener('beforeunload', this.windowBeforeUnloadHandler);
+
+    this.socket.on('connect', () => this.handleConnect());
+
+    if (this.socket.connected) {
+      this.handleConnect();
+    } else {
+      this.socket.connect();
+    }
+  }
+  disconnect(): void {
+    removeAwarenessStates(
+      this.awareness,
+      [this.awareness.clientID],
+      'disconnect'
+    );
+    this.awareness.off('update', this.awarenessUpdate);
+    this.socket.emit('client-leave-awareness', this.workspaceId);
+    this.socket.off('server-awareness-broadcast', this.awarenessBroadcast);
+    this.socket.off(
+      'new-client-awareness-init',
+      this.newClientAwarenessInitHandler
+    );
+    this.socket.off('connect', this.handleConnect);
+    window.removeEventListener('unload', this.windowBeforeUnloadHandler);
+  }
+
+  awarenessBroadcast = ({
     workspaceId: wsId,
     awarenessUpdate,
   }: {
     workspaceId: string;
     awarenessUpdate: string;
   }) => {
-    if (wsId !== workspaceId) {
+    if (wsId !== this.workspaceId) {
       return;
     }
     applyAwarenessUpdate(
-      awareness,
+      this.awareness,
       base64ToUint8Array(awarenessUpdate),
       'remote'
     );
   };
 
-  const awarenessUpdate = (changes: AwarenessChanges, origin: unknown) => {
+  awarenessUpdate = (changes: AwarenessChanges, origin: unknown) => {
     if (origin === 'remote') {
       return;
     }
@@ -46,63 +90,41 @@ export function createCloudAwarenessProvider(
       res.concat(cur)
     );
 
-    const update = encodeAwarenessUpdate(awareness, changedClients);
+    const update = encodeAwarenessUpdate(this.awareness, changedClients);
     uint8ArrayToBase64(update)
       .then(encodedUpdate => {
-        socket.emit('awareness-update', {
-          workspaceId: workspaceId,
+        this.socket.emit('awareness-update', {
+          workspaceId: this.workspaceId,
           awarenessUpdate: encodedUpdate,
         });
       })
       .catch(err => logger.error(err));
   };
 
-  const newClientAwarenessInitHandler = () => {
-    const awarenessUpdate = encodeAwarenessUpdate(awareness, [
-      awareness.clientID,
+  newClientAwarenessInitHandler = () => {
+    const awarenessUpdate = encodeAwarenessUpdate(this.awareness, [
+      this.awareness.clientID,
     ]);
     uint8ArrayToBase64(awarenessUpdate)
       .then(encodedAwarenessUpdate => {
-        socket.emit('awareness-update', {
-          guid: workspaceId,
+        this.socket.emit('awareness-update', {
+          guid: this.workspaceId,
           awarenessUpdate: encodedAwarenessUpdate,
         });
       })
       .catch(err => logger.error(err));
   };
 
-  const windowBeforeUnloadHandler = () => {
-    removeAwarenessStates(awareness, [awareness.clientID], 'window unload');
+  windowBeforeUnloadHandler = () => {
+    removeAwarenessStates(
+      this.awareness,
+      [this.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);
-      socket.on('new-client-awareness-init', newClientAwarenessInitHandler);
-      awareness.on('update', awarenessUpdate);
-
-      window.addEventListener('beforeunload', windowBeforeUnloadHandler);
-
-      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);
-    },
+  handleConnect = () => {
+    this.socket.emit('client-handshake-awareness', this.workspaceId);
+    this.socket.emit('awareness-init', this.workspaceId);
   };
 }

packages/frontend/workspace-impl/src/cloud/sync/batch-sync-sender.ts

@@ -1,107 +0,0 @@
-interface SyncUpdateSender {
-  (
-    guid: string,
-    updates: Uint8Array[]
-  ): Promise<{
-    accepted: boolean;
-    retry: boolean;
-  }>;
-}
-
-/**
- * BatchSyncSender is simple wrapper with vanilla update sync with several advanced features:
- *   - ACK mechanism, send updates sequentially with previous sync request correctly responds with ACK
- *   - batching updates, when waiting for previous ACK, new updates will be buffered and sent in single sync request
- *   - retryable, allow retry when previous sync request failed but with retry flag been set to true
- */
-export class BatchSyncSender {
-  private readonly buffered: Uint8Array[] = [];
-  private job: Promise<void> | null = null;
-  private started = true;
-
-  constructor(
-    private readonly guid: string,
-    private readonly rawSender: SyncUpdateSender
-  ) {}
-
-  send(update: Uint8Array) {
-    this.buffered.push(update);
-    this.next();
-    return Promise.resolve();
-  }
-
-  stop() {
-    this.started = false;
-  }
-
-  start() {
-    this.started = true;
-    this.next();
-  }
-
-  private next() {
-    if (!this.started || this.job || !this.buffered.length) {
-      return;
-    }
-
-    const lastIndex = Math.min(
-      this.buffered.length - 1,
-      99 /* max batch updates size */
-    );
-    const updates = this.buffered.slice(0, lastIndex + 1);
-
-    if (updates.length) {
-      this.job = this.rawSender(this.guid, updates)
-        .then(({ accepted, retry }) => {
-          // remove pending updates if updates are accepted
-          if (accepted) {
-            this.buffered.splice(0, lastIndex + 1);
-          }
-
-          // stop when previous sending failed and non-recoverable
-          if (accepted || retry) {
-            // avoid call stack overflow
-            setTimeout(() => {
-              this.next();
-            }, 0);
-          } else {
-            this.stop();
-          }
-        })
-        .catch(() => {
-          this.stop();
-        })
-        .finally(() => {
-          this.job = null;
-        });
-    }
-  }
-}
-
-export class MultipleBatchSyncSender {
-  private senders: Record<string, BatchSyncSender> = {};
-
-  constructor(private readonly rawSender: SyncUpdateSender) {}
-
-  async send(guid: string, update: Uint8Array) {
-    return this.getSender(guid).send(update);
-  }
-
-  private getSender(guid: string) {
-    let sender = this.senders[guid];
-    if (!sender) {
-      sender = new BatchSyncSender(guid, this.rawSender);
-      this.senders[guid] = sender;
-    }
-
-    return sender;
-  }
-
-  start() {
-    Object.values(this.senders).forEach(sender => sender.start());
-  }
-
-  stop() {
-    Object.values(this.senders).forEach(sender => sender.stop());
-  }
-}

packages/frontend/workspace-impl/src/cloud/sync/index.ts

@@ -4,165 +4,148 @@ import type { SyncStorage } from '@affine/workspace';
 
 import { getIoManager } from '../../utils/affine-io';
 import { base64ToUint8Array, uint8ArrayToBase64 } from '../../utils/base64';
-import { MultipleBatchSyncSender } from './batch-sync-sender';
 
 const logger = new DebugLogger('affine:storage:socketio');
 
+export class AffineSyncStorage implements SyncStorage {
+  name = 'affine-cloud';
+
+  SEND_TIMEOUT = 30000;
+
+  socket = getIoManager().socket('/');
+
+  constructor(private readonly workspaceId: string) {
+    this.socket.on('connect', this.handleConnect);
+
+    if (this.socket.connected) {
+      this.socket.emit('client-handshake-sync', this.workspaceId);
+    } else {
+      this.socket.connect();
+    }
+  }
+
+  handleConnect = () => {
+    this.socket.emit('client-handshake-sync', this.workspaceId);
+  };
+
+  async pull(
+    docId: string,
+    state: Uint8Array
+  ): Promise<{ data: Uint8Array; state?: Uint8Array } | null> {
+    const stateVector = state ? await uint8ArrayToBase64(state) : undefined;
+
+    logger.debug('doc-load-v2', {
+      workspaceId: this.workspaceId,
+      guid: docId,
+      stateVector,
+    });
+
+    const response:
+      | { error: any }
+      | { data: { missing: string; state: string } } = await this.socket
+      .timeout(this.SEND_TIMEOUT)
+      .emitWithAck('doc-load-v2', {
+        workspaceId: this.workspaceId,
+        guid: docId,
+        stateVector,
+      });
+
+    logger.debug('doc-load callback', {
+      workspaceId: this.workspaceId,
+      guid: docId,
+      stateVector,
+      response,
+    });
+
+    if ('error' in response) {
+      // TODO: result `EventError` with server
+      if (response.error.code === 'DOC_NOT_FOUND') {
+        return null;
+      } else {
+        throw new Error(response.error.message);
+      }
+    } else {
+      return {
+        data: base64ToUint8Array(response.data.missing),
+        state: response.data.state
+          ? base64ToUint8Array(response.data.state)
+          : undefined,
+      };
+    }
+  }
+
+  async push(docId: string, update: Uint8Array) {
+    logger.debug('client-update-v2', {
+      workspaceId: this.workspaceId,
+      guid: docId,
+      update,
+    });
+
+    const payload = await uint8ArrayToBase64(update);
+
+    const response: {
+      // TODO: reuse `EventError` with server
+      error?: any;
+      data: any;
+    } = await this.socket
+      .timeout(this.SEND_TIMEOUT)
+      .emitWithAck('client-update-v2', {
+        workspaceId: this.workspaceId,
+        guid: docId,
+        updates: [payload],
+      });
+
+    // TODO: raise error with different code to users
+    if (response.error) {
+      logger.error('client-update-v2 error', {
+        workspaceId: this.workspaceId,
+        guid: docId,
+        response,
+      });
+
+      throw new Error(response.error);
+    }
+  }
+
+  async subscribe(
+    cb: (docId: string, data: Uint8Array) => void,
+    disconnect: (reason: string) => void
+  ) {
+    const handleUpdate = async (message: {
+      workspaceId: string;
+      guid: string;
+      updates: string[];
+    }) => {
+      if (message.workspaceId === this.workspaceId) {
+        message.updates.forEach(update => {
+          cb(message.guid, base64ToUint8Array(update));
+        });
+      }
+    };
+    const handleDisconnect = (reason: string) => {
+      this.socket.off('server-updates', handleUpdate);
+      disconnect(reason);
+    };
+    this.socket.on('server-updates', handleUpdate);
+
+    this.socket.on('disconnect', handleDisconnect);
+
+    return () => {
+      this.socket.off('server-updates', handleUpdate);
+      this.socket.off('disconnect', handleDisconnect);
+    };
+  }
+
+  disconnect() {
+    this.socket.emit('client-leave-sync', this.workspaceId);
+    this.socket.off('connect', this.handleConnect);
+  }
+}
+
 export function createAffineStorage(
   workspaceId: string
 ): SyncStorage & { disconnect: () => void } {
-  logger.debug('createAffineStorage', workspaceId);
-  const socket = getIoManager().socket('/');
-
-  const syncSender = new MultipleBatchSyncSender(async (guid, updates) => {
-    const payload = await Promise.all(
-      updates.map(update => uint8ArrayToBase64(update))
-    );
-
-    return new Promise(resolve => {
-      socket.emit(
-        'client-update-v2',
-        {
-          workspaceId,
-          guid,
-          updates: payload,
-        },
-        (response: {
-          // TODO: reuse `EventError` with server
-          error?: any;
-          data: any;
-        }) => {
-          // TODO: raise error with different code to users
-          if (response.error) {
-            logger.error('client-update-v2 error', {
-              workspaceId,
-              guid,
-              response,
-            });
-          }
-
-          resolve({
-            accepted: !response.error,
-            // TODO: reuse `EventError` with server
-            retry: response.error?.code === 'INTERNAL',
-          });
-        }
-      );
-    });
-  });
-
-  function handleConnect() {
-    socket.emit(
-      'client-handshake-sync',
-      workspaceId,
-      (response: { error?: any }) => {
-        if (!response.error) {
-          syncSender.start();
-        }
-      }
-    );
-  }
-
-  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) {
-      const stateVector = state ? await uint8ArrayToBase64(state) : undefined;
-
-      return new Promise((resolve, reject) => {
-        logger.debug('doc-load-v2', {
-          workspaceId: workspaceId,
-          guid: docId,
-          stateVector,
-        });
-        socket.emit(
-          'doc-load-v2',
-          {
-            workspaceId: workspaceId,
-            guid: docId,
-            stateVector,
-          },
-          (
-            response: // TODO: reuse `EventError` with server
-            { error: any } | { data: { missing: string; state: string } }
-          ) => {
-            logger.debug('doc-load callback', {
-              workspaceId: workspaceId,
-              guid: docId,
-              stateVector,
-              response,
-            });
-
-            if ('error' in response) {
-              // TODO: result `EventError` with server
-              if (response.error.code === 'DOC_NOT_FOUND') {
-                resolve(null);
-              } else {
-                reject(new Error(response.error.message));
-              }
-            } else {
-              resolve({
-                data: base64ToUint8Array(response.data.missing),
-                state: response.data.state
-                  ? base64ToUint8Array(response.data.state)
-                  : undefined,
-              });
-            }
-          }
-        );
-      });
-    },
-    async push(docId, update) {
-      logger.debug('client-update-v2', {
-        workspaceId,
-        guid: docId,
-        update,
-      });
-
-      await syncSender.send(docId, update);
-    },
-    async subscribe(cb, disconnect) {
-      const handleUpdate = async (message: {
-        workspaceId: string;
-        guid: string;
-        updates: string[];
-      }) => {
-        if (message.workspaceId === workspaceId) {
-          message.updates.forEach(update => {
-            cb(message.guid, base64ToUint8Array(update));
-          });
-        }
-      };
-      socket.on('server-updates', handleUpdate);
-
-      socket.on('disconnect', reason => {
-        socket.off('server-updates', handleUpdate);
-        disconnect(reason);
-      });
-
-      return () => {
-        socket.off('server-updates', handleUpdate);
-      };
-    },
-    disconnect() {
-      syncSender.stop();
-      socket.emit('client-leave-sync', workspaceId);
-      socket.off('connect', handleConnect);
-    },
-  };
+  return new AffineSyncStorage(workspaceId);
 }
 
 export function createAffineStaticStorage(workspaceId: string): SyncStorage {