refactor(infra): directory structure (#4615)

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

@@ -0,0 +1,235 @@
+import { setTimeout } from 'node:timers/promises';
+
+import { describe, expect, test, vi } from 'vitest';
+import { applyUpdate, Doc, encodeStateAsUpdate, encodeStateVector } from 'yjs';
+
+import type { DocDataSource } from '../data-source';
+import { createLazyProvider } from '../lazy-provider';
+import { getDoc } from '../utils';
+
+const createMemoryDatasource = (rootDoc: Doc) => {
+  const selfUpdateOrigin = Symbol('self-origin');
+  const listeners = new Set<(guid: string, update: Uint8Array) => void>();
+
+  function trackDoc(doc: Doc) {
+    doc.on('update', (update, origin) => {
+      if (origin === selfUpdateOrigin) {
+        return;
+      }
+      for (const listener of listeners) {
+        listener(doc.guid, update);
+      }
+    });
+
+    doc.on('subdocs', () => {
+      for (const subdoc of rootDoc.subdocs) {
+        trackDoc(subdoc);
+      }
+    });
+  }
+
+  trackDoc(rootDoc);
+
+  const adapter = {
+    queryDocState: async (guid, options) => {
+      const subdoc = getDoc(rootDoc, guid);
+      if (!subdoc) {
+        return false;
+      }
+      return {
+        missing: encodeStateAsUpdate(subdoc, options?.stateVector),
+        state: encodeStateVector(subdoc),
+      };
+    },
+    sendDocUpdate: async (guid, update) => {
+      const subdoc = getDoc(rootDoc, guid);
+      if (!subdoc) {
+        return;
+      }
+      applyUpdate(subdoc, update, selfUpdateOrigin);
+    },
+    onDocUpdate: callback => {
+      listeners.add(callback);
+      return () => {
+        listeners.delete(callback);
+      };
+    },
+  } satisfies DocDataSource;
+  return {
+    rootDoc, // expose rootDoc for testing
+    ...adapter,
+  };
+};
+
+describe('y-provider', () => {
+  test('should sync a subdoc if it is loaded after connect', async () => {
+    const remoteRootDoc = new Doc(); // this is the remote doc lives in remote
+    const datasource = createMemoryDatasource(remoteRootDoc);
+
+    const remotesubdoc = new Doc();
+    remotesubdoc.getText('text').insert(0, 'test-subdoc-value');
+    // populate remote doc with simple data
+    remoteRootDoc.getMap('map').set('test-0', 'test-0-value');
+    remoteRootDoc.getMap('map').set('subdoc', remotesubdoc);
+
+    const rootDoc = new Doc({ guid: remoteRootDoc.guid }); // this is the doc that we want to sync
+    const provider = createLazyProvider(rootDoc, datasource);
+
+    provider.connect();
+
+    await setTimeout(); // wait for the provider to sync
+
+    const subdoc = rootDoc.getMap('map').get('subdoc') as Doc;
+
+    expect(rootDoc.getMap('map').get('test-0')).toBe('test-0-value');
+    expect(subdoc.getText('text').toJSON()).toBe('');
+
+    // onload, the provider should sync the subdoc
+    subdoc.load();
+    await setTimeout();
+    expect(subdoc.getText('text').toJSON()).toBe('test-subdoc-value');
+
+    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 () => {
+    const remoteRootDoc = new Doc(); // this is the remote doc lives in remote
+    const datasource = createMemoryDatasource(remoteRootDoc);
+
+    const remotesubdoc = new Doc();
+    remotesubdoc.getText('text').insert(0, 'test-subdoc-value');
+
+    // populate remote doc with simple data
+    remoteRootDoc.getMap('map').set('test-0', 'test-0-value');
+    remoteRootDoc.getMap('map').set('subdoc', remotesubdoc);
+
+    const rootDoc = new Doc({ guid: remoteRootDoc.guid }); // this is the doc that we want to sync
+    applyUpdate(rootDoc, encodeStateAsUpdate(remoteRootDoc)); // sync rootDoc with remoteRootDoc
+
+    const subdoc = rootDoc.getMap('map').get('subdoc') as Doc;
+    expect(subdoc.getText('text').toJSON()).toBe('');
+
+    subdoc.load();
+    const provider = createLazyProvider(rootDoc, datasource);
+
+    provider.connect();
+    await setTimeout(); // wait for the provider to sync
+    expect(subdoc.getText('text').toJSON()).toBe('test-subdoc-value');
+  });
+
+  test('should send existing local update to remote on connect', async () => {
+    const remoteRootDoc = new Doc(); // this is the remote doc lives in remote
+    const datasource = createMemoryDatasource(remoteRootDoc);
+
+    const rootDoc = new Doc({ guid: remoteRootDoc.guid }); // this is the doc that we want to sync
+    applyUpdate(rootDoc, encodeStateAsUpdate(remoteRootDoc)); // sync rootDoc with remoteRootDoc
+
+    rootDoc.getText('text').insert(0, 'test-value');
+    const provider = createLazyProvider(rootDoc, datasource);
+    provider.connect();
+    await setTimeout(); // wait for the provider to sync
+
+    expect(remoteRootDoc.getText('text').toJSON()).toBe('test-value');
+  });
+
+  test('should send local update to remote for subdoc after connect', async () => {
+    const remoteRootDoc = new Doc(); // this is the remote doc lives in remote
+    const datasource = createMemoryDatasource(remoteRootDoc);
+
+    const rootDoc = new Doc({ guid: remoteRootDoc.guid }); // this is the doc that we want to sync
+    const provider = createLazyProvider(rootDoc, datasource);
+
+    provider.connect();
+
+    await setTimeout(); // wait for the provider to sync
+
+    const subdoc = new Doc();
+    rootDoc.getMap('map').set('subdoc', subdoc);
+    subdoc.getText('text').insert(0, 'test-subdoc-value');
+
+    await setTimeout(); // wait for the provider to sync
+
+    const remoteSubdoc = remoteRootDoc.getMap('map').get('subdoc') as Doc;
+    expect(remoteSubdoc.getText('text').toJSON()).toBe('test-subdoc-value');
+  });
+
+  test('should not send local update to remote for subdoc after disconnect', async () => {
+    const remoteRootDoc = new Doc(); // this is the remote doc lives in remote
+    const datasource = createMemoryDatasource(remoteRootDoc);
+
+    const rootDoc = new Doc({ guid: remoteRootDoc.guid }); // this is the doc that we want to sync
+    const provider = createLazyProvider(rootDoc, datasource);
+
+    provider.connect();
+
+    await setTimeout(); // wait for the provider to sync
+
+    const subdoc = new Doc();
+    rootDoc.getMap('map').set('subdoc', subdoc);
+
+    await setTimeout(); // wait for the provider to sync
+
+    const remoteSubdoc = remoteRootDoc.getMap('map').get('subdoc') as Doc;
+    expect(remoteSubdoc.getText('text').toJSON()).toBe('');
+
+    provider.disconnect();
+    subdoc.getText('text').insert(0, 'test-subdoc-value');
+    await setTimeout();
+    expect(remoteSubdoc.getText('text').toJSON()).toBe('');
+
+    expect(provider.connected).toBe(false);
+  });
+
+  test('should not send remote update back', async () => {
+    const remoteRootDoc = new Doc(); // this is the remote doc lives in remote
+    const datasource = createMemoryDatasource(remoteRootDoc);
+    const spy = vi.spyOn(datasource, 'sendDocUpdate');
+
+    const rootDoc = new Doc({ guid: remoteRootDoc.guid }); // this is the doc that we want to sync
+    const provider = createLazyProvider(rootDoc, datasource);
+
+    provider.connect();
+
+    remoteRootDoc.getText('text').insert(0, 'test-value');
+
+    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/common/y-provider/src/data-source.ts

@@ -0,0 +1,102 @@
+import type { Doc as YDoc } from 'yjs';
+import { applyUpdate, encodeStateAsUpdate } from 'yjs';
+
+import type { DocState } from './types';
+
+export interface DocDataSource {
+  /**
+   * 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;
+}
+
+export async function syncDocFromDataSource(
+  rootDoc: YDoc,
+  datasource: DocDataSource
+) {
+  const downloadDocStateRecursively = async (doc: YDoc) => {
+    const docState = await datasource.queryDocState(doc.guid);
+    if (docState) {
+      applyUpdate(doc, docState.missing, 'sync-doc-from-datasource');
+    }
+    await Promise.all(
+      [...doc.subdocs].map(async subdoc => {
+        await downloadDocStateRecursively(subdoc);
+      })
+    );
+  };
+  await downloadDocStateRecursively(rootDoc);
+}
+
+export async function syncDataSourceFromDoc(
+  rootDoc: YDoc,
+  datasource: DocDataSource
+) {
+  const uploadDocStateRecursively = async (doc: YDoc) => {
+    await datasource.sendDocUpdate(doc.guid, encodeStateAsUpdate(doc));
+    await Promise.all(
+      [...doc.subdocs].map(async subdoc => {
+        await uploadDocStateRecursively(subdoc);
+      })
+    );
+  };
+
+  await uploadDocStateRecursively(rootDoc);
+}
+
+/**
+ * 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: DocDataSource,
+  localDataSource: DocDataSource
+) {
+  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/common/y-provider/src/index.ts

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

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

@@ -0,0 +1,380 @@
+import { assertExists } from '@blocksuite/global/utils';
+import {
+  applyUpdate,
+  type Doc,
+  encodeStateAsUpdate,
+  encodeStateVector,
+} from 'yjs';
+
+import type { DocDataSource } from './data-source';
+import type { DataSourceAdapter } from './types';
+import type { Status } from './types';
+
+function getDoc(doc: Doc, guid: string): Doc | undefined {
+  if (doc.guid === guid) {
+    return doc;
+  }
+  for (const subdoc of doc.subdocs) {
+    const found = getDoc(subdoc, guid);
+    if (found) {
+      return found;
+    }
+  }
+  return undefined;
+}
+
+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.
+ */
+export const createLazyProvider = (
+  rootDoc: Doc,
+  datasource: DocDataSource,
+  options: LazyProviderOptions = {}
+): DocProvider & DataSourceAdapter => {
+  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 abortController: AbortController | null = null;
+
+  const { origin = 'lazy-provider' } = options;
+
+  // todo: should we use a real state machine here like `xstate`?
+  let currentStatus: Status = {
+    type: 'idle',
+  };
+  let syncingStack = 0;
+  const callbackSet = new Set<() => void>();
+  const changeStatus = (newStatus: Status) => {
+    // simulate a stack, each syncing and synced should be paired
+    if (newStatus.type === 'syncing') {
+      syncingStack++;
+    } else if (newStatus.type === 'synced' || newStatus.type === 'error') {
+      syncingStack--;
+    }
+
+    if (syncingStack < 0) {
+      console.error(
+        'syncingStatus < 0, this should not happen',
+        options.origin
+      );
+    }
+
+    if (syncingStack === 0) {
+      currentStatus = newStatus;
+    }
+    if (newStatus.type !== 'synced') {
+      currentStatus = newStatus;
+    }
+    if (syncingStack === 0) {
+      if (!connected) {
+        currentStatus = {
+          type: 'idle',
+        };
+      } else {
+        currentStatus = {
+          type: 'synced',
+        };
+      }
+    }
+    callbackSet.forEach(cb => cb());
+  };
+
+  async function syncDoc(doc: Doc) {
+    const guid = doc.guid;
+    {
+      // backport from `@blocksuite/store`
+      const prefixId = guid.startsWith('space:') ? guid.slice(6) : guid;
+      const possible1 = `${rootDoc.guid}:space:${prefixId}`;
+      const possible2 = `space:${prefixId}`;
+      const update1 = await datasource.queryDocState(possible1);
+      const update2 = await datasource.queryDocState(possible2);
+      let hasUpdate = false;
+      if (
+        update1 &&
+        update1.missing.length !== 2 &&
+        update1.missing[0] !== 0 &&
+        update1.missing[1] !== 0
+      ) {
+        applyUpdate(doc, update1.missing, origin);
+        hasUpdate = true;
+      }
+      if (
+        update2 &&
+        update2.missing.length !== 2 &&
+        update2.missing[0] !== 0 &&
+        update2.missing[1] !== 0
+      ) {
+        applyUpdate(doc, update2.missing, origin);
+        hasUpdate = true;
+      }
+      if (hasUpdate) {
+        await datasource.sendDocUpdate(
+          guid,
+          encodeStateAsUpdate(
+            doc,
+            update1 ? update1.state : update2 ? update2.state : undefined
+          )
+        );
+      }
+    }
+    if (!connected) {
+      return;
+    }
+
+    changeStatus({
+      type: 'syncing',
+    });
+    const remoteUpdate = await datasource
+      .queryDocState(guid, {
+        stateVector: encodeStateVector(doc),
+      })
+      .then(remoteUpdate => {
+        changeStatus({
+          type: 'synced',
+        });
+        return remoteUpdate;
+      })
+      .catch(error => {
+        changeStatus({
+          type: 'error',
+          error,
+        });
+        throw error;
+      });
+
+    pendingMap.set(guid, []);
+
+    if (remoteUpdate) {
+      applyUpdate(doc, remoteUpdate.missing, origin);
+    }
+
+    if (!connected) {
+      return;
+    }
+
+    // perf: optimize me
+    // it is possible the doc is only in memory but not yet in the datasource
+    // we need to send the whole update to the datasource
+    await datasource.sendDocUpdate(
+      guid,
+      encodeStateAsUpdate(doc, remoteUpdate ? remoteUpdate.state : undefined)
+    );
+
+    doc.emit('sync', []);
+  }
+
+  /**
+   * Sets up event listeners for a Yjs document.
+   * @param doc - The Yjs document to set up listeners for.
+   */
+  function setupDocListener(doc: Doc) {
+    const disposables = new Set<() => void>();
+    disposableMap.set(doc.guid, disposables);
+    const updateHandler = async (update: Uint8Array, updateOrigin: unknown) => {
+      if (origin === updateOrigin) {
+        return;
+      }
+      changeStatus({
+        type: 'syncing',
+      });
+      datasource
+        .sendDocUpdate(doc.guid, update)
+        .then(() => {
+          changeStatus({
+            type: 'synced',
+          });
+        })
+        .catch(error => {
+          changeStatus({
+            type: 'error',
+            error,
+          });
+          console.error(error);
+        });
+    };
+
+    const subdocsHandler = (event: {
+      loaded: Set<Doc>;
+      removed: Set<Doc>;
+      added: Set<Doc>;
+    }) => {
+      event.loaded.forEach(subdoc => {
+        connectDoc(subdoc).catch(console.error);
+      });
+      event.removed.forEach(subdoc => {
+        disposeDoc(subdoc);
+      });
+    };
+
+    doc.on('update', updateHandler);
+    doc.on('subdocs', subdocsHandler);
+    // todo: handle destroy?
+    disposables.add(() => {
+      doc.off('update', updateHandler);
+      doc.off('subdocs', subdocsHandler);
+    });
+  }
+
+  /**
+   * Sets up event listeners for the datasource.
+   * Specifically, listens for updates to documents and applies them to the corresponding Yjs document.
+   */
+  function setupDatasourceListeners() {
+    assertExists(abortController, 'abortController should be defined');
+    const unsubscribe = datasource.onDocUpdate?.((guid, update) => {
+      changeStatus({
+        type: 'syncing',
+      });
+      const doc = getDoc(rootDoc, guid);
+      if (doc) {
+        applyUpdate(doc, update, origin);
+        //
+        if (pendingMap.has(guid)) {
+          pendingMap
+            .get(guid)
+            ?.forEach(update => applyUpdate(doc, update, origin));
+          pendingMap.delete(guid);
+        }
+      } else {
+        // 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('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
+  async function connectDoc(doc: Doc) {
+    // skip if already connected
+    if (connectedDocs.has(doc.guid)) {
+      return;
+    }
+    connectedDocs.add(doc.guid);
+    setupDocListener(doc);
+    await syncDoc(doc);
+
+    await Promise.all(
+      [...doc.subdocs]
+        .filter(subdoc => subdoc.shouldLoad)
+        .map(subdoc => connectDoc(subdoc))
+    );
+  }
+
+  function disposeDoc(doc: Doc) {
+    connectedDocs.delete(doc.guid);
+    const disposables = disposableMap.get(doc.guid);
+    if (disposables) {
+      disposables.forEach(dispose => dispose());
+      disposableMap.delete(doc.guid);
+    }
+    // also dispose all subdocs
+    doc.subdocs.forEach(disposeDoc);
+  }
+
+  function disposeAll() {
+    disposableMap.forEach(disposables => {
+      disposables.forEach(dispose => dispose());
+    });
+    disposableMap.clear();
+    connectedDocs.clear();
+  }
+
+  /**
+   * Connects to the datasource and sets up event listeners for document updates.
+   */
+  function connect() {
+    connected = true;
+    abortController = new AbortController();
+
+    changeStatus({
+      type: 'syncing',
+    });
+    // root doc should be already loaded,
+    // but we want to populate the cache for later update events
+    connectDoc(rootDoc)
+      .then(() => {
+        changeStatus({
+          type: 'synced',
+        });
+      })
+      .catch(error => {
+        changeStatus({
+          type: 'error',
+          error,
+        });
+        console.error(error);
+      });
+    setupDatasourceListeners();
+  }
+
+  async function disconnect() {
+    connected = false;
+    disposeAll();
+    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() {
+      return currentStatus;
+    },
+    subscribeStatusChange(cb: () => void) {
+      callbackSet.add(cb);
+      return () => {
+        callbackSet.delete(cb);
+      };
+    },
+    get connected() {
+      return connected;
+    },
+    passive: true,
+    connect,
+    disconnect,
+
+    datasource,
+  };
+};

packages/common/y-provider/src/types.ts

@@ -0,0 +1,35 @@
+import type { DocDataSource } from './data-source';
+
+export type Status =
+  | {
+      type: 'idle';
+    }
+  | {
+      type: 'syncing';
+    }
+  | {
+      type: 'synced';
+    }
+  | {
+      type: 'error';
+      error: unknown;
+    };
+
+export interface DataSourceAdapter {
+  datasource: DocDataSource;
+  readonly status: Status;
+
+  subscribeStatusChange(onStatusChange: () => void): () => void;
+}
+
+export interface DocState {
+  /**
+   * The missing structs of client queries with self state.
+   */
+  missing: Uint8Array;
+
+  /**
+   * The full state of remote, used to prepare for diff sync.
+   */
+  state?: Uint8Array;
+}

packages/common/y-provider/src/utils.ts

@@ -0,0 +1,30 @@
+import type { Doc } from 'yjs';
+
+export function getDoc(doc: Doc, guid: string): Doc | undefined {
+  if (doc.guid === guid) {
+    return doc;
+  }
+  for (const subdoc of doc.subdocs) {
+    const found = getDoc(subdoc, guid);
+    if (found) {
+      return found;
+    }
+  }
+  return undefined;
+}
+
+const saveAlert = (event: BeforeUnloadEvent) => {
+  event.preventDefault();
+  return (event.returnValue =
+    'Data is not saved. Are you sure you want to leave?');
+};
+
+export const writeOperation = async (op: Promise<unknown>) => {
+  window.addEventListener('beforeunload', saveAlert, {
+    capture: true,
+  });
+  await op;
+  window.removeEventListener('beforeunload', saveAlert, {
+    capture: true,
+  });
+};