refactor(store): add ref count to doc

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();
     }
   }
 
@@ -397,7 +391,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

@@ -539,6 +539,7 @@ export class Doc {
   }
 
   dispose() {
+    // this._blockCollection.dispose();
     this._disposeBlockUpdated.dispose();
     this.slots.ready.dispose();
     this.slots.blockUpdated.dispose();

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;
+  }
+}

packages/presets/src/editors/editor-container.ts

@@ -1,4 +1,4 @@
-import type { BlockModel, Doc } from '@blocksuite/store';
+import type { BlockModel, Doc, RcRef } from '@blocksuite/store';
 
 import {
   BlockStdScope,
@@ -12,6 +12,7 @@ import {
   PageEditorBlockSpecs,
   ThemeProvider,
 } from '@blocksuite/blocks';
+import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
 import { SignalWatcher, Slot, WithDisposable } from '@blocksuite/global/utils';
 import { computed, signal } from '@preact/signals-core';
 import { css, html } from 'lit';
@@ -89,7 +90,7 @@ export class AffineEditorContainer
     }
   `;
 
-  private _doc = signal<Doc>();
+  private _docRef = signal<RcRef<Doc> | null>(null);
 
   private _edgelessSpecs = signal<ExtensionType[]>(EdgelessEditorBlockSpecs);
 
@@ -122,11 +123,17 @@ export class AffineEditorContainer
   };
 
   get doc() {
-    return this._doc.value as Doc;
+    if (!this._docRef.value) {
+      throw new BlockSuiteError(ErrorCode.DocCollectionError, 'Doc not found');
+    }
+    return this._docRef.value.obj;
   }
 
   set doc(doc: Doc) {
-    this._doc.value = doc;
+    if (this._docRef.value) {
+      this._docRef.value.release();
+    }
+    this._docRef.value = doc.collection.getDocRef(doc.id);
   }
 
   set edgelessSpecs(specs: ExtensionType[]) {
@@ -180,6 +187,14 @@ export class AffineEditorContainer
     );
   }
 
+  override disconnectedCallback(): void {
+    super.disconnectedCallback();
+    if (this._docRef.value) {
+      this._docRef.value.release();
+      this._docRef.value = null;
+    }
+  }
+
   override firstUpdated() {
     if (this.mode === 'page') {
       setTimeout(() => {