feat(nbstore): add nbstore worker (#9185)

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

@@ -0,0 +1,294 @@
+import type { OpClient } from '@toeverything/infra/op';
+
+import { DummyConnection } from '../connection';
+import { DocFrontend } from '../frontend/doc';
+import {
+  type AwarenessRecord,
+  type AwarenessStorage,
+  type BlobRecord,
+  type BlobStorage,
+  type DocRecord,
+  type DocStorage,
+  type DocUpdate,
+  type ListedBlobRecord,
+  type StorageOptions,
+  universalId,
+} from '../storage';
+import type { AwarenessSync } from '../sync/awareness';
+import type { BlobSync } from '../sync/blob';
+import type { DocSync } from '../sync/doc';
+import type { WorkerOps } from './ops';
+
+export class WorkerClient {
+  constructor(
+    private readonly client: OpClient<WorkerOps>,
+    private readonly options: StorageOptions
+  ) {}
+
+  readonly docStorage = new WorkerDocStorage(this.client, this.options);
+  readonly blobStorage = new WorkerBlobStorage(this.client, this.options);
+  readonly awarenessStorage = new WorkerAwarenessStorage(
+    this.client,
+    this.options
+  );
+  readonly docSync = new WorkerDocSync(this.client);
+  readonly blobSync = new WorkerBlobSync(this.client);
+  readonly awarenessSync = new WorkerAwarenessSync(this.client);
+
+  readonly docFrontend = new DocFrontend(this.docStorage, this.docSync);
+}
+
+class WorkerDocStorage implements DocStorage {
+  constructor(
+    private readonly client: OpClient<WorkerOps>,
+    private readonly options: StorageOptions
+  ) {}
+
+  readonly peer = this.options.peer;
+  readonly spaceType = this.options.type;
+  readonly spaceId = this.options.id;
+  readonly universalId = universalId(this.options);
+  readonly storageType = 'doc';
+
+  async getDoc(docId: string) {
+    return this.client.call('docStorage.getDoc', docId);
+  }
+
+  async getDocDiff(docId: string, state?: Uint8Array) {
+    return this.client.call('docStorage.getDocDiff', { docId, state });
+  }
+
+  async pushDocUpdate(update: DocUpdate, origin?: string) {
+    return this.client.call('docStorage.pushDocUpdate', { update, origin });
+  }
+
+  async getDocTimestamp(docId: string) {
+    return this.client.call('docStorage.getDocTimestamp', docId);
+  }
+
+  async getDocTimestamps(after?: Date) {
+    return this.client.call('docStorage.getDocTimestamps', after ?? null);
+  }
+
+  async deleteDoc(docId: string) {
+    return this.client.call('docStorage.deleteDoc', docId);
+  }
+
+  subscribeDocUpdate(callback: (update: DocRecord, origin?: string) => void) {
+    const subscription = this.client
+      .ob$('docStorage.subscribeDocUpdate')
+      .subscribe(value => {
+        callback(value.update, value.origin);
+      });
+    return () => {
+      subscription.unsubscribe();
+    };
+  }
+
+  connection = new WorkerDocConnection(this.client);
+}
+
+class WorkerDocConnection extends DummyConnection {
+  constructor(private readonly client: OpClient<WorkerOps>) {
+    super();
+  }
+
+  override waitForConnected(signal?: AbortSignal): Promise<void> {
+    return new Promise((resolve, reject) => {
+      const abortListener = () => {
+        reject(signal?.reason);
+        subscription.unsubscribe();
+      };
+
+      signal?.addEventListener('abort', abortListener);
+
+      const subscription = this.client
+        .ob$('docStorage.waitForConnected')
+        .subscribe({
+          next() {
+            signal?.removeEventListener('abort', abortListener);
+            resolve();
+          },
+          error(err) {
+            signal?.removeEventListener('abort', abortListener);
+            reject(err);
+          },
+        });
+    });
+  }
+}
+
+class WorkerBlobStorage implements BlobStorage {
+  constructor(
+    private readonly client: OpClient<WorkerOps>,
+    private readonly options: StorageOptions
+  ) {}
+
+  readonly storageType = 'blob';
+  readonly peer = this.options.peer;
+  readonly spaceType = this.options.type;
+  readonly spaceId = this.options.id;
+  readonly universalId = universalId(this.options);
+
+  get(key: string, _signal?: AbortSignal): Promise<BlobRecord | null> {
+    return this.client.call('blobStorage.getBlob', key);
+  }
+  set(blob: BlobRecord, _signal?: AbortSignal): Promise<void> {
+    return this.client.call('blobStorage.setBlob', blob);
+  }
+
+  delete(
+    key: string,
+    permanently: boolean,
+    _signal?: AbortSignal
+  ): Promise<void> {
+    return this.client.call('blobStorage.deleteBlob', { key, permanently });
+  }
+
+  release(_signal?: AbortSignal): Promise<void> {
+    return this.client.call('blobStorage.releaseBlobs');
+  }
+
+  list(_signal?: AbortSignal): Promise<ListedBlobRecord[]> {
+    return this.client.call('blobStorage.listBlobs');
+  }
+
+  connection = new DummyConnection();
+}
+
+class WorkerAwarenessStorage implements AwarenessStorage {
+  constructor(
+    private readonly client: OpClient<WorkerOps>,
+    private readonly options: StorageOptions
+  ) {}
+
+  readonly storageType = 'awareness';
+  readonly peer = this.options.peer;
+  readonly spaceType = this.options.type;
+  readonly spaceId = this.options.id;
+  readonly universalId = universalId(this.options);
+
+  update(record: AwarenessRecord, origin?: string): Promise<void> {
+    return this.client.call('awarenessStorage.update', {
+      awareness: record,
+      origin,
+    });
+  }
+  subscribeUpdate(
+    id: string,
+    onUpdate: (update: AwarenessRecord, origin?: string) => void,
+    onCollect: () => Promise<AwarenessRecord | null>
+  ): () => void {
+    const subscription = this.client
+      .ob$('awarenessStorage.subscribeUpdate', id)
+      .subscribe({
+        next: update => {
+          if (update.type === 'awareness-update') {
+            onUpdate(update.awareness, update.origin);
+          }
+          if (update.type === 'awareness-collect') {
+            onCollect()
+              .then(record => {
+                if (record) {
+                  this.client
+                    .call('awarenessStorage.collect', {
+                      awareness: record,
+                      collectId: update.collectId,
+                    })
+                    .catch(err => {
+                      console.error('error feedback collected awareness', err);
+                    });
+                }
+              })
+              .catch(err => {
+                console.error('error collecting awareness', err);
+              });
+          }
+        },
+      });
+    return () => {
+      subscription.unsubscribe();
+    };
+  }
+  connection = new DummyConnection();
+}
+
+class WorkerDocSync implements DocSync {
+  constructor(private readonly client: OpClient<WorkerOps>) {}
+
+  readonly state$ = this.client.ob$('docSync.state');
+
+  docState$(docId: string) {
+    return this.client.ob$('docSync.docState', docId);
+  }
+
+  addPriority(docId: string, priority: number) {
+    const subscription = this.client
+      .ob$('docSync.addPriority', { docId, priority })
+      .subscribe();
+    return () => {
+      subscription.unsubscribe();
+    };
+  }
+}
+
+class WorkerBlobSync implements BlobSync {
+  constructor(private readonly client: OpClient<WorkerOps>) {}
+  downloadBlob(
+    blobId: string,
+    _signal?: AbortSignal
+  ): Promise<BlobRecord | null> {
+    return this.client.call('blobSync.downloadBlob', blobId);
+  }
+  uploadBlob(blob: BlobRecord, _signal?: AbortSignal): Promise<void> {
+    return this.client.call('blobSync.uploadBlob', blob);
+  }
+}
+
+class WorkerAwarenessSync implements AwarenessSync {
+  constructor(private readonly client: OpClient<WorkerOps>) {}
+
+  update(record: AwarenessRecord, origin?: string): Promise<void> {
+    return this.client.call('awarenessSync.update', {
+      awareness: record,
+      origin,
+    });
+  }
+
+  subscribeUpdate(
+    id: string,
+    onUpdate: (update: AwarenessRecord, origin?: string) => void,
+    onCollect: () => Promise<AwarenessRecord | null>
+  ): () => void {
+    const subscription = this.client
+      .ob$('awarenessSync.subscribeUpdate', id)
+      .subscribe({
+        next: update => {
+          if (update.type === 'awareness-update') {
+            onUpdate(update.awareness, update.origin);
+          }
+          if (update.type === 'awareness-collect') {
+            onCollect()
+              .then(record => {
+                if (record) {
+                  this.client
+                    .call('awarenessSync.collect', {
+                      awareness: record,
+                      collectId: update.collectId,
+                    })
+                    .catch(err => {
+                      console.error('error feedback collected awareness', err);
+                    });
+                }
+              })
+              .catch(err => {
+                console.error('error collecting awareness', err);
+              });
+          }
+        },
+      });
+    return () => {
+      subscription.unsubscribe();
+    };
+  }
+}

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

@@ -0,0 +1,256 @@
+import type { OpConsumer } from '@toeverything/infra/op';
+import { Observable } from 'rxjs';
+
+import { getAvailableStorageImplementations } from '../impls';
+import { SpaceStorage, type StorageOptions } from '../storage';
+import type { AwarenessRecord } from '../storage/awareness';
+import { Sync } from '../sync';
+import type { WorkerOps } from './ops';
+
+export class WorkerConsumer {
+  private remotes: SpaceStorage[] = [];
+  private local: SpaceStorage | null = null;
+  private sync: Sync | null = null;
+
+  get ensureLocal() {
+    if (!this.local) {
+      throw new Error('Not initialized');
+    }
+    return this.local;
+  }
+
+  get ensureSync() {
+    if (!this.sync) {
+      throw new Error('Not initialized');
+    }
+    return this.sync;
+  }
+
+  get docStorage() {
+    return this.ensureLocal.get('doc');
+  }
+
+  get docSync() {
+    const docSync = this.ensureSync.doc;
+    if (!docSync) {
+      throw new Error('Doc sync not initialized');
+    }
+    return docSync;
+  }
+
+  get blobStorage() {
+    return this.ensureLocal.get('blob');
+  }
+
+  get blobSync() {
+    const blobSync = this.ensureSync.blob;
+    if (!blobSync) {
+      throw new Error('Blob sync not initialized');
+    }
+    return blobSync;
+  }
+
+  get syncStorage() {
+    return this.ensureLocal.get('sync');
+  }
+
+  get awarenessStorage() {
+    return this.ensureLocal.get('awareness');
+  }
+
+  get awarenessSync() {
+    const awarenessSync = this.ensureSync.awareness;
+    if (!awarenessSync) {
+      throw new Error('Awareness sync not initialized');
+    }
+    return awarenessSync;
+  }
+
+  constructor(private readonly consumer: OpConsumer<WorkerOps>) {}
+
+  listen() {
+    this.registerHandlers();
+    this.consumer.listen();
+  }
+
+  async init(init: {
+    local: { name: string; opts: StorageOptions }[];
+    remotes: { name: string; opts: StorageOptions }[][];
+  }) {
+    this.local = new SpaceStorage(
+      init.local.map(opt => {
+        const Storage = getAvailableStorageImplementations(opt.name);
+        return new Storage(opt.opts);
+      })
+    );
+    this.remotes = init.remotes.map(opts => {
+      return new SpaceStorage(
+        opts.map(opt => {
+          const Storage = getAvailableStorageImplementations(opt.name);
+          return new Storage(opt.opts);
+        })
+      );
+    });
+    this.sync = new Sync(this.local, this.remotes);
+    this.local.connect();
+    for (const remote of this.remotes) {
+      remote.connect();
+    }
+    this.sync.start();
+  }
+
+  async destroy() {
+    this.sync?.stop();
+    this.local?.disconnect();
+    await this.local?.destroy();
+    for (const remote of this.remotes) {
+      remote.disconnect();
+      await remote.destroy();
+    }
+  }
+
+  private registerHandlers() {
+    const collectJobs = new Map<
+      string,
+      (awareness: AwarenessRecord | null) => void
+    >();
+    let collectId = 0;
+    this.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),
+      'docStorage.pushDocUpdate': ({ update, origin }) =>
+        this.docStorage.pushDocUpdate(update, origin),
+      'docStorage.getDocTimestamps': after =>
+        this.docStorage.getDocTimestamps(after ?? undefined),
+      'docStorage.getDocTimestamp': docId =>
+        this.docStorage.getDocTimestamp(docId),
+      'docStorage.deleteDoc': (docId: string) =>
+        this.docStorage.deleteDoc(docId),
+      'docStorage.subscribeDocUpdate': () =>
+        new Observable(subscriber => {
+          return this.docStorage.subscribeDocUpdate((update, origin) => {
+            subscriber.next({ update, origin });
+          });
+        }),
+      'docStorage.waitForConnected': () =>
+        new Observable(subscriber => {
+          const abortController = new AbortController();
+          this.docStorage.connection
+            .waitForConnected(abortController.signal)
+            .then(() => {
+              subscriber.next(true);
+              subscriber.complete();
+            })
+            .catch(error => {
+              subscriber.error(error);
+            });
+          return () => abortController.abort();
+        }),
+      'blobStorage.getBlob': key => this.blobStorage.get(key),
+      'blobStorage.setBlob': blob => this.blobStorage.set(blob),
+      'blobStorage.deleteBlob': ({ key, permanently }) =>
+        this.blobStorage.delete(key, permanently),
+      'blobStorage.releaseBlobs': () => this.blobStorage.release(),
+      'blobStorage.listBlobs': () => this.blobStorage.list(),
+      'syncStorage.clearClocks': () => this.syncStorage.clearClocks(),
+      'syncStorage.getPeerPulledRemoteClock': ({ peer, docId }) =>
+        this.syncStorage.getPeerPulledRemoteClock(peer, docId),
+      'syncStorage.getPeerPulledRemoteClocks': ({ peer }) =>
+        this.syncStorage.getPeerPulledRemoteClocks(peer),
+      'syncStorage.setPeerPulledRemoteClock': ({ peer, clock }) =>
+        this.syncStorage.setPeerPulledRemoteClock(peer, clock),
+      'syncStorage.getPeerRemoteClock': ({ peer, docId }) =>
+        this.syncStorage.getPeerRemoteClock(peer, docId),
+      'syncStorage.getPeerRemoteClocks': ({ peer }) =>
+        this.syncStorage.getPeerRemoteClocks(peer),
+      'syncStorage.setPeerRemoteClock': ({ peer, clock }) =>
+        this.syncStorage.setPeerRemoteClock(peer, clock),
+      'syncStorage.getPeerPushedClock': ({ peer, docId }) =>
+        this.syncStorage.getPeerPushedClock(peer, docId),
+      'syncStorage.getPeerPushedClocks': ({ peer }) =>
+        this.syncStorage.getPeerPushedClocks(peer),
+      'syncStorage.setPeerPushedClock': ({ peer, clock }) =>
+        this.syncStorage.setPeerPushedClock(peer, clock),
+      'awarenessStorage.update': ({ awareness, origin }) =>
+        this.awarenessStorage.update(awareness, origin),
+      'awarenessStorage.subscribeUpdate': docId =>
+        new Observable(subscriber => {
+          return this.awarenessStorage.subscribeUpdate(
+            docId,
+            (update, origin) => {
+              subscriber.next({
+                type: 'awareness-update',
+                awareness: update,
+                origin,
+              });
+            },
+            () => {
+              const currentCollectId = collectId++;
+              const promise = new Promise<AwarenessRecord | null>(resolve => {
+                collectJobs.set(currentCollectId.toString(), awareness => {
+                  resolve(awareness);
+                  collectJobs.delete(currentCollectId.toString());
+                });
+              });
+              return promise;
+            }
+          );
+        }),
+      'awarenessStorage.collect': ({ collectId, awareness }) =>
+        collectJobs.get(collectId)?.(awareness),
+      'docSync.state': () =>
+        new Observable(subscriber => {
+          const subscription = this.docSync.state$.subscribe(state => {
+            subscriber.next(state);
+          });
+          return () => subscription.unsubscribe();
+        }),
+      'docSync.docState': docId =>
+        new Observable(subscriber => {
+          const subscription = this.docSync
+            .docState$(docId)
+            .subscribe(state => {
+              subscriber.next(state);
+            });
+          return () => subscription.unsubscribe();
+        }),
+      'docSync.addPriority': ({ docId, priority }) =>
+        new Observable(() => {
+          const undo = this.docSync.addPriority(docId, priority);
+          return () => undo();
+        }),
+      'blobSync.downloadBlob': key => this.blobSync.downloadBlob(key),
+      'blobSync.uploadBlob': blob => this.blobSync.uploadBlob(blob),
+      'awarenessSync.update': ({ awareness, origin }) =>
+        this.awarenessSync.update(awareness, origin),
+      'awarenessSync.subscribeUpdate': docId =>
+        new Observable(subscriber => {
+          return this.awarenessStorage.subscribeUpdate(
+            docId,
+            (update, origin) => {
+              subscriber.next({
+                type: 'awareness-update',
+                awareness: update,
+                origin,
+              });
+            },
+            () => {
+              const currentCollectId = collectId++;
+              const promise = new Promise<AwarenessRecord | null>(resolve => {
+                collectJobs.set(currentCollectId.toString(), awareness => {
+                  resolve(awareness);
+                  collectJobs.delete(currentCollectId.toString());
+                });
+              });
+              return promise;
+            }
+          );
+        }),
+      'awarenessSync.collect': ({ collectId, awareness }) =>
+        collectJobs.get(collectId)?.(awareness),
+    });
+  }
+}

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

@@ -0,0 +1,122 @@
+import type {
+  BlobRecord,
+  DocClock,
+  DocClocks,
+  DocDiff,
+  DocRecord,
+  DocUpdate,
+  ListedBlobRecord,
+  StorageOptions,
+} from '../storage';
+import type { AwarenessRecord } from '../storage/awareness';
+import type { DocSyncDocState, DocSyncState } from '../sync/doc';
+
+interface GroupedWorkerOps {
+  worker: {
+    init: [
+      {
+        local: { name: string; opts: StorageOptions }[];
+        remotes: { name: string; opts: StorageOptions }[][];
+      },
+      void,
+    ];
+    destroy: [void, void];
+  };
+
+  docStorage: {
+    getDoc: [string, DocRecord | null];
+    getDocDiff: [{ docId: string; state?: Uint8Array }, DocDiff | null];
+    pushDocUpdate: [{ update: DocUpdate; origin?: string }, DocClock];
+    getDocTimestamps: [Date | null, DocClocks];
+    getDocTimestamp: [string, DocClock | null];
+    deleteDoc: [string, void];
+    subscribeDocUpdate: [void, { update: DocRecord; origin?: string }];
+    waitForConnected: [void, boolean];
+  };
+
+  blobStorage: {
+    getBlob: [string, BlobRecord | null];
+    setBlob: [BlobRecord, void];
+    deleteBlob: [{ key: string; permanently: boolean }, void];
+    releaseBlobs: [void, void];
+    listBlobs: [void, ListedBlobRecord[]];
+  };
+
+  syncStorage: {
+    getPeerPulledRemoteClocks: [{ peer: string }, DocClocks];
+    getPeerPulledRemoteClock: [
+      { peer: string; docId: string },
+      DocClock | null,
+    ];
+    setPeerPulledRemoteClock: [{ peer: string; clock: DocClock }, void];
+    getPeerRemoteClocks: [{ peer: string }, DocClocks];
+    getPeerRemoteClock: [{ peer: string; docId: string }, DocClock | null];
+    setPeerRemoteClock: [{ peer: string; clock: DocClock }, void];
+    getPeerPushedClocks: [{ peer: string }, DocClocks];
+    getPeerPushedClock: [{ peer: string; docId: string }, DocClock | null];
+    setPeerPushedClock: [{ peer: string; clock: DocClock }, void];
+    clearClocks: [void, void];
+  };
+
+  awarenessStorage: {
+    update: [{ awareness: AwarenessRecord; origin?: string }, void];
+    subscribeUpdate: [
+      string,
+      (
+        | {
+            type: 'awareness-update';
+            awareness: AwarenessRecord;
+            origin?: string;
+          }
+        | { type: 'awareness-collect'; collectId: string }
+      ),
+    ];
+    collect: [{ collectId: string; awareness: AwarenessRecord }, void];
+  };
+
+  docSync: {
+    state: [void, DocSyncState];
+    docState: [string, DocSyncDocState];
+    addPriority: [{ docId: string; priority: number }, boolean];
+  };
+
+  blobSync: {
+    downloadBlob: [string, BlobRecord | null];
+    uploadBlob: [BlobRecord, void];
+  };
+
+  awarenessSync: {
+    update: [{ awareness: AwarenessRecord; origin?: string }, void];
+    subscribeUpdate: [
+      string,
+      (
+        | {
+            type: 'awareness-update';
+            awareness: AwarenessRecord;
+            origin?: string;
+          }
+        | { type: 'awareness-collect'; collectId: string }
+      ),
+    ];
+    collect: [{ collectId: string; awareness: AwarenessRecord }, void];
+  };
+}
+
+type Values<T> = T extends { [k in keyof T]: any } ? T[keyof T] : never;
+type UnionToIntersection<U> = (U extends any ? (x: U) => void : never) extends (
+  x: infer I
+) => void
+  ? I
+  : never;
+
+export type WorkerOps = UnionToIntersection<
+  Values<
+    Values<{
+      [k in keyof GroupedWorkerOps]: {
+        [k2 in keyof GroupedWorkerOps[k]]: k2 extends string
+          ? Record<`${k}.${k2}`, GroupedWorkerOps[k][k2]>
+          : never;
+      };
+    }>
+  >
+>;