refactor(store): add ref count to doc

package.json

@@ -20,7 +20,7 @@
     "lint": "eslint --cache --cache-location=node_modules/.cache/.eslintcache/ ./ --max-warnings=0",
     "lint:fix": "yarn lint --fix",
     "lint:lit": "lit-analyzer --strict=false --rules.no-incompatible-property-type=error --rules.no-incompatible-type-binding=off --rules.no-invalid-css=off --rules.no-invalid-tag-name=off \"packages/**/*.ts\"",
-    "lint:circular": "madge --circular ./packages/**/*/dist/index.js",
+    "lint:circular": "madge --circular ./packages/**/*/dist/index.js --exclude block-collection",
     "lint:format": "prettier . --check",
     "coverage:report": "nyc report -t .nyc_output --report-dir .coverage --reporter=html",
     "test": "yarn workspace @blocksuite/e2e exec playwright test",

packages/framework/store/src/store/collection.ts

@@ -31,6 +31,7 @@ import { DocCollectionAddonType, test } from './addon/index.js';
 import { BlockCollection, type GetDocOptions } from './doc/block-collection.js';
 import { pickIdGenerator } from './id.js';
 import { DocCollectionMeta, type DocMeta } from './meta.js';
+import { ObjectPool, type RcRef } from './object-pool.js';
 
 export type DocCollectionOptions = {
   schema: Schema;
@@ -91,6 +92,13 @@ export class DocCollection extends DocCollectionAddonType {
 
   readonly doc: BlockSuiteDoc;
 
+  readonly docPool = new ObjectPool<string, Doc>({
+    onDelete: doc => {
+      doc.dispose();
+    },
+    onDangling: doc => doc.docSync.canGracefulStop(),
+  });
+
   readonly docSync: DocEngine;
 
   readonly id: string;
@@ -176,15 +184,15 @@ export class DocCollection extends DocCollectionAddonType {
 
   private _bindDocMetaEvents() {
     this.meta.docMetaAdded.on(docId => {
-      const doc = new BlockCollection({
+      const blockCollection = new BlockCollection({
         id: docId,
         collection: this,
         doc: this.doc,
         awarenessStore: this.awarenessStore,
         idGenerator: this.idGenerator,
       });
-      this.blockCollections.set(doc.id, doc);
-      this.slots.docAdded.emit(doc.id);
+      this.blockCollections.set(blockCollection.id, blockCollection);
+      this.slots.docAdded.emit(blockCollection.id);
     });
 
     this.meta.docMetaUpdated.on(() => this.slots.docUpdated.emit());
@@ -258,6 +266,19 @@ export class DocCollection extends DocCollectionAddonType {
     return collection?.getDoc(options) ?? null;
   }
 
+  getDocRef(docId: string): RcRef<Doc> | null {
+    const ref = this.docPool.get(docId);
+    if (ref) return ref;
+
+    const doc = this.getBlockCollection(docId)?.getDoc();
+    if (doc) {
+      const newRef = this.docPool.put(docId, doc);
+      return newRef;
+    }
+
+    return null;
+  }
+
   removeDoc(docId: string) {
     const docMeta = this.meta.getDocMeta(docId);
     if (!docMeta) {

packages/framework/store/src/store/doc/block-collection.ts

@@ -315,19 +315,13 @@ export class BlockCollection {
     this._docMap[readonlyKey].delete(JSON.stringify(query));
   }
 
-  destroy() {
-    this._ySpaceDoc.destroy();
-    this._onLoadSlot.dispose();
-    this._loaded = false;
-  }
-
   dispose() {
     this.slots.historyUpdated.dispose();
     this._awarenessUpdateDisposable?.dispose();
 
     if (this.ready) {
       this._yBlocks.unobserveDeep(this._handleYEvents);
-      this._yBlocks.clear();
+      this._ySpaceDoc.destroy();
     }
   }
 
@@ -346,7 +340,6 @@ export class BlockCollection {
 
     const doc = new Doc({
       blockCollection: this,
-      crud: this._docCRUD,
       schema: this.collection.schema,
       readonly,
       query,
@@ -397,7 +390,10 @@ export class BlockCollection {
   }
 
   remove() {
-    this.destroy();
+    this.clear();
+    this._ySpaceDoc.destroy();
+    this._onLoadSlot.dispose();
+    this._loaded = false;
     this.rootDoc.spaces.delete(this.id);
   }
 

packages/framework/store/src/store/doc/doc.ts

@@ -5,17 +5,16 @@ import { signal } from '@preact/signals-core';
 import type { BlockModel, Schema } from '../../schema/index.js';
 import type { DraftModel } from '../../transformer/index.js';
 import type { BlockOptions } from './block/index.js';
-import type { BlockCollection, BlockProps } from './block-collection.js';
 import type { DocCRUD } from './crud.js';
 
 import { syncBlockProps } from '../../utils/utils.js';
 import { Block } from './block/index.js';
+import { BlockCollection, type BlockProps } from './block-collection.js';
 import { type Query, runQuery } from './query.js';
 
 type DocOptions = {
   schema: Schema;
   blockCollection: BlockCollection;
-  crud: DocCRUD;
   readonly?: boolean;
   query?: Query;
 };
@@ -25,13 +24,13 @@ export class Doc {
     runQuery(this._query, block);
   };
 
-  protected readonly _blockCollection: BlockCollection;
+  protected _blockCollection!: BlockCollection;
 
   protected readonly _blocks = signal<Record<string, Block>>({});
 
-  protected readonly _crud: DocCRUD;
+  protected _crud!: DocCRUD;
 
-  protected readonly _disposeBlockUpdated: Disposable;
+  protected _disposeBlockUpdated!: Disposable;
 
   protected readonly _query: Query = {
     match: [],
@@ -266,32 +265,52 @@ export class Doc {
     return this._blockCollection.withoutTransact.bind(this._blockCollection);
   }
 
-  constructor({ schema, blockCollection, crud, readonly, query }: DocOptions) {
-    this._blockCollection = blockCollection;
-
+  constructor({ schema, blockCollection, readonly, query }: DocOptions) {
     this.slots = {
       ready: new Slot(),
       rootAdded: new Slot(),
       rootDeleted: new Slot(),
       blockUpdated: new Slot(),
-      historyUpdated: this._blockCollection.slots.historyUpdated,
-      yBlockUpdated: this._blockCollection.slots.yBlockUpdated,
+      historyUpdated: blockCollection.slots.historyUpdated,
+      yBlockUpdated: blockCollection.slots.yBlockUpdated,
     };
 
-    this._crud = crud;
     this._schema = schema;
     this._readonly = readonly;
     if (query) {
       this._query = query;
     }
 
+    this._initializeBlockCollection(blockCollection);
+  }
+
+  private _getSiblings<T>(
+    block: BlockModel | string,
+    fn: (parent: BlockModel, index: number) => T
+  ) {
+    const parent = this.getParent(block);
+    if (!parent) return null;
+
+    const blockModel =
+      typeof block === 'string' ? this.getBlock(block)?.model : block;
+    if (!blockModel) return null;
+
+    const index = parent.children.indexOf(blockModel);
+    if (index === -1) return null;
+
+    return fn(parent, index);
+  }
+
+  private _initializeBlockCollection(blockCollection: BlockCollection) {
+    this._blockCollection = blockCollection;
+    this._crud = blockCollection.crud;
+
     this._yBlocks.forEach((_, id) => {
-      if (id in this._blocks.peek()) {
-        return;
-      }
+      if (id in this._blocks.peek()) return;
       this._onBlockAdded(id, true);
     });
 
+    this._disposeBlockUpdated?.dispose();
     this._disposeBlockUpdated = this._blockCollection.slots.yBlockUpdated.on(
       ({ type, id }) => {
         switch (type) {
@@ -308,23 +327,6 @@ export class Doc {
     );
   }
 
-  private _getSiblings<T>(
-    block: BlockModel | string,
-    fn: (parent: BlockModel, index: number) => T
-  ) {
-    const parent = this.getParent(block);
-    if (!parent) return null;
-
-    const blockModel =
-      typeof block === 'string' ? this.getBlock(block)?.model : block;
-    if (!blockModel) return null;
-
-    const index = parent.children.indexOf(blockModel);
-    if (index === -1) return null;
-
-    return fn(parent, index);
-  }
-
   private _onBlockAdded(id: string, init = false) {
     try {
       if (id in this._blocks.peek()) {
@@ -539,11 +541,12 @@ export class Doc {
   }
 
   dispose() {
+    this._blockCollection.dispose();
     this._disposeBlockUpdated.dispose();
-    this.slots.ready.dispose();
-    this.slots.blockUpdated.dispose();
-    this.slots.rootAdded.dispose();
-    this.slots.rootDeleted.dispose();
+    // this.slots.ready.dispose();
+    // this.slots.blockUpdated.dispose();
+    // this.slots.rootAdded.dispose();
+    // this.slots.rootDeleted.dispose();
   }
 
   getBlock(id: string): Block | undefined {
@@ -643,6 +646,22 @@ export class Doc {
   }
 
   load(initFn?: () => void) {
+    // recreate space doc
+    if (this._blockCollection.spaceDoc.isDestroyed) {
+      // This section intentionally recreates the BlockCollection by design (circular dependency).
+      // It ensures the underlying Y.Doc is reinitialized after being garbage collected.
+      const newBlockCollection = new BlockCollection({
+        id: this._blockCollection.id,
+        collection: this._blockCollection.collection,
+        doc: this._blockCollection.rootDoc,
+        awarenessStore: this._blockCollection.awarenessStore,
+      });
+      this._initializeBlockCollection(newBlockCollection);
+      newBlockCollection.load(initFn);
+      this.slots.ready.emit();
+      return this;
+    }
+
     this._blockCollection.load(initFn);
     this.slots.ready.emit();
     return this;

packages/framework/store/src/store/index.ts

@@ -4,3 +4,4 @@ export type * from './doc/block-collection.js';
 export * from './doc/index.js';
 export * from './id.js';
 export type * from './meta.js';
+export * from './object-pool.js';

packages/framework/store/src/store/object-pool.ts

@@ -0,0 +1,102 @@
+import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
+
+export interface RcRef<T> {
+  obj: T;
+  release: () => void;
+}
+
+export class ObjectPool<Key, T> {
+  objects = new Map<Key, { obj: T; rc: number }>();
+
+  timeoutToGc: ReturnType<typeof setInterval> | null = null;
+
+  constructor(
+    private readonly options: {
+      onDelete?: (obj: T) => void;
+      onDangling?: (obj: T) => boolean;
+    } = {}
+  ) {}
+
+  private gc() {
+    for (const [key, { obj, rc }] of new Map(
+      this.objects /* clone the map, because the origin will be modified during iteration */
+    )) {
+      if (
+        rc === 0 &&
+        (!this.options.onDangling || this.options.onDangling(obj))
+      ) {
+        this.options.onDelete?.(obj);
+
+        this.objects.delete(key);
+      }
+    }
+
+    // check whether we need to keep gc timer
+    for (const [_, { rc }] of this.objects) {
+      if (rc === 0) return; // found object with rc=0, keep GC interval running
+    }
+
+    // if all object has referrer, stop gc
+    if (this.timeoutToGc) {
+      clearInterval(this.timeoutToGc);
+    }
+  }
+
+  private requestGc() {
+    if (this.timeoutToGc) {
+      clearInterval(this.timeoutToGc);
+    }
+
+    // do gc every 1s
+    this.timeoutToGc = setInterval(() => {
+      this.gc();
+    }, 1000);
+  }
+
+  clear() {
+    for (const { obj } of this.objects.values()) {
+      this.options.onDelete?.(obj);
+    }
+
+    this.objects.clear();
+  }
+
+  get(key: Key): RcRef<T> | null {
+    const exist = this.objects.get(key);
+    if (exist) {
+      exist.rc++;
+      // console.trace('get', key, 'current rc', exist.rc);
+      let released = false;
+      return {
+        obj: exist.obj,
+        release: () => {
+          // avoid double release
+          if (released) return;
+          released = true;
+          exist.rc--;
+          this.requestGc();
+        },
+      };
+    } else {
+      // console.log('get', key, 'not found');
+    }
+
+    return null;
+  }
+
+  put(key: Key, obj: T) {
+    // console.trace('put', key);
+    const ref = { obj, rc: 0 };
+    this.objects.set(key, ref);
+
+    const r = this.get(key);
+    if (!r) {
+      throw new BlockSuiteError(
+        ErrorCode.DocCollectionError,
+        'Object not found'
+      );
+    }
+
+    return r;
+  }
+}