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

packages/common/nbstore/src/worker/client.ts

@@ -1,4 +1,4 @@
-import type { OpClient } from '@toeverything/infra/op';
+import { OpClient, transfer } from '@toeverything/infra/op';
 
 import { DummyConnection } from '../connection';
 import { AwarenessFrontend, BlobFrontend, DocFrontend } from '../frontend';
@@ -14,18 +14,68 @@ import {
 import type { AwarenessSync } from '../sync/awareness';
 import type { BlobSync } from '../sync/blob';
 import type { DocSync } from '../sync/doc';
-import type { WorkerInitOptions, WorkerOps } from './ops';
+import type { StoreInitOptions, WorkerManagerOps, WorkerOps } from './ops';
 
-export type { WorkerInitOptions } from './ops';
+export type { StoreInitOptions as WorkerInitOptions } from './ops';
 
-export class WorkerClient {
-  constructor(
-    private readonly client: OpClient<WorkerOps>,
-    options: WorkerInitOptions
-  ) {
-    this.client.call('worker.init', options).catch(err => {
-      console.error('error initializing worker', err);
+export class StoreManagerClient {
+  private readonly connections = new Map<
+    string,
+    {
+      store: StoreClient;
+      dispose: () => void;
+    }
+  >();
+
+  constructor(private readonly client: OpClient<WorkerManagerOps>) {}
+
+  open(key: string, options: StoreInitOptions) {
+    const { port1, port2 } = new MessageChannel();
+
+    const client = new OpClient<WorkerOps>(port1);
+    const closeKey = crypto.randomUUID();
+
+    this.client
+      .call(
+        'open',
+        transfer(
+          {
+            key,
+            closeKey,
+            options,
+            port: port2,
+          },
+          [port2]
+        )
+      )
+      .catch(err => {
+        console.error('error opening', err);
+      });
+
+    const connection = {
+      store: new StoreClient(client),
+      dispose: () => {
+        this.client.call('close', closeKey).catch(err => {
+          console.error('error closing', err);
+        });
+        this.connections.delete(closeKey);
+      },
+    };
+
+    this.connections.set(closeKey, connection);
+
+    return connection;
+  }
+
+  dispose() {
+    this.connections.forEach(connection => {
+      connection.dispose();
     });
+  }
+}
+
+export class StoreClient {
+  constructor(private readonly client: OpClient<WorkerOps>) {
     this.docStorage = new WorkerDocStorage(this.client);
     this.blobStorage = new WorkerBlobStorage(this.client);
     this.docSync = new WorkerDocSync(this.client);

packages/common/nbstore/src/worker/consumer.ts

@@ -1,5 +1,5 @@
 import { MANUALLY_STOP } from '@toeverything/infra';
-import type { OpConsumer } from '@toeverything/infra/op';
+import { OpConsumer } from '@toeverything/infra/op';
 import { Observable } from 'rxjs';
 
 import { type StorageConstructor } from '../impls';
@@ -7,14 +7,13 @@ import { SpaceStorage } from '../storage';
 import type { AwarenessRecord } from '../storage/awareness';
 import { Sync } from '../sync';
 import type { PeerStorageOptions } from '../sync/types';
-import type { WorkerInitOptions, WorkerOps } from './ops';
+import type { StoreInitOptions, WorkerManagerOps, WorkerOps } from './ops';
 
-export type { WorkerOps };
+export type { WorkerManagerOps };
 
-export class WorkerConsumer {
-  private inited = false;
-  private storages: PeerStorageOptions<SpaceStorage> | null = null;
-  private sync: Sync | null = null;
+class StoreConsumer {
+  private readonly storages: PeerStorageOptions<SpaceStorage>;
+  private readonly sync: Sync;
 
   get ensureLocal() {
     if (!this.storages) {
@@ -59,18 +58,9 @@ export class WorkerConsumer {
   }
 
   constructor(
-    private readonly availableStorageImplementations: StorageConstructor[]
-  ) {}
-
-  bindConsumer(consumer: OpConsumer<WorkerOps>) {
-    this.registerHandlers(consumer);
-  }
-
-  init(init: WorkerInitOptions) {
-    if (this.inited) {
-      return;
-    }
-    this.inited = true;
+    private readonly availableStorageImplementations: StorageConstructor[],
+    init: StoreInitOptions
+  ) {
     this.storages = {
       local: new SpaceStorage(
         Object.fromEntries(
@@ -122,6 +112,10 @@ export class WorkerConsumer {
     this.sync.start();
   }
 
+  bindConsumer(consumer: OpConsumer<WorkerOps>) {
+    this.registerHandlers(consumer);
+  }
+
   async destroy() {
     this.sync?.stop();
     this.storages?.local.disconnect();
@@ -139,8 +133,6 @@ export class WorkerConsumer {
     >();
     let collectId = 0;
     consumer.registerAll({
-      'worker.init': this.init.bind(this),
-      'worker.destroy': this.destroy.bind(this),
       'docStorage.getDoc': (docId: string) => this.docStorage.getDoc(docId),
       'docStorage.getDocDiff': ({ docId, state }) =>
         this.docStorage.getDocDiff(docId, state),
@@ -298,3 +290,61 @@ export class WorkerConsumer {
     });
   }
 }
+
+export class StoreManagerConsumer {
+  private readonly storeDisposers = new Map<string, () => void>();
+  private readonly storePool = new Map<
+    string,
+    { store: StoreConsumer; refCount: number }
+  >();
+
+  constructor(
+    private readonly availableStorageImplementations: StorageConstructor[]
+  ) {}
+
+  bindConsumer(consumer: OpConsumer<WorkerManagerOps>) {
+    this.registerHandlers(consumer);
+  }
+
+  private registerHandlers(consumer: OpConsumer<WorkerManagerOps>) {
+    consumer.registerAll({
+      open: ({ port, key, closeKey, options }) => {
+        console.debug('open store', key, closeKey);
+        let storeRef = this.storePool.get(key);
+
+        if (!storeRef) {
+          const store = new StoreConsumer(
+            this.availableStorageImplementations,
+            options
+          );
+          storeRef = { store, refCount: 0 };
+        }
+        storeRef.refCount++;
+
+        const workerConsumer = new OpConsumer<WorkerOps>(port);
+        storeRef.store.bindConsumer(workerConsumer);
+
+        this.storeDisposers.set(closeKey, () => {
+          storeRef.refCount--;
+          if (storeRef.refCount === 0) {
+            storeRef.store.destroy().catch(error => {
+              console.error(error);
+            });
+            this.storePool.delete(key);
+          }
+        });
+        this.storePool.set(key, storeRef);
+        return closeKey;
+      },
+      close: key => {
+        console.debug('close store', key);
+        const workerDisposer = this.storeDisposers.get(key);
+        if (!workerDisposer) {
+          throw new Error('Worker not found');
+        }
+        workerDisposer();
+        this.storeDisposers.delete(key);
+      },
+    });
+  }
+}

packages/common/nbstore/src/worker/ops.ts

@@ -20,17 +20,12 @@ type StorageInitOptions = Values<{
   };
 }>;
 
-export interface WorkerInitOptions {
+export interface StoreInitOptions {
   local: { [key in StorageType]?: StorageInitOptions };
   remotes: Record<string, { [key in StorageType]?: StorageInitOptions }>;
 }
 
 interface GroupedWorkerOps {
-  worker: {
-    init: [WorkerInitOptions, void];
-    destroy: [void, void];
-  };
-
   docStorage: {
     getDoc: [string, DocRecord | null];
     getDocDiff: [{ docId: string; state?: Uint8Array }, DocDiff | null];
@@ -132,3 +127,16 @@ export type WorkerOps = UnionToIntersection<
     }>
   >
 >;
+
+export type WorkerManagerOps = {
+  open: [
+    {
+      port: MessagePort;
+      key: string;
+      closeKey: string;
+      options: StoreInitOptions;
+    },
+    string,
+  ];
+  close: [string, void];
+};