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

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

@@ -14,7 +14,7 @@ import { PopupWindowProvider } from '@affine/core/modules/url';
 import { ClientSchemeProvider } from '@affine/core/modules/url/providers/client-schema';
 import { configureBrowserWorkbenchModule } from '@affine/core/modules/workbench';
 import { configureBrowserWorkspaceFlavours } from '@affine/core/modules/workspace-engine';
-import { WorkerClient } from '@affine/nbstore/worker/client';
+import { StoreManagerClient } from '@affine/nbstore/worker/client';
 import { App as CapacitorApp } from '@capacitor/app';
 import { InAppBrowser } from '@capgo/inappbrowser';
 import { Framework, FrameworkRoot, getCurrentStore } from '@toeverything/infra';
@@ -22,6 +22,17 @@ import { OpClient } from '@toeverything/infra/op';
 import { Suspense } from 'react';
 import { RouterProvider } from 'react-router-dom';
 
+const storeManagerClient = new StoreManagerClient(
+  new OpClient(
+    new Worker(
+      new URL(/* webpackChunkName: "nbstore" */ './nbstore.ts', import.meta.url)
+    )
+  )
+);
+window.addEventListener('beforeunload', () => {
+  storeManagerClient.dispose();
+});
+
 const future = {
   v7_startTransition: true,
 } as const;
@@ -33,15 +44,12 @@ configureLocalStorageStateStorageImpls(framework);
 configureBrowserWorkspaceFlavours(framework);
 configureMobileModules(framework);
 framework.impl(NbstoreProvider, {
-  openStore(_key, options) {
-    const worker = new Worker(
-      new URL(/* webpackChunkName: "nbstore" */ './nbstore.ts', import.meta.url)
-    );
-    const client = new WorkerClient(new OpClient(worker), options);
+  openStore(key, options) {
+    const { store, dispose } = storeManagerClient.open(key, options);
     return {
-      store: client,
+      store,
       dispose: () => {
-        worker.terminate();
+        dispose();
       },
     };
   },

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

@@ -4,18 +4,18 @@ import { broadcastChannelStorages } from '@affine/nbstore/broadcast-channel';
 import { cloudStorages } from '@affine/nbstore/cloud';
 import { idbStorages } from '@affine/nbstore/idb';
 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,
   ...broadcastChannelStorages,
   ...cloudStorages,
 ]);
 
-const opConsumer = new OpConsumer<WorkerOps>(
+const opConsumer = new OpConsumer<WorkerManagerOps>(
   globalThis as MessageCommunicapable
 );
 

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

packages/frontend/apps/electron/src/main/windows-manager/main-window.ts

@@ -269,12 +269,6 @@ export async function openUrlInHiddenWindow(urlObj: URL) {
     win.webContents.openDevTools();
   }
 
-  win.on('close', e => {
-    e.preventDefault();
-    if (win && !win.isDestroyed()) {
-      win.destroy();
-    }
-  });
   logger.info('loading page at', url);
   win.loadURL(url).catch(e => {
     logger.error('failed to load url', e);

packages/frontend/apps/electron/src/main/windows-manager/tab-views.ts

@@ -395,7 +395,9 @@ export class WebContentViewsManager {
 
       if (this.mainWindow && view) {
         this.mainWindow.contentView.removeChildView(view);
-        view?.webContents.close();
+        view?.webContents.close({
+          waitForBeforeUnload: true,
+        });
       }
     }, 500); // delay a bit to get rid of the flicker
   };

packages/frontend/apps/electron/src/main/worker/pool.ts

@@ -43,13 +43,9 @@ export class WorkerManager {
         show: false,
       });
       let disconnectHelperProcess: (() => void) | null = null;
-      worker.on('close', e => {
-        e.preventDefault();
-        if (worker && !worker.isDestroyed()) {
-          worker.destroy();
-          this.workers.delete(key);
-          disconnectHelperProcess?.();
-        }
+      worker.on('closed', () => {
+        this.workers.delete(key);
+        disconnectHelperProcess?.();
       });
       worker.loadURL(backgroundWorkerViewUrl).catch(e => {
         logger.error('failed to load url', e);
@@ -74,6 +70,7 @@ export class WorkerManager {
       this.disconnectWorker(key, portId);
     });
     const worker = await this.getOrCreateWorker(key);
+    worker.ports.add(portId);
     const { port1: portForWorker, port2: portForRenderer } =
       new MessageChannelMain();
 

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

@@ -27,7 +27,7 @@ import { configureBrowserWorkbenchModule } from '@affine/core/modules/workbench'
 import { WorkspacesService } from '@affine/core/modules/workspace';
 import { configureBrowserWorkspaceFlavours } from '@affine/core/modules/workspace-engine';
 import { I18n } from '@affine/i18n';
-import { WorkerClient } from '@affine/nbstore/worker/client';
+import { StoreManagerClient } from '@affine/nbstore/worker/client';
 import {
   defaultBlockMarkdownAdapterMatchers,
   docLinkBaseURLMiddleware,
@@ -56,6 +56,11 @@ import { Intelligents } from './plugins/intelligents';
 import { NbStoreNativeDBApis } from './plugins/nbstore';
 import { enableNavigationGesture$ } from './web-navigation-control';
 
+const storeManagerClient = createStoreManagerClient();
+window.addEventListener('beforeunload', () => {
+  storeManagerClient.dispose();
+});
+
 const future = {
   v7_startTransition: true,
 } as const;
@@ -67,46 +72,12 @@ configureLocalStorageStateStorageImpls(framework);
 configureBrowserWorkspaceFlavours(framework);
 configureMobileModules(framework);
 framework.impl(NbstoreProvider, {
-  openStore(_key, options) {
-    const worker = new Worker(
-      new URL(
-        /* webpackChunkName: "nbstore-worker" */ './worker.ts',
-        import.meta.url
-      )
-    );
-    const { port1: nativeDBApiChannelServer, port2: nativeDBApiChannelClient } =
-      new MessageChannel();
-    AsyncCall<typeof NbStoreNativeDBApis>(NbStoreNativeDBApis, {
-      channel: {
-        on(listener) {
-          const f = (e: MessageEvent<any>) => {
-            listener(e.data);
-          };
-          nativeDBApiChannelServer.addEventListener('message', f);
-          return () => {
-            nativeDBApiChannelServer.removeEventListener('message', f);
-          };
-        },
-        send(data) {
-          nativeDBApiChannelServer.postMessage(data);
-        },
-      },
-      log: false,
-    });
-    nativeDBApiChannelServer.start();
-    worker.postMessage(
-      {
-        type: 'native-db-api-channel',
-        port: nativeDBApiChannelClient,
-      },
-      [nativeDBApiChannelClient]
-    );
-    const client = new WorkerClient(new OpClient(worker), options);
+  openStore(key, options) {
+    const { store, dispose } = storeManagerClient.open(key, options);
     return {
-      store: client,
+      store,
       dispose: () => {
-        worker.terminate();
-        nativeDBApiChannelServer.close();
+        dispose();
       },
     };
   },
@@ -336,3 +307,40 @@ export function App() {
     </Suspense>
   );
 }
+
+function createStoreManagerClient() {
+  const worker = new Worker(
+    new URL(
+      /* webpackChunkName: "nbstore-worker" */ './worker.ts',
+      import.meta.url
+    )
+  );
+  const { port1: nativeDBApiChannelServer, port2: nativeDBApiChannelClient } =
+    new MessageChannel();
+  AsyncCall<typeof NbStoreNativeDBApis>(NbStoreNativeDBApis, {
+    channel: {
+      on(listener) {
+        const f = (e: MessageEvent<any>) => {
+          listener(e.data);
+        };
+        nativeDBApiChannelServer.addEventListener('message', f);
+        return () => {
+          nativeDBApiChannelServer.removeEventListener('message', f);
+        };
+      },
+      send(data) {
+        nativeDBApiChannelServer.postMessage(data);
+      },
+    },
+    log: false,
+  });
+  nativeDBApiChannelServer.start();
+  worker.postMessage(
+    {
+      type: 'native-db-api-channel',
+      port: nativeDBApiChannelClient,
+    },
+    [nativeDBApiChannelClient]
+  );
+  return new StoreManagerClient(new OpClient(worker));
+}

packages/frontend/apps/ios/src/worker.ts

@@ -8,8 +8,8 @@ import {
   sqliteStorages,
 } from '@affine/nbstore/sqlite';
 import {
-  WorkerConsumer,
-  type WorkerOps,
+  StoreManagerConsumer,
+  type WorkerManagerOps,
 } from '@affine/nbstore/worker/consumer';
 import { type MessageCommunicapable, OpConsumer } from '@toeverything/infra/op';
 import { AsyncCall } from 'async-call-rpc';
@@ -41,12 +41,14 @@ globalThis.addEventListener('message', e => {
   }
 });
 
-const consumer = new OpConsumer<WorkerOps>(globalThis as MessageCommunicapable);
+const consumer = new OpConsumer<WorkerManagerOps>(
+  globalThis as MessageCommunicapable
+);
 
-const worker = new WorkerConsumer([
+const storeManager = new StoreManagerConsumer([
   ...sqliteStorages,
   ...broadcastChannelStorages,
   ...cloudStorages,
 ]);
 
-worker.bindConsumer(consumer);
+storeManager.bindConsumer(consumer);

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

@@ -13,12 +13,30 @@ import {
 import { PopupWindowProvider } from '@affine/core/modules/url';
 import { configureBrowserWorkbenchModule } from '@affine/core/modules/workbench';
 import { configureBrowserWorkspaceFlavours } from '@affine/core/modules/workspace-engine';
-import { WorkerClient } from '@affine/nbstore/worker/client';
+import { StoreManagerClient } from '@affine/nbstore/worker/client';
 import { Framework, FrameworkRoot, getCurrentStore } from '@toeverything/infra';
 import { OpClient } from '@toeverything/infra/op';
 import { Suspense } from 'react';
 import { RouterProvider } from 'react-router-dom';
 
+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;
@@ -31,38 +49,13 @@ configureBrowserWorkspaceFlavours(framework);
 configureMobileModules(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);
-      worker.port.start();
-      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();
-        },
-      };
-    }
+    const { store, dispose } = storeManagerClient.open(key, options);
+    return {
+      store: store,
+      dispose: () => {
+        dispose();
+      },
+    };
   },
 });
 framework.impl(PopupWindowProvider, {

packages/frontend/apps/mobile/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);
 }

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

packages/frontend/core/src/bootstrap/browser.ts

@@ -1,5 +1,4 @@
 // ORDER MATTERS
 import './env';
 import './public-path';
-import './shared-worker';
 import './polyfill/browser';

packages/frontend/core/src/bootstrap/electron.ts

@@ -1,5 +1,4 @@
 // ORDER MATTERS
 import './env';
 import './public-path';
-import './shared-worker';
 import './polyfill/electron';

packages/frontend/core/src/bootstrap/shared-worker.ts

@@ -1,25 +0,0 @@
-/**
- * This is a wrapper for SharedWorker,
- * added the `name` parameter to the `SharedWorker` URL so that
- * multiple `SharedWorkers` can share one script file.
- */
-const rawSharedWorker = globalThis.SharedWorker;
-
-// TODO(@eyhn): remove this when we can use single shared worker for all workspaces
-function PatchedSharedWorker(
-  urlParam: URL | string,
-  options?: string | { name: string }
-) {
-  const url = typeof urlParam === 'string' ? new URL(urlParam) : urlParam;
-  if (options) {
-    url.searchParams.append(
-      typeof options === 'string' ? options : options.name,
-      ''
-    );
-  }
-  return new rawSharedWorker(url, options);
-}
-// if SharedWorker is not supported, do nothing
-if (rawSharedWorker) {
-  globalThis.SharedWorker = PatchedSharedWorker as any;
-}

packages/frontend/core/src/modules/cloud/constant.ts

@@ -33,7 +33,9 @@ export const BUILD_IN_SERVERS: (ServerMetadata & { config: ServerConfig })[] =
       ? [
           {
             id: 'affine-cloud',
-            baseUrl: location.origin,
+            baseUrl: BUILD_CONFIG.isElectron
+              ? 'http://localhost:8080'
+              : location.origin,
             config: {
               serverName: 'Affine Cloud',
               features: [

packages/frontend/core/src/modules/storage/providers/nbstore.ts

@@ -1,5 +1,5 @@
 import type {
-  WorkerClient,
+  StoreClient,
   WorkerInitOptions,
 } from '@affine/nbstore/worker/client';
 import { createIdentifier } from '@toeverything/infra';
@@ -17,7 +17,7 @@ export interface NbstoreProvider {
     key: string,
     options: WorkerInitOptions
   ): {
-    store: WorkerClient;
+    store: StoreClient;
     dispose: () => void;
   };
 }

packages/frontend/core/src/modules/userspace/entities/current-user-db.ts

@@ -1,8 +1,9 @@
 import { Entity, LiveData } from '@toeverything/infra';
-import { finalize, of, switchMap } from 'rxjs';
+import { Observable, of, switchMap } from 'rxjs';
 
 import type { AuthService } from '../../cloud';
 import type { UserspaceService } from '../services/userspace';
+import type { UserDBWithTables } from './user-db';
 
 export class CurrentUserDB extends Entity {
   constructor(
@@ -19,11 +20,12 @@ export class CurrentUserDB extends Entity {
         switchMap(userId => {
           if (userId) {
             const ref = this.userDBService.openDB(userId);
-            return of(ref.obj).pipe(
-              finalize(() => {
+            return new Observable<UserDBWithTables>(subscriber => {
+              subscriber.next(ref.obj);
+              return () => {
                 ref.release();
-              })
-            );
+              };
+            });
           } else {
             return of(null);
           }

packages/frontend/core/src/modules/userspace/entities/user-db-engine.ts

@@ -1,6 +1,6 @@
 import { IndexedDBDocStorage } from '@affine/nbstore/idb';
 import { SqliteDocStorage } from '@affine/nbstore/sqlite';
-import type { WorkerClient } from '@affine/nbstore/worker/client';
+import type { StoreClient } from '@affine/nbstore/worker/client';
 import { Entity } from '@toeverything/infra';
 
 import type { ServerService } from '../../cloud';
@@ -10,7 +10,7 @@ export class UserDBEngine extends Entity<{
   userId: string;
 }> {
   private readonly userId = this.props.userId;
-  readonly client: WorkerClient;
+  readonly client: StoreClient;
 
   DocStorageType =
     BUILD_CONFIG.isElectron || BUILD_CONFIG.isIOS

packages/frontend/core/src/modules/userspace/entities/user-db.ts

@@ -39,6 +39,10 @@ export class UserDB extends Entity<{
       });
     });
   }
+
+  override dispose(): void {
+    this.engine.dispose();
+  }
 }
 
 export type UserDBWithTables = UserDB & {

packages/frontend/core/src/modules/workspace/entities/engine.ts

@@ -1,5 +1,5 @@
 import type {
-  WorkerClient,
+  StoreClient,
   WorkerInitOptions,
 } from '@affine/nbstore/worker/client';
 import { Entity } from '@toeverything/infra';
@@ -12,7 +12,7 @@ export class WorkspaceEngine extends Entity<{
   isSharedMode?: boolean;
   engineWorkerInitOptions: WorkerInitOptions;
 }> {
-  worker?: WorkerClient;
+  client?: StoreClient;
   started = false;
 
   constructor(
@@ -23,24 +23,24 @@ export class WorkspaceEngine extends Entity<{
   }
 
   get doc() {
-    if (!this.worker) {
+    if (!this.client) {
       throw new Error('Engine is not initialized');
     }
-    return this.worker.docFrontend;
+    return this.client.docFrontend;
   }
 
   get blob() {
-    if (!this.worker) {
+    if (!this.client) {
       throw new Error('Engine is not initialized');
     }
-    return this.worker.blobFrontend;
+    return this.client.blobFrontend;
   }
 
   get awareness() {
-    if (!this.worker) {
+    if (!this.client) {
       throw new Error('Engine is not initialized');
     }
-    return this.worker.awarenessFrontend;
+    return this.client.awarenessFrontend;
   }
 
   start() {
@@ -54,7 +54,7 @@ export class WorkspaceEngine extends Entity<{
         `workspace:${this.workspaceService.workspace.flavour}:${this.workspaceService.workspace.id}`,
       this.props.engineWorkerInitOptions
     );
-    this.worker = store;
+    this.client = store;
     this.disposables.push(dispose);
     this.eventBus.emit(WorkspaceEngineBeforeStart, this);