feat: init doc monitor (#3320)

packages/workspace/package.json

@@ -3,7 +3,7 @@
   "private": true,
   "exports": {
     "./atom": "./src/atom.ts",
-    "./utils": "./src/utils.ts",
+    "./manager": "./src/manager/index.ts",
     "./type": "./src/type.ts",
     "./migration": "./src/migration/index.ts",
     "./local/crud": "./src/local/crud.ts",

packages/workspace/src/atom.ts

@@ -1,6 +1,6 @@
 import type { WorkspaceAdapter } from '@affine/env/workspace';
 import { WorkspaceFlavour, WorkspaceVersion } from '@affine/env/workspace';
-import { createEmptyBlockSuiteWorkspace } from '@affine/workspace/utils';
+import { getOrCreateWorkspace } from '@affine/workspace/manager';
 import type { BlockHub } from '@blocksuite/blocks';
 import { assertExists } from '@blocksuite/global/utils';
 import { atom } from 'jotai';
@@ -145,7 +145,7 @@ const rootWorkspacesMetadataPromiseAtom = atom<
         meta.flavour === WorkspaceFlavour.AFFINE_CLOUD ||
         meta.flavour === WorkspaceFlavour.LOCAL
       ) {
-        createEmptyBlockSuiteWorkspace(id, meta.flavour);
+        getOrCreateWorkspace(id, meta.flavour);
       } else {
         throw new Error(`unknown flavour ${meta.flavour}`);
       }

packages/workspace/src/local/crud.ts

@@ -6,7 +6,7 @@ import { createIndexedDBProvider } from '@toeverything/y-indexeddb';
 import { createJSONStorage } from 'jotai/utils';
 import { z } from 'zod';
 
-import { createEmptyBlockSuiteWorkspace } from '../utils';
+import { getOrCreateWorkspace } from '../manager';
 
 const getStorage = () => createJSONStorage(() => localStorage);
 
@@ -41,7 +41,7 @@ export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
     if (!id) {
       return null;
     }
-    const blockSuiteWorkspace = createEmptyBlockSuiteWorkspace(
+    const blockSuiteWorkspace = getOrCreateWorkspace(
       id,
       WorkspaceFlavour.LOCAL
     );
@@ -59,7 +59,7 @@ export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
       storage.setItem(kStoreKey, []);
     const binary = BlockSuiteWorkspace.Y.encodeStateAsUpdate(doc);
     const id = nanoid();
-    const blockSuiteWorkspace = createEmptyBlockSuiteWorkspace(
+    const blockSuiteWorkspace = getOrCreateWorkspace(
       id,
       WorkspaceFlavour.LOCAL
     );

packages/workspace/src/manager/index.ts

@@ -13,9 +13,11 @@ import {
   Workspace,
 } from '@blocksuite/store';
 import { INTERNAL_BLOCKSUITE_HASH_MAP } from '@toeverything/plugin-infra/__internal__/workspace';
+import type { Doc } from 'yjs';
+import type { Transaction } from 'yjs';
 
-import { createStaticStorage } from './blob/local-static-storage';
-import { createSQLiteStorage } from './blob/sqlite-blob-storage';
+import { createStaticStorage } from '../blob/local-static-storage';
+import { createSQLiteStorage } from '../blob/sqlite-blob-storage';
 
 function setEditorFlags(workspace: Workspace) {
   Object.entries(runtimeConfig.editorFlags).forEach(([key, value]) => {
@@ -30,11 +32,57 @@ function setEditorFlags(workspace: Workspace) {
   );
 }
 
-export function createEmptyBlockSuiteWorkspace(
-  id: string,
-  flavour: WorkspaceFlavour.AFFINE_CLOUD | WorkspaceFlavour.LOCAL
-): Workspace;
-export function createEmptyBlockSuiteWorkspace(
+type UpdateCallback = (
+  update: Uint8Array,
+  origin: string | number | null,
+  doc: Doc,
+  transaction: Transaction
+) => void;
+
+type SubdocEvent = {
+  loaded: Set<Doc>;
+  removed: Set<Doc>;
+  added: Set<Doc>;
+};
+
+const docUpdateCallbackWeakMap = new WeakMap<Doc, UpdateCallback>();
+
+const createMonitor = (doc: Doc) => {
+  const onUpdate: UpdateCallback = (update, origin) => {
+    if (process.env.NODE_ENV === 'development') {
+      if (typeof origin !== 'string' && typeof origin !== 'number') {
+        console.warn(
+          'origin is not a string or number, this will cause problems in the future',
+          origin
+        );
+      }
+    } else {
+      // todo: add monitor in the future
+    }
+  };
+  docUpdateCallbackWeakMap.set(doc, onUpdate);
+  doc.on('update', onUpdate);
+  const onSubdocs = (event: SubdocEvent) => {
+    event.added.forEach(subdoc => {
+      if (!docUpdateCallbackWeakMap.has(subdoc)) {
+        createMonitor(subdoc);
+      }
+    });
+    event.removed.forEach(subdoc => {
+      if (docUpdateCallbackWeakMap.has(subdoc)) {
+        docUpdateCallbackWeakMap.delete(subdoc);
+      }
+    });
+  };
+  doc.on('subdocs', onSubdocs);
+  doc.on('destroy', () => {
+    docUpdateCallbackWeakMap.delete(doc);
+    doc.off('update', onSubdocs);
+  });
+};
+
+// if not exist, create a new workspace
+export function getOrCreateWorkspace(
   id: string,
   flavour: WorkspaceFlavour
 ): Workspace {
@@ -76,6 +124,7 @@ export function createEmptyBlockSuiteWorkspace(
   })
     .register(AffineSchemas)
     .register(__unstableSchemas);
+  createMonitor(workspace.doc);
   setEditorFlags(workspace);
   INTERNAL_BLOCKSUITE_HASH_MAP.set(id, workspace);
   return workspace;

packages/workspace/src/migration/index.ts

@@ -1,7 +1,7 @@
 import { migrateToSubdoc } from '@affine/env/blocksuite';
 import type { LocalWorkspace } from '@affine/env/workspace';
 import { WorkspaceFlavour } from '@affine/env/workspace';
-import { createEmptyBlockSuiteWorkspace } from '@affine/workspace/utils';
+import { getOrCreateWorkspace } from '@affine/workspace/manager';
 import { nanoid, Workspace } from '@blocksuite/store';
 import { createIndexeddbStorage } from '@blocksuite/store';
 const Y = Workspace.Y;
@@ -14,7 +14,7 @@ export function upgradeV1ToV2(oldWorkspace: LocalWorkspace): LocalWorkspace {
     return oldWorkspace;
   } else {
     const id = nanoid();
-    const newBlockSuiteWorkspace = createEmptyBlockSuiteWorkspace(
+    const newBlockSuiteWorkspace = getOrCreateWorkspace(
       id,
       WorkspaceFlavour.LOCAL
     );

packages/workspace/src/providers/index.ts

@@ -62,6 +62,7 @@ const createIndexedDBBackgroundProvider: DocProviderCreator = (
 };
 
 const cache: WeakMap<Doc, Uint8Array> = new WeakMap();
+const indexedDBDownloadOrigin = 'indexeddb-download-provider';
 
 const createIndexedDBDownloadProvider: DocProviderCreator = (
   id,
@@ -76,11 +77,11 @@ const createIndexedDBDownloadProvider: DocProviderCreator = (
   async function downloadBinaryRecursively(doc: Doc) {
     if (cache.has(doc)) {
       const binary = cache.get(doc) as Uint8Array;
-      Y.applyUpdate(doc, binary);
+      Y.applyUpdate(doc, binary, indexedDBDownloadOrigin);
     } else {
       const binary = await downloadBinary(doc.guid);
       if (binary) {
-        Y.applyUpdate(doc, binary);
+        Y.applyUpdate(doc, binary, indexedDBDownloadOrigin);
         cache.set(doc, binary);
       }
     }