fix(infra): memory leak (#9013)

packages/frontend/core/src/modules/docs-search/entities/docs-indexer.ts

@@ -36,7 +36,7 @@ export class DocsIndexer extends Entity {
   /**
    * increase this number to re-index all docs
    */
-  static INDEXER_VERSION = 10;
+  static INDEXER_VERSION = 11;
 
   private readonly jobQueue: JobQueue<IndexerJobPayload> =
     new IndexedDBJobQueue<IndexerJobPayload>(
@@ -85,24 +85,26 @@ export class DocsIndexer extends Entity {
   }
 
   setupListener() {
-    this.workspaceEngine.doc.storage.eventBus.on(event => {
-      if (WorkspaceDBService.isDBDocId(event.docId)) {
-        // skip db doc
-        return;
-      }
-      if (event.clientId === this.workspaceEngine.doc.clientId) {
-        this.jobQueue
-          .enqueue([
-            {
-              batchKey: event.docId,
-              payload: { storageDocId: event.docId },
-            },
-          ])
-          .catch(err => {
-            console.error('Error enqueueing job', err);
-          });
-      }
-    });
+    this.disposables.push(
+      this.workspaceEngine.doc.storage.eventBus.on(event => {
+        if (WorkspaceDBService.isDBDocId(event.docId)) {
+          // skip db doc
+          return;
+        }
+        if (event.clientId === this.workspaceEngine.doc.clientId) {
+          this.jobQueue
+            .enqueue([
+              {
+                batchKey: event.docId,
+                payload: { storageDocId: event.docId },
+              },
+            ])
+            .catch(err => {
+              console.error('Error enqueueing job', err);
+            });
+        }
+      })
+    );
   }
 
   async execJob(jobs: Job<IndexerJobPayload>[], signal: AbortSignal) {
@@ -298,6 +300,8 @@ export class DocsIndexer extends Entity {
   }
 
   override dispose(): void {
+    super.dispose();
     this.runner.stop();
+    this.worker?.dispose();
   }
 }

packages/frontend/core/src/modules/docs-search/schema.ts

@@ -1,6 +1,7 @@
 import { defineSchema } from '@toeverything/infra';
 
 export const docIndexSchema = defineSchema({
+  docId: 'String',
   title: 'FullText',
   // summary of the doc, used for preview
   summary: { type: 'String', index: false },

packages/frontend/core/src/modules/docs-search/services/docs-search.ts

@@ -632,10 +632,31 @@ export class DocsSearchService extends Service {
       );
   }
 
-  async getDocTitle(docId: string) {
-    const doc = await this.indexer.docIndex.get(docId);
-    const title = doc?.get('title');
-    return typeof title === 'string' ? title : title?.[0];
+  watchDocSummary(docId: string) {
+    return this.indexer.docIndex
+      .search$(
+        {
+          type: 'match',
+          field: 'docId',
+          match: docId,
+        },
+        {
+          fields: ['summary'],
+          pagination: {
+            limit: 1,
+          },
+        }
+      )
+      .pipe(
+        map(({ nodes }) => {
+          const node = nodes.at(0);
+          return (
+            (typeof node?.fields.summary === 'string'
+              ? node?.fields.summary
+              : node?.fields.summary[0]) ?? null
+          );
+        })
+      );
   }
 
   override dispose(): void {

packages/frontend/core/src/modules/docs-search/worker/in-worker.ts

@@ -105,6 +105,11 @@ const bookmarkFlavours = new Set([
   'affine:embed-loom',
 ]);
 
+const markdownPreviewDocCollection = new DocCollection({
+  id: 'indexer',
+  schema: blocksuiteSchema,
+});
+
 function generateMarkdownPreviewBuilder(
   yRootDoc: YDoc,
   workspaceId: string,
@@ -164,10 +169,7 @@ function generateMarkdownPreviewBuilder(
 
   const markdownAdapter = new MarkdownAdapter(
     new Job({
-      collection: new DocCollection({
-        id: 'indexer',
-        schema: blocksuiteSchema,
-      }),
+      collection: markdownPreviewDocCollection,
       middlewares: [docLinkBaseURLMiddleware, titleMiddleware],
     })
   );
@@ -875,6 +877,7 @@ async function crawlingDocData({
         {
           id: docId,
           doc: Document.from<DocIndexSchema>(docId, {
+            docId,
             title: docTitle,
             summary,
           }),

packages/frontend/core/src/modules/docs-search/worker/out-worker.ts

@@ -95,8 +95,8 @@ export async function createWorker(abort: AbortSignal) {
       });
     },
     dispose: () => {
-      worker.terminate();
       terminateAbort.abort(MANUALLY_STOP);
+      worker.terminate();
     },
   };
 }

packages/frontend/core/src/modules/permissions/entities/permission.ts

@@ -159,4 +159,8 @@ export class WorkspacePermission extends Entity {
       permission
     );
   }
+
+  override dispose(): void {
+    this.revalidate.unsubscribe();
+  }
 }

packages/frontend/core/src/modules/permissions/services/permission.ts

@@ -15,6 +15,10 @@ export class WorkspacePermissionService extends Service {
     super();
   }
 
+  override dispose(): void {
+    this.permission?.dispose();
+  }
+
   async leaveWorkspace() {
     await this.store.leaveWorkspace(this.workspaceService.workspace.id);
     this.workspacesService.list.revalidate();

packages/frontend/core/src/modules/quota/services/quota.ts

@@ -4,4 +4,8 @@ import { WorkspaceQuota } from '../entities/quota';
 
 export class WorkspaceQuotaService extends Service {
   quota = this.framework.createEntity(WorkspaceQuota);
+
+  override dispose(): void {
+    this.quota.dispose();
+  }
 }

packages/frontend/core/src/modules/share-doc/entities/share-docs-list.ts

@@ -66,4 +66,8 @@ export class ShareDocsList extends Entity {
       )
     )
   );
+
+  override dispose(): void {
+    this.revalidate.unsubscribe();
+  }
 }

packages/frontend/core/src/modules/share-doc/services/share-docs-list.ts

@@ -12,4 +12,8 @@ export class ShareDocsListService extends Service {
     this.workspaceService.workspace.flavour !== 'local'
       ? this.framework.createEntity(ShareDocsList)
       : null;
+
+  override dispose(): void {
+    this.shareDocs?.dispose();
+  }
 }

packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts

@@ -30,7 +30,7 @@ import {
 import { isEqual } from 'lodash-es';
 import { nanoid } from 'nanoid';
 import { EMPTY, map, mergeMap, Observable, switchMap } from 'rxjs';
-import { applyUpdate, encodeStateAsUpdate } from 'yjs';
+import { encodeStateAsUpdate } from 'yjs';
 
 import type { Server, ServersService } from '../../cloud';
 import {
@@ -48,6 +48,7 @@ import { CloudBlobStorage } from './engine/blob-cloud';
 import { StaticBlobStorage } from './engine/blob-static';
 import { CloudDocEngineServer } from './engine/doc-cloud';
 import { CloudStaticDocStorage } from './engine/doc-cloud-static';
+import { getWorkspaceProfileWorker } from './out-worker';
 
 const getCloudWorkspaceCacheKey = (serverId: string) => {
   if (serverId === 'affine-cloud') {
@@ -123,21 +124,25 @@ class CloudWorkspaceFlavourProvider implements WorkspaceFlavourProvider {
       },
     });
 
-    // apply initial state
-    await initial(docCollection, blobStorage, docStorage);
+    try {
+      // apply initial state
+      await initial(docCollection, blobStorage, docStorage);
 
-    // save workspace to local storage, should be vary fast
-    await docStorage.doc.set(
-      workspaceId,
-      encodeStateAsUpdate(docCollection.doc)
-    );
-    for (const subdocs of docCollection.doc.getSubdocs()) {
-      await docStorage.doc.set(subdocs.guid, encodeStateAsUpdate(subdocs));
+      // save workspace to local storage, should be vary fast
+      await docStorage.doc.set(
+        workspaceId,
+        encodeStateAsUpdate(docCollection.doc)
+      );
+      for (const subdocs of docCollection.doc.getSubdocs()) {
+        await docStorage.doc.set(subdocs.guid, encodeStateAsUpdate(subdocs));
+      }
+
+      this.revalidate();
+      await this.waitForLoaded();
+    } finally {
+      docCollection.dispose();
     }
 
-    this.revalidate();
-    await this.waitForLoaded();
-
     return {
       id: workspaceId,
       flavour: this.server.id,
@@ -229,7 +234,7 @@ class CloudWorkspaceFlavourProvider implements WorkspaceFlavourProvider {
     const docStorage = this.storageProvider.getDocStorage(id);
     // download root doc
     const localData = await docStorage.doc.get(id);
-    const cloudData = await cloudStorage.pull(id);
+    const cloudData = (await cloudStorage.pull(id))?.data;
 
     const info = await this.getWorkspaceInfo(id, signal);
 
@@ -241,17 +246,16 @@ class CloudWorkspaceFlavourProvider implements WorkspaceFlavourProvider {
       };
     }
 
-    const bs = new DocCollection({
-      id,
-      schema: getAFFiNEWorkspaceSchema(),
-    });
+    const client = getWorkspaceProfileWorker();
 
-    if (localData) applyUpdate(bs.doc, localData);
-    if (cloudData) applyUpdate(bs.doc, cloudData.data);
+    const result = await client.call(
+      'renderWorkspaceProfile',
+      [localData, cloudData].filter(Boolean) as Uint8Array[]
+    );
 
     return {
-      name: bs.meta.name,
-      avatar: bs.meta.avatar,
+      name: result.name,
+      avatar: result.avatar,
       isOwner: info.isOwner,
       isAdmin: info.isAdmin,
       isTeam: info.workspace.team,

packages/frontend/core/src/modules/workspace-engine/impls/in-worker.ts

@@ -0,0 +1,25 @@
+import { type MessageCommunicapable, OpConsumer } from '@toeverything/infra/op';
+import { applyUpdate, Doc as YDoc } from 'yjs';
+
+import type { WorkerOps } from './worker-ops';
+
+const consumer = new OpConsumer<WorkerOps>(globalThis as MessageCommunicapable);
+
+consumer.register('renderWorkspaceProfile', data => {
+  const doc = new YDoc({
+    guid: 'workspace',
+  });
+  for (const update of data) {
+    applyUpdate(doc, update);
+  }
+  const meta = doc.getMap('meta');
+  const name = meta.get('name');
+  const avatar = meta.get('avatar');
+
+  return {
+    name: typeof name === 'string' ? name : undefined,
+    avatar: typeof avatar === 'string' ? avatar : undefined,
+  };
+});
+
+consumer.listen();

packages/frontend/core/src/modules/workspace-engine/impls/local.ts

@@ -18,12 +18,13 @@ import {
 import { isEqual } from 'lodash-es';
 import { nanoid } from 'nanoid';
 import { Observable } from 'rxjs';
-import { applyUpdate, encodeStateAsUpdate } from 'yjs';
+import { encodeStateAsUpdate } from 'yjs';
 
 import { DesktopApiService } from '../../desktop-api';
 import type { WorkspaceEngineStorageProvider } from '../providers/engine';
 import { BroadcastChannelAwarenessConnection } from './engine/awareness-broadcast-channel';
 import { StaticBlobStorage } from './engine/blob-static';
+import { getWorkspaceProfileWorker } from './out-worker';
 
 export const LOCAL_WORKSPACE_LOCAL_STORAGE_KEY = 'affine-local-workspace';
 const LOCAL_WORKSPACE_CHANGED_BROADCAST_CHANNEL_KEY =
@@ -97,21 +98,25 @@ class LocalWorkspaceFlavourProvider implements WorkspaceFlavourProvider {
       blobSources: { main: blobStorage },
     });
 
-    // apply initial state
-    await initial(docCollection, blobStorage, docStorage);
+    try {
+      // apply initial state
+      await initial(docCollection, blobStorage, docStorage);
 
-    // save workspace to local storage, should be vary fast
-    await docStorage.doc.set(id, encodeStateAsUpdate(docCollection.doc));
-    for (const subdocs of docCollection.doc.getSubdocs()) {
-      await docStorage.doc.set(subdocs.guid, encodeStateAsUpdate(subdocs));
+      // save workspace to local storage, should be vary fast
+      await docStorage.doc.set(id, encodeStateAsUpdate(docCollection.doc));
+      for (const subdocs of docCollection.doc.getSubdocs()) {
+        await docStorage.doc.set(subdocs.guid, encodeStateAsUpdate(subdocs));
+      }
+
+      // save workspace id to local storage
+      setLocalWorkspaceIds(ids => [...ids, id]);
+
+      // notify all browser tabs, so they can update their workspace list
+      this.notifyChannel.postMessage(id);
+    } finally {
+      docCollection.dispose();
     }
 
-    // save workspace id to local storage
-    setLocalWorkspaceIds(ids => [...ids, id]);
-
-    // notify all browser tabs, so they can update their workspace list
-    this.notifyChannel.postMessage(id);
-
     return { id, flavour: 'local' };
   }
   workspaces$ = LiveData.from(
@@ -158,16 +163,16 @@ class LocalWorkspaceFlavourProvider implements WorkspaceFlavourProvider {
       };
     }
 
-    const bs = new DocCollection({
-      id,
-      schema: getAFFiNEWorkspaceSchema(),
-    });
+    const client = getWorkspaceProfileWorker();
 
-    if (localData) applyUpdate(bs.doc, localData);
+    const result = await client.call(
+      'renderWorkspaceProfile',
+      [localData].filter(Boolean) as Uint8Array[]
+    );
 
     return {
-      name: bs.meta.name,
-      avatar: bs.meta.avatar,
+      name: result.name,
+      avatar: result.avatar,
       isOwner: true,
     };
   }

packages/frontend/core/src/modules/workspace-engine/impls/out-worker.ts

@@ -0,0 +1,22 @@
+import { OpClient } from '@toeverything/infra/op';
+
+import type { WorkerOps } from './worker-ops';
+
+let worker: OpClient<WorkerOps> | undefined;
+
+export function getWorkspaceProfileWorker() {
+  if (worker) {
+    return worker;
+  }
+
+  const rawWorker = new Worker(
+    new URL(
+      /* webpackChunkName: "workspace-profile-worker" */ './in-worker.ts',
+      import.meta.url
+    )
+  );
+
+  worker = new OpClient<WorkerOps>(rawWorker);
+  worker.listen();
+  return worker;
+}

packages/frontend/core/src/modules/workspace-engine/impls/worker-ops.ts

@@ -0,0 +1,5 @@
+import type { OpSchema } from '@toeverything/infra/op';
+
+export interface WorkerOps extends OpSchema {
+  renderWorkspaceProfile: [Uint8Array[], { name?: string; avatar?: string }];
+}