perf: use lazy load provider for IDB and SQLITE (#3351)

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

@@ -3,18 +3,18 @@
  */
 import 'fake-indexeddb/auto';
 
+import { setTimeout } from 'node:timers/promises';
+
 import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models';
 import { assertExists } from '@blocksuite/global/utils';
 import type { Page } from '@blocksuite/store';
 import { uuidv4, Workspace } from '@blocksuite/store';
 import { openDB } from 'idb';
 import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest';
-import { IndexeddbPersistence } from 'y-indexeddb';
 import { applyUpdate, Doc, encodeStateAsUpdate } from 'yjs';
 
 import type { WorkspacePersist } from '../index';
 import {
-  CleanupWhenConnectingError,
   createIndexedDBProvider,
   dbVersion,
   DEFAULT_DB_NAME,
@@ -43,7 +43,7 @@ function initEmptyPage(page: Page) {
 
 async function getUpdates(id: string): Promise<Uint8Array[]> {
   const db = await openDB(rootDBName, dbVersion);
-  const store = await db
+  const store = db
     .transaction('workspace', 'readonly')
     .objectStore('workspace');
   const data = (await store.get(id)) as WorkspacePersist | undefined;
@@ -74,7 +74,10 @@ describe('indexeddb provider', () => {
   test('connect', async () => {
     const provider = createIndexedDBProvider(workspace.doc);
     provider.connect();
-    await provider.whenSynced;
+
+    // todo: has a better way to know when data is synced
+    await setTimeout(200);
+
     const db = await openDB(rootDBName, dbVersion);
     {
       const store = db
@@ -96,9 +99,9 @@ describe('indexeddb provider', () => {
       const frameId = page.addBlock('affine:note', {}, pageBlockId);
       page.addBlock('affine:paragraph', {}, frameId);
     }
-    await new Promise(resolve => setTimeout(resolve, 1000));
+    await setTimeout(200);
     {
-      const store = await db
+      const store = db
         .transaction('workspace', 'readonly')
         .objectStore('workspace');
       const data = (await store.get(id)) as WorkspacePersist | undefined;
@@ -130,22 +133,11 @@ describe('indexeddb provider', () => {
     }
   });
 
-  test('disconnect suddenly', async () => {
-    const provider = createIndexedDBProvider(workspace.doc, rootDBName);
-    const fn = vi.fn();
-    provider.connect();
-    provider.disconnect();
-    expect(fn).toBeCalledTimes(0);
-    await provider.whenSynced.catch(fn);
-    expect(fn).toBeCalledTimes(1);
-  });
-
   test('connect and disconnect', async () => {
     const provider = createIndexedDBProvider(workspace.doc, rootDBName);
     provider.connect();
     expect(provider.connected).toBe(true);
-    const p1 = provider.whenSynced;
-    await p1;
+    await setTimeout(200);
     const snapshot = encodeStateAsUpdate(workspace.doc);
     provider.disconnect();
     expect(provider.connected).toBe(false);
@@ -164,8 +156,7 @@ describe('indexeddb provider', () => {
     expect(provider.connected).toBe(false);
     provider.connect();
     expect(provider.connected).toBe(true);
-    const p2 = provider.whenSynced;
-    await p2;
+    await setTimeout(200);
     {
       const updates = await getUpdates(workspace.id);
       expect(updates).not.toEqual([]);
@@ -173,13 +164,12 @@ describe('indexeddb provider', () => {
     expect(provider.connected).toBe(true);
     provider.disconnect();
     expect(provider.connected).toBe(false);
-    expect(p1).not.toBe(p2);
   });
 
   test('cleanup', async () => {
     const provider = createIndexedDBProvider(workspace.doc);
     provider.connect();
-    await provider.whenSynced;
+    await setTimeout(200);
     const db = await openDB(rootDBName, dbVersion);
 
     {
@@ -190,8 +180,8 @@ describe('indexeddb provider', () => {
       expect(keys).contain(workspace.id);
     }
 
-    provider.disconnect();
     await provider.cleanup();
+    provider.disconnect();
 
     {
       const store = db
@@ -202,17 +192,6 @@ describe('indexeddb provider', () => {
     }
   });
 
-  test('cleanup when connecting', async () => {
-    const provider = createIndexedDBProvider(workspace.doc);
-    provider.connect();
-    await expect(() => provider.cleanup()).rejects.toThrowError(
-      CleanupWhenConnectingError
-    );
-    await provider.whenSynced;
-    provider.disconnect();
-    await provider.cleanup();
-  });
-
   test('merge', async () => {
     setMergeCount(5);
     const provider = createIndexedDBProvider(workspace.doc, rootDBName);
@@ -226,7 +205,7 @@ describe('indexeddb provider', () => {
         page.addBlock('affine:paragraph', {}, frameId);
       }
     }
-    await provider.whenSynced;
+    await setTimeout(200);
     {
       const updates = await getUpdates(id);
       expect(updates.length).lessThanOrEqual(5);
@@ -242,14 +221,12 @@ describe('indexeddb provider', () => {
     {
       const provider = createIndexedDBProvider(doc, rootDBName);
       provider.connect();
-      await provider.whenSynced;
       provider.disconnect();
     }
     {
       const newDoc = new Workspace.Y.Doc();
       const provider = createIndexedDBProvider(newDoc, rootDBName);
       provider.connect();
-      await provider.whenSynced;
       provider.disconnect();
       newDoc.getMap('map').forEach((value, key) => {
         expect(value).toBe(parseInt(key));
@@ -257,42 +234,6 @@ describe('indexeddb provider', () => {
     }
   });
 
-  test('migration', async () => {
-    {
-      const yDoc = new Doc();
-      yDoc.getMap().set('foo', 'bar');
-      const persistence = new IndexeddbPersistence('test', yDoc);
-      await persistence.whenSynced;
-      await persistence.destroy();
-    }
-    {
-      const yDoc = new Doc({
-        guid: 'test',
-      });
-      const provider = createIndexedDBProvider(yDoc);
-      provider.connect();
-      await provider.whenSynced;
-      await new Promise(resolve => setTimeout(resolve, 0));
-      expect(yDoc.getMap().get('foo')).toBe('bar');
-    }
-    localStorage.clear();
-    {
-      indexedDB.databases = vi.fn(async () => {
-        throw new Error('not supported');
-      });
-      await expect(indexedDB.databases).rejects.toThrow('not supported');
-      const yDoc = new Doc({
-        guid: 'test',
-      });
-      expect(indexedDB.databases).toBeCalledTimes(1);
-      const provider = createIndexedDBProvider(yDoc);
-      provider.connect();
-      await provider.whenSynced;
-      expect(indexedDB.databases).toBeCalledTimes(2);
-      expect(yDoc.getMap().get('foo')).toBe('bar');
-    }
-  });
-
   test('beforeunload', async () => {
     const oldAddEventListener = window.addEventListener;
     window.addEventListener = vi.fn((event: string, fn, options) => {
@@ -311,7 +252,8 @@ describe('indexeddb provider', () => {
     const map = doc.getMap('map');
     map.set('1', 1);
     provider.connect();
-    await provider.whenSynced;
+
+    await setTimeout(200);
 
     expect(window.addEventListener).toBeCalledTimes(1);
     expect(window.removeEventListener).toBeCalledTimes(1);
@@ -396,7 +338,7 @@ describe('subDoc', () => {
       map.set('2', 'test');
       const provider = createIndexedDBProvider(doc);
       provider.connect();
-      await provider.whenSynced;
+      await setTimeout(200);
       provider.disconnect();
       json1 = doc.toJSON();
     }
@@ -406,7 +348,7 @@ describe('subDoc', () => {
       });
       const provider = createIndexedDBProvider(doc);
       provider.connect();
-      await provider.whenSynced;
+      await setTimeout(200);
       const map = doc.getMap();
       const subDoc = map.get('1') as Doc;
       subDoc.load();
@@ -431,7 +373,7 @@ describe('subDoc', () => {
     });
     await page1.waitForLoaded();
     const { paragraphBlockId: paragraphBlockIdPage2 } = initEmptyPage(page1);
-    await new Promise(resolve => setTimeout(resolve, 1000));
+    await setTimeout(200);
     provider.disconnect();
     {
       const newWorkspace = new Workspace({
@@ -441,15 +383,17 @@ describe('subDoc', () => {
       newWorkspace.register(AffineSchemas).register(__unstableSchemas);
       const provider = createIndexedDBProvider(newWorkspace.doc, rootDBName);
       provider.connect();
-      await provider.whenSynced;
+      await setTimeout(200);
       const page0 = newWorkspace.getPage('page0') as Page;
       await page0.waitForLoaded();
+      await setTimeout(200);
       {
         const block = page0.getBlockById(paragraphBlockIdPage1);
         assertExists(block);
       }
       const page1 = newWorkspace.getPage('page1') as Page;
       await page1.waitForLoaded();
+      await setTimeout(200);
       {
         const block = page1.getBlockById(paragraphBlockIdPage2);
         assertExists(block);
@@ -465,7 +409,7 @@ describe('utils', () => {
     initEmptyPage(page);
     const provider = createIndexedDBProvider(workspace.doc, rootDBName);
     provider.connect();
-    await provider.whenSynced;
+    await setTimeout(200);
     provider.disconnect();
     const update = (await downloadBinary(
       workspace.id,
@@ -478,16 +422,12 @@ describe('utils', () => {
     });
     newWorkspace.register(AffineSchemas).register(__unstableSchemas);
     applyUpdate(newWorkspace.doc, update);
-    await new Promise<void>(resolve =>
-      setTimeout(() => {
-        expect(workspace.doc.toJSON()['meta']).toEqual(
-          newWorkspace.doc.toJSON()['meta']
-        );
-        expect(Object.keys(workspace.doc.toJSON()['spaces'])).toEqual(
-          Object.keys(newWorkspace.doc.toJSON()['spaces'])
-        );
-        resolve();
-      }, 0)
+    await setTimeout();
+    expect(workspace.doc.toJSON()['meta']).toEqual(
+      newWorkspace.doc.toJSON()['meta']
+    );
+    expect(Object.keys(workspace.doc.toJSON()['spaces'])).toEqual(
+      Object.keys(newWorkspace.doc.toJSON()['spaces'])
     );
   });
 

packages/y-indexeddb/src/index.ts

@@ -1,26 +1,17 @@
 import { openDB } from 'idb';
 import {
   applyUpdate,
-  diffUpdate,
   Doc,
   encodeStateAsUpdate,
   encodeStateVector,
   UndoManager,
 } from 'yjs';
 
-import type {
-  BlockSuiteBinaryDB,
-  IndexedDBProvider,
-  WorkspaceMilestone,
-} from './shared';
+import type { BlockSuiteBinaryDB, WorkspaceMilestone } from './shared';
 import { dbVersion, DEFAULT_DB_NAME, upgradeDB } from './shared';
-import { tryMigrate } from './utils';
 
-const indexeddbOrigin = 'indexeddb-provider-origin';
 const snapshotOrigin = 'snapshot-origin';
 
-let mergeCount = 500;
-
 /**
  * @internal
  */
@@ -40,10 +31,6 @@ export const writeOperation = async (op: Promise<unknown>) => {
   });
 };
 
-export function setMergeCount(count: number) {
-  mergeCount = count;
-}
-
 export function revertUpdate(
   doc: Doc,
   snapshotUpdate: Uint8Array,
@@ -142,257 +129,10 @@ export const getMilestones = async (
   return milestone.milestone;
 };
 
-type SubDocsEvent = {
-  added: Set<Doc>;
-  removed: Set<Doc>;
-  loaded: Set<Doc>;
-};
-
 /**
  * We use `doc.guid` as the unique key, please make sure it not changes.
  */
-export const createIndexedDBProvider = (
-  doc: Doc,
-  dbName: string = DEFAULT_DB_NAME,
-  /**
-   * In the future, migrate will be removed and there will be a separate function
-   */
-  migrate = true
-): IndexedDBProvider => {
-  let resolve: () => void;
-  let reject: (reason?: unknown) => void;
-  let early = true;
-  let connected = false;
-  const dbPromise = openDB<BlockSuiteBinaryDB>(dbName, dbVersion, {
-    upgrade: upgradeDB,
-  });
-
-  const updateHandlerMap = new WeakMap<
-    Doc,
-    (update: Uint8Array, origin: unknown) => void
-  >();
-  const destroyHandlerMap = new WeakMap<Doc, () => void>();
-  const subDocsHandlerMap = new WeakMap<Doc, (event: SubDocsEvent) => void>();
-
-  const createOrGetHandleUpdate = (id: string, doc: Doc) => {
-    if (updateHandlerMap.has(doc)) {
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      return updateHandlerMap.get(doc)!;
-    }
-    const fn = async function handleUpdate(
-      update: Uint8Array,
-      origin: unknown
-    ) {
-      const db = await dbPromise;
-      if (!connected) {
-        return;
-      }
-      if (origin === indexeddbOrigin) {
-        return;
-      }
-      const store = db
-        .transaction('workspace', 'readwrite')
-        .objectStore('workspace');
-      let data = await store.get(id);
-      if (!data) {
-        data = {
-          id,
-          updates: [],
-        };
-      }
-      data.updates.push({
-        timestamp: Date.now(),
-        update,
-      });
-      if (data.updates.length > mergeCount) {
-        const updates = data.updates.map(({ update }) => update);
-        const doc = new Doc();
-        doc.transact(() => {
-          updates.forEach(update => {
-            applyUpdate(doc, update, indexeddbOrigin);
-          });
-        }, indexeddbOrigin);
-
-        const update = encodeStateAsUpdate(doc);
-        data = {
-          id,
-          updates: [
-            {
-              timestamp: Date.now(),
-              update,
-            },
-          ],
-        };
-        await writeOperation(store.put(data));
-      } else {
-        await writeOperation(store.put(data));
-      }
-    };
-    updateHandlerMap.set(doc, fn);
-    return fn;
-  };
-
-  /* deepscan-disable UNUSED_PARAM */
-  const createOrGetHandleDestroy = (_: string, doc: Doc) => {
-    if (destroyHandlerMap.has(doc)) {
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      return destroyHandlerMap.get(doc)!;
-    }
-    const fn = async function handleDestroy() {
-      unTrackDoc(doc.guid, doc);
-    };
-    destroyHandlerMap.set(doc, fn);
-    return fn;
-  };
-
-  /* deepscan-disable UNUSED_PARAM */
-  const createOrGetHandleSubDocs = (_: string, doc: Doc) => {
-    if (subDocsHandlerMap.has(doc)) {
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      return subDocsHandlerMap.get(doc)!;
-    }
-    const fn = async function handleSubDocs(event: SubDocsEvent) {
-      event.removed.forEach(doc => {
-        unTrackDoc(doc.guid, doc);
-      });
-      event.loaded.forEach(doc => {
-        trackDoc(doc.guid, doc);
-      });
-    };
-    subDocsHandlerMap.set(doc, fn);
-    return fn;
-  };
-
-  function trackDoc(id: string, doc: Doc) {
-    doc.on('update', createOrGetHandleUpdate(id, doc));
-    doc.on('destroy', createOrGetHandleDestroy(id, doc));
-    doc.on('subdocs', createOrGetHandleSubDocs(id, doc));
-
-    doc.subdocs.forEach(doc => {
-      trackDoc(doc.guid, doc);
-    });
-  }
-
-  function unTrackDoc(id: string, doc: Doc) {
-    doc.subdocs.forEach(doc => {
-      unTrackDoc(doc.guid, doc);
-    });
-    doc.off('update', createOrGetHandleUpdate(id, doc));
-    doc.off('destroy', createOrGetHandleDestroy(id, doc));
-    doc.off('subdocs', createOrGetHandleSubDocs(id, doc));
-  }
-
-  async function saveDocOperation(id: string, doc: Doc) {
-    const db = await dbPromise;
-    const store = db
-      .transaction('workspace', 'readwrite')
-      .objectStore('workspace');
-    const data = await store.get(id);
-    if (!connected) {
-      return;
-    }
-    if (!data) {
-      await writeOperation(
-        db.put('workspace', {
-          id,
-          updates: [
-            {
-              timestamp: Date.now(),
-              update: encodeStateAsUpdate(doc),
-            },
-          ],
-        })
-      );
-    } else {
-      const updates = data.updates.map(({ update }) => update);
-      const fakeDoc = new Doc();
-      fakeDoc.transact(() => {
-        updates.forEach(update => {
-          applyUpdate(fakeDoc, update, indexeddbOrigin);
-        });
-      }, indexeddbOrigin);
-      const newUpdate = diffUpdate(
-        encodeStateAsUpdate(doc),
-        encodeStateAsUpdate(fakeDoc)
-      );
-      await writeOperation(
-        store.put({
-          ...data,
-          updates: [
-            ...data.updates,
-            {
-              timestamp: Date.now(),
-              update: newUpdate,
-            },
-          ],
-        })
-      );
-      doc.transact(() => {
-        updates.forEach(update => {
-          applyUpdate(doc, update, indexeddbOrigin);
-        });
-      }, indexeddbOrigin);
-    }
-  }
-
-  const apis = {
-    connect: async () => {
-      if (connected) return;
-
-      apis.whenSynced = new Promise<void>((_resolve, _reject) => {
-        early = true;
-        resolve = _resolve;
-        reject = _reject;
-      });
-      connected = true;
-      trackDoc(doc.guid, doc);
-
-      // only the runs `await` below, otherwise the logic is incorrect
-      const db = await dbPromise;
-      if (migrate) {
-        // Tips:
-        //  this is only backward compatible with the yjs official version of y-indexeddb
-        await tryMigrate(db, doc.guid, dbName);
-      }
-      if (!connected) {
-        return;
-      }
-
-      // recursively save all docs into indexeddb
-      const docs: [string, Doc][] = [];
-      docs.push([doc.guid, doc]);
-      while (docs.length > 0) {
-        const [id, doc] = docs.pop() as [string, Doc];
-        await saveDocOperation(id, doc);
-        doc.subdocs.forEach(doc => {
-          docs.push([doc.guid, doc]);
-        });
-      }
-
-      early = false;
-      resolve();
-    },
-    disconnect() {
-      connected = false;
-      if (early) {
-        reject(new EarlyDisconnectError());
-      }
-      unTrackDoc(doc.guid, doc);
-    },
-    async cleanup() {
-      if (connected) {
-        throw new CleanupWhenConnectingError();
-      }
-      await (await dbPromise).delete('workspace', doc.guid);
-    },
-    whenSynced: Promise.resolve(),
-    get connected() {
-      return connected;
-    },
-  };
-
-  return apis;
-};
 
+export * from './provider';
 export * from './shared';
 export * from './utils';

packages/y-indexeddb/src/provider.ts

@@ -0,0 +1,133 @@
+import {
+  createLazyProvider,
+  type DatasourceDocAdapter,
+  writeOperation,
+} from '@affine/y-provider';
+import { openDB } from 'idb';
+import type { Doc } from 'yjs';
+import { diffUpdate, mergeUpdates } from 'yjs';
+
+import {
+  type BlockSuiteBinaryDB,
+  dbVersion,
+  DEFAULT_DB_NAME,
+  type IndexedDBProvider,
+  type UpdateMessage,
+  upgradeDB,
+} from './shared';
+
+let mergeCount = 500;
+
+export function setMergeCount(count: number) {
+  mergeCount = count;
+}
+
+const createDatasource = ({
+  dbName,
+  mergeCount,
+}: {
+  dbName: string;
+  mergeCount?: number;
+}) => {
+  const dbPromise = openDB<BlockSuiteBinaryDB>(dbName, dbVersion, {
+    upgrade: upgradeDB,
+  });
+  const adapter = {
+    queryDocState: async (guid, options) => {
+      try {
+        const db = await dbPromise;
+        const store = db
+          .transaction('workspace', 'readonly')
+          .objectStore('workspace');
+        const data = await store.get(guid);
+
+        if (!data) {
+          return false;
+        }
+
+        const { updates } = data;
+        const update = mergeUpdates(updates.map(({ update }) => update));
+
+        const diff = options?.stateVector
+          ? diffUpdate(update, options?.stateVector)
+          : update;
+
+        return diff;
+      } catch (err: any) {
+        if (!err.message?.includes('The database connection is closing.')) {
+          throw err;
+        }
+        return false;
+      }
+    },
+    sendDocUpdate: async (guid, update) => {
+      try {
+        const db = await dbPromise;
+        const store = db
+          .transaction('workspace', 'readwrite')
+          .objectStore('workspace');
+
+        // TODO: maybe we do not need to get data every time
+        const { updates } = (await store.get(guid)) ?? { updates: [] };
+        let rows: UpdateMessage[] = [
+          ...updates,
+          { timestamp: Date.now(), update },
+        ];
+        if (mergeCount && rows.length >= mergeCount) {
+          const merged = mergeUpdates(rows.map(({ update }) => update));
+          rows = [{ timestamp: Date.now(), update: merged }];
+        }
+
+        await writeOperation(
+          store.put({
+            id: guid,
+            updates: rows,
+          })
+        );
+      } catch (err: any) {
+        if (!err.message?.includes('The database connection is closing.')) {
+          throw err;
+        }
+      }
+    },
+  } satisfies DatasourceDocAdapter;
+
+  return {
+    ...adapter,
+    disconnect: () => {
+      dbPromise.then(db => db.close()).catch(console.error);
+    },
+    cleanup: async () => {
+      const db = await dbPromise;
+      await db.clear('workspace');
+    },
+  };
+};
+
+export const createIndexedDBProvider = (
+  doc: Doc,
+  dbName: string = DEFAULT_DB_NAME
+): IndexedDBProvider => {
+  let datasource: ReturnType<typeof createDatasource> | null = null;
+  let provider: ReturnType<typeof createLazyProvider> | null = null;
+
+  return {
+    connect: () => {
+      datasource = createDatasource({ dbName, mergeCount });
+      provider = createLazyProvider(doc, datasource);
+      provider.connect();
+    },
+    disconnect: () => {
+      datasource?.disconnect();
+      provider?.disconnect();
+      datasource = null;
+      provider = null;
+    },
+    cleanup: async () => {
+      await datasource?.cleanup();
+    },
+    get connected() {
+      return provider?.connected || false;
+    },
+  };
+};

packages/y-indexeddb/src/shared.ts

@@ -12,7 +12,6 @@ export interface IndexedDBProvider {
   connect: () => void;
   disconnect: () => void;
   cleanup: () => Promise<void>;
-  whenSynced: Promise<void>;
   readonly connected: boolean;
 }