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

packages/frontend/apps/electron-renderer/src/app.tsx

@@ -36,7 +36,7 @@ import { WorkspacesService } from '@affine/core/modules/workspace';
 import { configureBrowserWorkspaceFlavours } from '@affine/core/modules/workspace-engine';
 import createEmotionCache from '@affine/core/utils/create-emotion-cache';
 import { apis, events } from '@affine/electron-api';
-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';
@@ -45,6 +45,11 @@ import { RouterProvider } from 'react-router-dom';
 
 import { DesktopThemeSync } from './theme-sync';
 
+const storeManagerClient = createStoreManagerClient();
+window.addEventListener('beforeunload', () => {
+  storeManagerClient.dispose();
+});
+
 const desktopWhiteList = [
   '/open-app/signin-redirect',
   '/open-app/url',
@@ -81,50 +86,12 @@ configureSpellCheckSettingModule(framework);
 configureDesktopBackupModule(framework);
 framework.impl(NbstoreProvider, {
   openStore(key, options) {
-    const { port1: portForOpClient, port2: portForWorker } =
-      new MessageChannel();
-    let portFromWorker: MessagePort | null = null;
-    let portId = crypto.randomUUID();
+    const { store, dispose } = storeManagerClient.open(key, options);
 
-    const handleMessage = (ev: MessageEvent) => {
-      if (
-        ev.data.type === 'electron:worker-connect' &&
-        ev.data.portId === portId
-      ) {
-        portFromWorker = ev.ports[0];
-        // connect portForWorker and portFromWorker
-        portFromWorker.addEventListener('message', ev => {
-          portForWorker.postMessage(ev.data);
-        });
-        portForWorker.addEventListener('message', ev => {
-          // oxlint-disable-next-line no-non-null-assertion
-          portFromWorker!.postMessage(ev.data);
-        });
-        portForWorker.start();
-        portFromWorker.start();
-      }
-    };
-
-    window.addEventListener('message', handleMessage);
-
-    // oxlint-disable-next-line no-non-null-assertion
-    apis!.worker.connectWorker(key, portId).catch(err => {
-      console.error('failed to connect worker', err);
-    });
-
-    const store = new WorkerClient(new OpClient(portForOpClient), options);
-    portForOpClient.start();
     return {
       store,
       dispose: () => {
-        window.removeEventListener('message', handleMessage);
-        portForOpClient.close();
-        portForWorker.close();
-        portFromWorker?.close();
-        // oxlint-disable-next-line no-non-null-assertion
-        apis!.worker.disconnectWorker(key, portId).catch(err => {
-          console.error('failed to disconnect worker', err);
-        });
+        dispose();
       },
     };
   },
@@ -238,3 +205,39 @@ export function App() {
     </Suspense>
   );
 }
+
+function createStoreManagerClient() {
+  const { port1: portForOpClient, port2: portForWorker } = new MessageChannel();
+  let portFromWorker: MessagePort | null = null;
+  let portId = crypto.randomUUID();
+
+  const handleMessage = (ev: MessageEvent) => {
+    if (
+      ev.data.type === 'electron:worker-connect' &&
+      ev.data.portId === portId
+    ) {
+      portFromWorker = ev.ports[0];
+      // connect portForWorker and portFromWorker
+      portFromWorker.addEventListener('message', ev => {
+        portForWorker.postMessage(ev.data, [...ev.ports]);
+      });
+      portForWorker.addEventListener('message', ev => {
+        // oxlint-disable-next-line no-non-null-assertion
+        portFromWorker!.postMessage(ev.data, [...ev.ports]);
+      });
+      portForWorker.start();
+      portFromWorker.start();
+    }
+  };
+
+  window.addEventListener('message', handleMessage);
+
+  // oxlint-disable-next-line no-non-null-assertion
+  apis!.worker.connectWorker('affine-shared-worker', portId).catch(err => {
+    console.error('failed to connect worker', err);
+  });
+
+  const storeManager = new StoreManagerClient(new OpClient(portForOpClient));
+  portForOpClient.start();
+  return storeManager;
+}

packages/frontend/apps/electron-renderer/src/background-worker/index.ts

@@ -9,8 +9,8 @@ import {
   sqliteV1Storages,
 } from '@affine/nbstore/sqlite/v1';
 import {
-  WorkerConsumer,
-  type WorkerOps,
+  StoreManagerConsumer,
+  type WorkerManagerOps,
 } from '@affine/nbstore/worker/consumer';
 import { OpConsumer } from '@toeverything/infra/op';
 
@@ -19,7 +19,7 @@ bindNativeDBApis(apis!.nbstore);
 // oxlint-disable-next-line no-non-null-assertion
 bindNativeDBV1Apis(apis!.db);
 
-const worker = new WorkerConsumer([
+const storeManager = new StoreManagerConsumer([
   ...sqliteStorages,
   ...sqliteV1Storages,
   ...broadcastChannelStorages,
@@ -30,7 +30,7 @@ window.addEventListener('message', ev => {
   if (ev.data.type === 'electron:worker-connect') {
     const port = ev.ports[0];
 
-    const consumer = new OpConsumer<WorkerOps>(port);
-    worker.bindConsumer(consumer);
+    const consumer = new OpConsumer<WorkerManagerOps>(port);
+    storeManager.bindConsumer(consumer);
   }
 });

packages/frontend/apps/electron-renderer/src/nbstore.ts

@@ -1,96 +0,0 @@
-import '@affine/core/bootstrap/electron';
-
-import type { ClientHandler } from '@affine/electron-api';
-import { broadcastChannelStorages } from '@affine/nbstore/broadcast-channel';
-import { cloudStorages } from '@affine/nbstore/cloud';
-import { bindNativeDBApis, sqliteStorages } from '@affine/nbstore/sqlite';
-import {
-  bindNativeDBV1Apis,
-  sqliteV1Storages,
-} from '@affine/nbstore/sqlite/v1';
-import {
-  WorkerConsumer,
-  type WorkerOps,
-} from '@affine/nbstore/worker/consumer';
-import { OpConsumer } from '@toeverything/infra/op';
-import { AsyncCall } from 'async-call-rpc';
-
-const worker = new WorkerConsumer([
-  ...sqliteStorages,
-  ...sqliteV1Storages,
-  ...broadcastChannelStorages,
-  ...cloudStorages,
-]);
-
-let activeConnectionCount = 0;
-let electronAPIsInitialized = false;
-
-function connectElectronAPIs(port: MessagePort) {
-  if (electronAPIsInitialized) {
-    return;
-  }
-  electronAPIsInitialized = true;
-  port.postMessage({ type: '__electron-apis-init__' });
-
-  const { promise, resolve } = Promise.withResolvers<MessagePort>();
-  port.addEventListener('message', event => {
-    if (event.data.type === '__electron-apis__') {
-      const [port] = event.ports;
-      resolve(port);
-    }
-  });
-
-  const rpc = AsyncCall<Record<string, any>>(null, {
-    channel: promise.then(p => ({
-      on(listener) {
-        p.onmessage = e => {
-          listener(e.data);
-        };
-        p.start();
-        return () => {
-          p.onmessage = null;
-          try {
-            p.close();
-          } catch (err) {
-            console.error('close port error', err);
-          }
-        };
-      },
-      send(data) {
-        p.postMessage(data);
-      },
-    })),
-    log: false,
-  });
-
-  const electronAPIs = new Proxy<ClientHandler>(rpc as any, {
-    get(_, namespace: string) {
-      return new Proxy(rpc as any, {
-        get(_, method: string) {
-          return rpc[`${namespace}:${method}`];
-        },
-      });
-    },
-  });
-
-  bindNativeDBApis(electronAPIs.nbstore);
-  bindNativeDBV1Apis(electronAPIs.db);
-}
-
-(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();
-      }
-    }
-  });
-
-  connectElectronAPIs(port);
-
-  const consumer = new OpConsumer<WorkerOps>(port);
-  worker.bindConsumer(consumer);
-};