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

packages/common/nbstore/src/impls/cloud/awareness.ts

@@ -3,7 +3,7 @@ import type { SocketOptions } from 'socket.io-client';
 import { share } from '../../connection';
 import {
   type AwarenessRecord,
-  AwarenessStorage,
+  AwarenessStorageBase,
   type AwarenessStorageOptions,
 } from '../../storage/awareness';
 import {
@@ -16,7 +16,7 @@ interface CloudAwarenessStorageOptions extends AwarenessStorageOptions {
   socketOptions: SocketOptions;
 }
 
-export class CloudAwarenessStorage extends AwarenessStorage<CloudAwarenessStorageOptions> {
+export class CloudAwarenessStorage extends AwarenessStorageBase<CloudAwarenessStorageOptions> {
   connection = share(
     new SocketConnection(this.peer, this.options.socketOptions)
   );
@@ -38,7 +38,7 @@ export class CloudAwarenessStorage extends AwarenessStorage<CloudAwarenessStorag
   override subscribeUpdate(
     id: string,
     onUpdate: (update: AwarenessRecord, origin?: string) => void,
-    onCollect: () => AwarenessRecord
+    onCollect: () => Promise<AwarenessRecord | null>
   ): () => void {
     // TODO: handle disconnect
     // leave awareness
@@ -92,14 +92,16 @@ export class CloudAwarenessStorage extends AwarenessStorage<CloudAwarenessStorag
         docId === id
       ) {
         (async () => {
-          const record = onCollect();
-          const encodedUpdate = await uint8ArrayToBase64(record.bin);
-          this.socket.emit('space:update-awareness', {
-            spaceType: this.spaceType,
-            spaceId: this.spaceId,
-            docId: record.docId,
-            awarenessUpdate: encodedUpdate,
-          });
+          const record = await onCollect();
+          if (record) {
+            const encodedUpdate = await uint8ArrayToBase64(record.bin);
+            this.socket.emit('space:update-awareness', {
+              spaceType: this.spaceType,
+              spaceId: this.spaceId,
+              docId: record.docId,
+              awarenessUpdate: encodedUpdate,
+            });
+          }
         })().catch(err => console.error('awareness upload failed', err));
       }
     };

packages/common/nbstore/src/impls/cloud/blob.ts

@@ -7,16 +7,31 @@ import {
 } from '@affine/graphql';
 
 import { DummyConnection } from '../../connection';
-import { type BlobRecord, BlobStorage } from '../../storage';
+import {
+  type BlobRecord,
+  BlobStorageBase,
+  type BlobStorageOptions,
+} from '../../storage';
 
-export class CloudBlobStorage extends BlobStorage {
-  private readonly gql = gqlFetcherFactory(this.options.peer + '/graphql');
+interface CloudBlobStorageOptions extends BlobStorageOptions {
+  apiBaseUrl: string;
+}
+
+export class CloudBlobStorage extends BlobStorageBase<CloudBlobStorageOptions> {
+  private readonly gql = gqlFetcherFactory(
+    this.options.apiBaseUrl + '/graphql'
+  );
   override connection = new DummyConnection();
 
   override async get(key: string) {
     const res = await fetch(
-      this.options.peer + '/api/workspaces/' + this.spaceId + '/blobs/' + key,
+      this.options.apiBaseUrl +
+        '/api/workspaces/' +
+        this.spaceId +
+        '/blobs/' +
+        key,
       {
+        cache: 'default',
         headers: {
           'x-affine-version': BUILD_CONFIG.appVersion,
         },

packages/common/nbstore/src/impls/cloud/doc.ts

@@ -1,10 +1,14 @@
-import type { SocketOptions } from 'socket.io-client';
+import type { Socket, SocketOptions } from 'socket.io-client';
 
-import { share } from '../../connection';
+import {
+  type Connection,
+  type ConnectionStatus,
+  share,
+} from '../../connection';
 import {
   type DocClock,
   type DocClocks,
-  DocStorage,
+  DocStorageBase,
   type DocStorageOptions,
   type DocUpdate,
 } from '../../storage';
@@ -17,63 +21,14 @@ import {
 
 interface CloudDocStorageOptions extends DocStorageOptions {
   socketOptions: SocketOptions;
+  serverBaseUrl: string;
 }
 
-export class CloudDocStorage extends DocStorage<CloudDocStorageOptions> {
-  connection = share(
-    new SocketConnection(this.peer, this.options.socketOptions)
-  );
-
-  private disposeConnectionStatusListener?: () => void;
-
-  private get socket() {
+export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
+  get socket() {
     return this.connection.inner;
   }
 
-  override connect() {
-    if (!this.disposeConnectionStatusListener) {
-      this.disposeConnectionStatusListener = this.connection.onStatusChanged(
-        status => {
-          if (status === 'connected') {
-            this.join().catch(err => {
-              console.error('doc storage join failed', err);
-            });
-            this.socket.on('space:broadcast-doc-update', this.onServerUpdate);
-          }
-        }
-      );
-    }
-    super.connect();
-  }
-
-  override disconnect() {
-    if (this.disposeConnectionStatusListener) {
-      this.disposeConnectionStatusListener();
-    }
-    this.socket.emit('space:leave', {
-      spaceType: this.spaceType,
-      spaceId: this.spaceId,
-    });
-    this.socket.off('space:broadcast-doc-update', this.onServerUpdate);
-    super.disconnect();
-  }
-
-  async join() {
-    try {
-      const res = await this.socket.emitWithAck('space:join', {
-        spaceType: this.spaceType,
-        spaceId: this.spaceId,
-        clientVersion: BUILD_CONFIG.appVersion,
-      });
-
-      if ('error' in res) {
-        this.connection.setStatus('closed', new Error(res.error.message));
-      }
-    } catch (e) {
-      this.connection.setStatus('error', e as Error);
-    }
-  }
-
   onServerUpdate: ServerEventsMap['space:broadcast-doc-update'] = message => {
     if (
       this.spaceType === message.spaceType &&
@@ -88,6 +43,11 @@ export class CloudDocStorage extends DocStorage<CloudDocStorageOptions> {
     }
   };
 
+  readonly connection = new CloudDocStorageConnection(
+    this.options,
+    this.onServerUpdate
+  );
+
   override async getDocSnapshot(docId: string) {
     const response = await this.socket.emitWithAck('space:load-doc', {
       spaceType: this.spaceType,
@@ -207,3 +167,84 @@ export class CloudDocStorage extends DocStorage<CloudDocStorageOptions> {
     return 0;
   }
 }
+
+class CloudDocStorageConnection implements Connection<Socket> {
+  connection = share(
+    new SocketConnection(
+      `${this.options.serverBaseUrl}/`,
+      this.options.socketOptions
+    )
+  );
+
+  private disposeConnectionStatusListener?: () => void;
+
+  private get socket() {
+    return this.connection.inner;
+  }
+
+  constructor(
+    private readonly options: CloudDocStorageOptions,
+    private readonly onServerUpdate: ServerEventsMap['space:broadcast-doc-update']
+  ) {}
+
+  get status() {
+    return this.connection.status;
+  }
+
+  get inner() {
+    return this.connection.inner;
+  }
+
+  connect(): void {
+    if (!this.disposeConnectionStatusListener) {
+      this.disposeConnectionStatusListener = this.connection.onStatusChanged(
+        status => {
+          if (status === 'connected') {
+            this.join().catch(err => {
+              console.error('doc storage join failed', err);
+            });
+            this.socket.on('space:broadcast-doc-update', this.onServerUpdate);
+          }
+        }
+      );
+    }
+    return this.connection.connect();
+  }
+
+  async join() {
+    try {
+      const res = await this.socket.emitWithAck('space:join', {
+        spaceType: this.options.type,
+        spaceId: this.options.id,
+        clientVersion: BUILD_CONFIG.appVersion,
+      });
+
+      if ('error' in res) {
+        this.connection.setStatus('closed', new Error(res.error.message));
+      }
+    } catch (e) {
+      this.connection.setStatus('error', e as Error);
+    }
+  }
+
+  disconnect() {
+    if (this.disposeConnectionStatusListener) {
+      this.disposeConnectionStatusListener();
+    }
+    this.socket.emit('space:leave', {
+      spaceType: this.options.type,
+      spaceId: this.options.id,
+    });
+    this.socket.off('space:broadcast-doc-update', this.onServerUpdate);
+    this.connection.disconnect();
+  }
+
+  waitForConnected(signal?: AbortSignal): Promise<void> {
+    return this.connection.waitForConnected(signal);
+  }
+  onStatusChanged(
+    cb: (status: ConnectionStatus, error?: Error) => void
+  ): () => void {
+    return this.connection.onStatusChanged(cb);
+  }
+}

packages/common/nbstore/src/impls/cloud/socket.ts

@@ -4,7 +4,10 @@ import {
   type SocketOptions,
 } from 'socket.io-client';
 
-import { Connection, type ConnectionStatus } from '../../connection';
+import {
+  AutoReconnectConnection,
+  type ConnectionStatus,
+} from '../../connection';
 
 // TODO(@forehalo): use [UserFriendlyError]
 interface EventError {
@@ -150,7 +153,7 @@ export function base64ToUint8Array(base64: string) {
   return new Uint8Array(binaryArray);
 }
 
-export class SocketConnection extends Connection<Socket> {
+export class SocketConnection extends AutoReconnectConnection<Socket> {
   manager = new SocketIOManager(this.endpoint, {
     autoConnect: false,
     transports: ['websocket'],