refactor(workspace): sync doc update in background using data source (#4081)

packages/workspace/src/affine/sync.ts

@@ -1,15 +1,20 @@
+import { syncDataSource } from '@affine/y-provider';
 import { createIndexeddbStorage } from '@blocksuite/store';
-import { pushBinary } from '@toeverything/y-indexeddb';
+import {
+  createIndexedDBDatasource,
+  DEFAULT_DB_NAME,
+  downloadBinary,
+} from '@toeverything/y-indexeddb';
 import type { Doc } from 'yjs';
 import { applyUpdate } from 'yjs';
 
 import { createCloudBlobStorage } from '../blob/cloud-blob-storage';
-import { downloadBinaryFromCloud } from '../providers/cloud';
+import { createAffineDataSource } from '.';
 import { CRUD } from './crud';
 
 let abortController: AbortController | undefined;
 
-const downloadRecursive = async (
+const downloadRootFromIndexedDB = async (
   rootGuid: string,
   doc: Doc,
   signal: AbortSignal
@@ -17,33 +22,43 @@ const downloadRecursive = async (
   if (signal.aborted) {
     return;
   }
-  const binary = await downloadBinaryFromCloud(rootGuid, doc.guid);
-  if (typeof binary !== 'boolean') {
-    const update = new Uint8Array(binary);
-    if (rootGuid === doc.guid) {
-      // only apply the root doc
-      applyUpdate(doc, update, 'affine-cloud-service');
-    } else {
-      await pushBinary(doc.guid, update);
-    }
+  const update = await downloadBinary(rootGuid);
+  if (update !== false) {
+    applyUpdate(doc, update);
   }
-  return Promise.all(
-    [...doc.subdocs.values()].map(subdoc =>
-      downloadRecursive(rootGuid, subdoc, signal)
-    )
-  ).then();
 };
 
 export async function startSync() {
-  if (!environment.isDesktop) {
-    return;
-  }
   abortController = new AbortController();
   const signal = abortController.signal;
   const workspaces = await CRUD.list();
-  const downloadCloudPromises = workspaces.map(workspace =>
-    downloadRecursive(workspace.id, workspace.blockSuiteWorkspace.doc, signal)
+  const syncDocPromises = workspaces.map(workspace =>
+    downloadRootFromIndexedDB(
+      workspace.id,
+      workspace.blockSuiteWorkspace.doc,
+      signal
+    )
   );
+  await Promise.all(syncDocPromises);
+  const syncPromises = workspaces.map(workspace => {
+    const remoteDataSource = createAffineDataSource(
+      workspace.id,
+      workspace.blockSuiteWorkspace.doc,
+      workspace.blockSuiteWorkspace.awarenessStore.awareness
+    );
+    const indexeddbDataSource = createIndexedDBDatasource({
+      dbName: DEFAULT_DB_NAME,
+    });
+    return syncDataSource(
+      (): string[] => [
+        workspace.blockSuiteWorkspace.doc.guid,
+        ...[...workspace.blockSuiteWorkspace.doc.subdocs].map(doc => doc.guid),
+      ],
+      remoteDataSource,
+      indexeddbDataSource
+    );
+  });
+
   const syncBlobPromises = workspaces.map(async workspace => {
     const cloudBlobStorage = createCloudBlobStorage(workspace.id);
     const indexeddbBlobStorage = createIndexeddbStorage(workspace.id);
@@ -88,7 +103,7 @@ export async function startSync() {
       ]);
     });
   });
-  await Promise.all([...downloadCloudPromises, ...syncBlobPromises]);
+  await Promise.all([...syncPromises, ...syncBlobPromises]);
 }
 
 export async function stopSync() {

packages/y-indexeddb/src/provider.ts

@@ -24,7 +24,7 @@ export function setMergeCount(count: number) {
   mergeCount = count;
 }
 
-const createDatasource = ({
+export const createIndexedDBDatasource = ({
   dbName,
   mergeCount,
 }: {
@@ -109,7 +109,7 @@ export const createIndexedDBProvider = (
   doc: Doc,
   dbName: string = DEFAULT_DB_NAME
 ): IndexedDBProvider => {
-  let datasource: ReturnType<typeof createDatasource> | null = null;
+  let datasource: ReturnType<typeof createIndexedDBDatasource> | null = null;
   let provider: ReturnType<typeof createLazyProvider> | null = null;
 
   const apis = {
@@ -125,7 +125,7 @@ export const createIndexedDBProvider = (
       if (apis.connected) {
         apis.disconnect();
       }
-      datasource = createDatasource({ dbName, mergeCount });
+      datasource = createIndexedDBDatasource({ dbName, mergeCount });
       provider = createLazyProvider(doc, datasource, { origin: 'idb' });
       provider.connect();
     },

packages/y-provider/src/__tests__/index.spec.ts

@@ -3,8 +3,8 @@ import { setTimeout } from 'node:timers/promises';
 import { describe, expect, test, vi } from 'vitest';
 import { applyUpdate, Doc, encodeStateAsUpdate, encodeStateVector } from 'yjs';
 
+import type { DatasourceDocAdapter } from '../data-source';
 import { createLazyProvider } from '../lazy-provider';
-import type { DatasourceDocAdapter } from '../types';
 import { getDoc } from '../utils';
 
 const createMemoryDatasource = (rootDoc: Doc) => {
@@ -92,6 +92,16 @@ describe('y-provider', () => {
     remotesubdoc.getText('text').insert(0, 'prefix-');
     await setTimeout();
     expect(subdoc.getText('text').toJSON()).toBe('prefix-test-subdoc-value');
+
+    // disconnect then reconnect
+    provider.disconnect();
+    remotesubdoc.getText('text').delete(0, 'prefix-'.length);
+    await setTimeout();
+    expect(subdoc.getText('text').toJSON()).toBe('prefix-test-subdoc-value');
+
+    provider.connect();
+    await setTimeout();
+    expect(subdoc.getText('text').toJSON()).toBe('test-subdoc-value');
   });
 
   test('should sync a shouldLoad=true subdoc on connect', async () => {
@@ -196,4 +206,30 @@ describe('y-provider', () => {
 
     expect(spy).not.toBeCalled();
   });
+
+  test('only sync', async () => {
+    const remoteRootDoc = new Doc(); // this is the remote doc lives in remote
+    const datasource = createMemoryDatasource(remoteRootDoc);
+    remoteRootDoc.getText().insert(0, 'hello, world!');
+
+    const rootDoc = new Doc({ guid: remoteRootDoc.guid }); // this is the doc that we want to sync
+    const provider = createLazyProvider(rootDoc, datasource);
+
+    await provider.sync(true);
+    expect(rootDoc.getText().toJSON()).toBe('hello, world!');
+
+    const remotesubdoc = new Doc();
+    remotesubdoc.getText('text').insert(0, 'test-subdoc-value');
+    remoteRootDoc.getMap('map').set('subdoc', remotesubdoc);
+    expect(rootDoc.subdocs.size).toBe(0);
+
+    await provider.sync(true);
+    expect(rootDoc.subdocs.size).toBe(1);
+    const subdoc = rootDoc.getMap('map').get('subdoc') as Doc;
+    expect(subdoc.getText('text').toJSON()).toBe('');
+    await provider.sync(true);
+    expect(subdoc.getText('text').toJSON()).toBe('');
+    await provider.sync(false);
+    expect(subdoc.getText('text').toJSON()).toBe('test-subdoc-value');
+  });
 });

packages/y-provider/src/data-source.ts

@@ -0,0 +1,65 @@
+import type { DocState, StatusAdapter } from './types';
+
+export interface DatasourceDocAdapter extends Partial<StatusAdapter> {
+  /**
+   * request diff update from other clients
+   */
+  queryDocState: (
+    guid: string,
+    options?: {
+      stateVector?: Uint8Array;
+      targetClientId?: number;
+    }
+  ) => Promise<DocState | false>;
+
+  /**
+   * send update to the datasource
+   */
+  sendDocUpdate: (guid: string, update: Uint8Array) => Promise<void>;
+
+  /**
+   * listen to update from the datasource. Returns a function to unsubscribe.
+   * this is optional because some datasource might not support it
+   */
+  onDocUpdate?(
+    callback: (guid: string, update: Uint8Array) => void
+  ): () => void;
+}
+
+/**
+ * query the datasource from source, and save the latest update to target
+ *
+ * @example
+ *  bindDataSource(socketIO, indexedDB)
+ *  bindDataSource(socketIO, sqlite)
+ */
+export async function syncDataSource(
+  listDocGuids: () => string[],
+  remoteDataSource: DatasourceDocAdapter,
+  localDataSource: DatasourceDocAdapter
+) {
+  const guids = listDocGuids();
+  await Promise.all(
+    guids.map(guid => {
+      return localDataSource.queryDocState(guid).then(async docState => {
+        const remoteDocState = await (async () => {
+          if (docState) {
+            return remoteDataSource.queryDocState(guid, {
+              stateVector: docState.state,
+            });
+          } else {
+            return remoteDataSource.queryDocState(guid);
+          }
+        })();
+        if (remoteDocState) {
+          const missing = remoteDocState.missing;
+          if (missing.length === 2 && missing[0] === 0 && missing[1] === 0) {
+            // empty update
+            return;
+          }
+          await localDataSource.sendDocUpdate(guid, remoteDocState.missing);
+        }
+      });
+    })
+  );
+}

packages/y-provider/src/index.ts

@@ -1,3 +1,4 @@
+export * from './data-source';
 export * from './lazy-provider';
 export * from './types';
 export * from './utils';

packages/y-provider/src/lazy-provider.ts

@@ -1,4 +1,4 @@
-import type { PassiveDocProvider } from '@blocksuite/store';
+import { assertExists } from '@blocksuite/global/utils';
 import {
   applyUpdate,
   type Doc,
@@ -6,7 +6,8 @@ import {
   encodeStateVector,
 } from 'yjs';
 
-import type { DatasourceDocAdapter, StatusAdapter } from './types';
+import type { DatasourceDocAdapter } from './data-source';
+import type { StatusAdapter } from './types';
 import type { Status } from './types';
 
 function getDoc(doc: Doc, guid: string): Doc | undefined {
@@ -26,6 +27,17 @@ interface LazyProviderOptions {
   origin?: string;
 }
 
+export type DocProvider = {
+  // backport from `@blocksuite/store`
+  passive: true;
+
+  sync(onlyRootDoc?: boolean): Promise<void>;
+
+  get connected(): boolean;
+  connect(): void;
+  disconnect(): void;
+};
+
 /**
  * Creates a lazy provider that connects to a datasource and synchronizes a root document.
  */
@@ -33,12 +45,12 @@ export const createLazyProvider = (
   rootDoc: Doc,
   datasource: DatasourceDocAdapter,
   options: LazyProviderOptions = {}
-): Omit<PassiveDocProvider, 'flavour'> & StatusAdapter => {
+): DocProvider & StatusAdapter => {
   let connected = false;
   const pendingMap = new Map<string, Uint8Array[]>(); // guid -> pending-updates
   const disposableMap = new Map<string, Set<() => void>>();
   const connectedDocs = new Set<string>();
-  let datasourceUnsub: (() => void) | undefined;
+  let abortController: AbortController | null = null;
 
   const { origin = 'lazy-provider' } = options;
 
@@ -187,7 +199,8 @@ export const createLazyProvider = (
    * Specifically, listens for updates to documents and applies them to the corresponding Yjs document.
    */
   function setupDatasourceListeners() {
-    datasourceUnsub = datasource.onDocUpdate?.((guid, update) => {
+    assertExists(abortController, 'abortController should be defined');
+    const unsubscribe = datasource.onDocUpdate?.((guid, update) => {
       if (!connected) {
         return;
       }
@@ -208,13 +221,16 @@ export const createLazyProvider = (
         // This case happens when the father doc is not yet updated,
         //  so that the child doc is not yet created.
         //  We need to put it into cache so that it can be applied later.
-        console.warn('idb: doc not found', guid);
+        console.warn('doc not found', guid);
         pendingMap.set(guid, (pendingMap.get(guid) ?? []).concat(update));
       }
       changeStatus({
         type: 'synced',
       });
     });
+    abortController.signal.addEventListener('abort', () => {
+      unsubscribe?.();
+    });
   }
 
   // when a subdoc is loaded, we need to sync it with the datasource and setup listeners
@@ -258,6 +274,7 @@ export const createLazyProvider = (
    */
   function connect() {
     connected = true;
+    abortController = new AbortController();
 
     changeStatus({
       type: 'syncing',
@@ -292,13 +309,32 @@ export const createLazyProvider = (
       type: 'idle',
     });
     disposeAll();
-    datasourceUnsub?.();
-    datasourceUnsub = undefined;
+    assertExists(abortController, 'abortController should be defined');
+    abortController.abort();
+    abortController = null;
   }
 
+  const syncDocRecursive = async (doc: Doc) => {
+    await syncDoc(doc);
+    await Promise.all(
+      [...doc.subdocs.values()].map(subdoc => syncDocRecursive(subdoc))
+    );
+  };
+
   return {
+    sync: async onlyRootDoc => {
+      connected = true;
+      try {
+        if (onlyRootDoc) {
+          await syncDoc(rootDoc);
+        } else {
+          await syncDocRecursive(rootDoc);
+        }
+      } finally {
+        connected = false;
+      }
+    },
     get status() {
-      console.log('currentStatus', currentStatus);
       return currentStatus;
     },
     subscribeStatusChange(cb: () => void) {

packages/y-provider/src/types.ts

@@ -29,29 +29,3 @@ export interface DocState {
    */
   state?: Uint8Array;
 }
-
-export interface DatasourceDocAdapter extends Partial<StatusAdapter> {
-  /**
-   * request diff update from other clients
-   */
-  queryDocState: (
-    guid: string,
-    options?: {
-      stateVector?: Uint8Array;
-      targetClientId?: number;
-    }
-  ) => Promise<DocState | false>;
-
-  /**
-   * send update to the datasource
-   */
-  sendDocUpdate: (guid: string, update: Uint8Array) => Promise<void>;
-
-  /**
-   * listen to update from the datasource. Returns a function to unsubscribe.
-   * this is optional because some datasource might not support it
-   */
-  onDocUpdate?(
-    callback: (guid: string, update: Uint8Array) => void
-  ): () => void;
-}