feat(nbstore): share worker between workspaces (#9947)

packages/frontend/apps/web/src/app.tsx

@@ -12,7 +12,7 @@ import { PopupWindowProvider } from '@affine/core/modules/url';
 import { configureBrowserWorkbenchModule } from '@affine/core/modules/workbench';
 import { configureBrowserWorkspaceFlavours } from '@affine/core/modules/workspace-engine';
 import createEmotionCache from '@affine/core/utils/create-emotion-cache';
-import { WorkerClient } from '@affine/nbstore/worker/client';
+import { StoreManagerClient } from '@affine/nbstore/worker/client';
 import { CacheProvider } from '@emotion/react';
 import { Framework, FrameworkRoot, getCurrentStore } from '@toeverything/infra';
 import { OpClient } from '@toeverything/infra/op';
@@ -21,6 +21,24 @@ import { RouterProvider } from 'react-router-dom';
 
 const cache = createEmotionCache();
 
+let storeManagerClient: StoreManagerClient;
+
+if (window.SharedWorker) {
+  const worker = new SharedWorker(
+    new URL(/* webpackChunkName: "nbstore" */ './nbstore.ts', import.meta.url),
+    { name: 'affine-shared-worker' }
+  );
+  storeManagerClient = new StoreManagerClient(new OpClient(worker.port));
+} else {
+  const worker = new Worker(
+    new URL(/* webpackChunkName: "nbstore" */ './nbstore.ts', import.meta.url)
+  );
+  storeManagerClient = new StoreManagerClient(new OpClient(worker));
+}
+window.addEventListener('beforeunload', () => {
+  storeManagerClient.dispose();
+});
+
 const future = {
   v7_startTransition: true,
 } as const;
@@ -32,37 +50,7 @@ configureLocalStorageStateStorageImpls(framework);
 configureBrowserWorkspaceFlavours(framework);
 framework.impl(NbstoreProvider, {
   openStore(key, options) {
-    if (window.SharedWorker) {
-      const worker = new SharedWorker(
-        new URL(
-          /* webpackChunkName: "nbstore" */ './nbstore.ts',
-          import.meta.url
-        ),
-        { name: key }
-      );
-      const client = new WorkerClient(new OpClient(worker.port), options);
-      return {
-        store: client,
-        dispose: () => {
-          worker.port.postMessage({ type: '__close__' });
-          worker.port.close();
-        },
-      };
-    } else {
-      const worker = new Worker(
-        new URL(
-          /* webpackChunkName: "nbstore" */ './nbstore.ts',
-          import.meta.url
-        )
-      );
-      const client = new WorkerClient(new OpClient(worker), options);
-      return {
-        store: client,
-        dispose: () => {
-          worker.terminate();
-        },
-      };
-    }
+    return storeManagerClient.open(key, options);
   },
 });
 framework.impl(PopupWindowProvider, {

packages/frontend/apps/web/src/nbstore.ts

@@ -5,12 +5,12 @@ import { cloudStorages } from '@affine/nbstore/cloud';
 import { idbStorages } from '@affine/nbstore/idb';
 import { idbV1Storages } from '@affine/nbstore/idb/v1';
 import {
-  WorkerConsumer,
-  type WorkerOps,
+  StoreManagerConsumer,
+  type WorkerManagerOps,
 } from '@affine/nbstore/worker/consumer';
 import { type MessageCommunicapable, OpConsumer } from '@toeverything/infra/op';
 
-const consumer = new WorkerConsumer([
+const consumer = new StoreManagerConsumer([
   ...idbStorages,
   ...idbV1Storages,
   ...broadcastChannelStorages,
@@ -19,28 +19,14 @@ const consumer = new WorkerConsumer([
 
 if ('onconnect' in globalThis) {
   // if in shared worker
-  let activeConnectionCount = 0;
 
   (globalThis as any).onconnect = (event: MessageEvent) => {
-    activeConnectionCount++;
     const port = event.ports[0];
-    port.addEventListener('message', (event: MessageEvent) => {
-      if (event.data.type === '__close__') {
-        activeConnectionCount--;
-        if (activeConnectionCount === 0) {
-          globalThis.close();
-        }
-      }
-    });
-
-    const opConsumer = new OpConsumer<WorkerOps>(port);
-    consumer.bindConsumer(opConsumer);
+    consumer.bindConsumer(new OpConsumer<WorkerManagerOps>(port));
   };
 } else {
   // if in worker
-  const opConsumer = new OpConsumer<WorkerOps>(
-    globalThis as MessageCommunicapable
+  consumer.bindConsumer(
+    new OpConsumer<WorkerManagerOps>(globalThis as MessageCommunicapable)
   );
-
-  consumer.bindConsumer(opConsumer);
 }