feat(core): new worker workspace engine (#9257)

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

@@ -1,5 +1,3 @@
-import type { SocketOptions } from 'socket.io-client';
-
 import { share } from '../../connection';
 import {
   type AwarenessRecord,
@@ -13,7 +11,6 @@ import {
 } from './socket';
 
 interface CloudAwarenessStorageOptions {
-  socketOptions?: SocketOptions;
   serverBaseUrl: string;
   type: SpaceType;
   id: string;
@@ -26,12 +23,7 @@ export class CloudAwarenessStorage extends AwarenessStorageBase {
     super();
   }
 
-  connection = share(
-    new SocketConnection(
-      `${this.options.serverBaseUrl}/`,
-      this.options.socketOptions
-    )
-  );
+  connection = share(new SocketConnection(`${this.options.serverBaseUrl}/`));
 
   private get socket() {
     return this.connection.inner;
@@ -52,9 +44,14 @@ export class CloudAwarenessStorage extends AwarenessStorageBase {
     onUpdate: (update: AwarenessRecord, origin?: string) => void,
     onCollect: () => Promise<AwarenessRecord | null>
   ): () => void {
-    // TODO: handle disconnect
     // leave awareness
     const leave = () => {
+      if (this.connection.status !== 'connected') return;
+      this.socket.off('space:collect-awareness', handleCollectAwareness);
+      this.socket.off(
+        'space:broadcast-awareness-update',
+        handleBroadcastAwarenessUpdate
+      );
       this.socket.emit('space:leave-awareness', {
         spaceType: this.options.type,
         spaceId: this.options.id,
@@ -64,6 +61,11 @@ export class CloudAwarenessStorage extends AwarenessStorageBase {
 
     // join awareness, and collect awareness from others
     const joinAndCollect = async () => {
+      this.socket.on('space:collect-awareness', handleCollectAwareness);
+      this.socket.on(
+        'space:broadcast-awareness-update',
+        handleBroadcastAwarenessUpdate
+      );
       await this.socket.emitWithAck('space:join-awareness', {
         spaceType: this.options.type,
         spaceId: this.options.id,
@@ -77,7 +79,11 @@ export class CloudAwarenessStorage extends AwarenessStorageBase {
       });
     };
 
-    joinAndCollect().catch(err => console.error('awareness join failed', err));
+    if (this.connection.status === 'connected') {
+      joinAndCollect().catch(err =>
+        console.error('awareness join failed', err)
+      );
+    }
 
     const unsubscribeConnectionStatusChanged = this.connection.onStatusChanged(
       status => {
@@ -141,18 +147,9 @@ export class CloudAwarenessStorage extends AwarenessStorageBase {
       }
     };
 
-    this.socket.on('space:collect-awareness', handleCollectAwareness);
-    this.socket.on(
-      'space:broadcast-awareness-update',
-      handleBroadcastAwarenessUpdate
-    );
     return () => {
       leave();
-      this.socket.off('space:collect-awareness', handleCollectAwareness);
-      this.socket.off(
-        'space:broadcast-awareness-update',
-        handleBroadcastAwarenessUpdate
-      );
+
       unsubscribeConnectionStatusChanged();
     };
   }

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

@@ -45,23 +45,28 @@ export class StaticCloudDocStorage extends DocStorageBase<CloudDocStorageOptions
   protected override async getDocSnapshot(
     docId: string
   ): Promise<DocRecord | null> {
-    const arrayBuffer = await this.connection.fetchArrayBuffer(
-      `/api/workspaces/${this.spaceId}/docs/${docId}`,
-      {
-        priority: 'high',
-        headers: {
-          Accept: 'application/octet-stream', // this is necessary for ios native fetch to return arraybuffer
-        },
+    try {
+      const arrayBuffer = await this.connection.fetchArrayBuffer(
+        `/api/workspaces/${this.spaceId}/docs/${docId}`,
+        {
+          priority: 'high',
+          headers: {
+            Accept: 'application/octet-stream', // this is necessary for ios native fetch to return arraybuffer
+          },
+        }
+      );
+      if (!arrayBuffer) {
+        return null;
       }
-    );
-    if (!arrayBuffer) {
+      return {
+        docId: docId,
+        bin: new Uint8Array(arrayBuffer),
+        timestamp: new Date(),
+      };
+    } catch (error) {
+      console.error(error);
       return null;
     }
-    return {
-      docId: docId,
-      bin: new Uint8Array(arrayBuffer),
-      timestamp: new Date(),
-    };
   }
   protected override setDocSnapshot(
     _snapshot: DocRecord,

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

@@ -1,10 +1,5 @@
-import type { Socket, SocketOptions } from 'socket.io-client';
+import type { Socket } from 'socket.io-client';
 
-import {
-  type Connection,
-  type ConnectionStatus,
-  share,
-} from '../../connection';
 import {
   type DocClock,
   type DocClocks,
@@ -12,6 +7,7 @@ import {
   type DocStorageOptions,
   type DocUpdate,
 } from '../../storage';
+import { getIdConverter, type IdConverter } from '../../utils/id-converter';
 import type { SpaceType } from '../../utils/universal-id';
 import {
   base64ToUint8Array,
@@ -21,7 +17,6 @@ import {
 } from './socket';
 
 interface CloudDocStorageOptions extends DocStorageOptions {
-  socketOptions?: SocketOptions;
   serverBaseUrl: string;
   type: SpaceType;
 }
@@ -32,7 +27,12 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
   get socket() {
     return this.connection.inner;
   }
-
+  get idConverter() {
+    if (!this.connection.idConverter) {
+      throw new Error('Id converter not initialized');
+    }
+    return this.connection.idConverter;
+  }
   readonly spaceType = this.options.type;
 
   onServerUpdate: ServerEventsMap['space:broadcast-doc-update'] = message => {
@@ -41,7 +41,7 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
       this.spaceId === message.spaceId
     ) {
       this.emit('update', {
-        docId: message.docId,
+        docId: this.idConverter.oldIdToNewId(message.docId),
         bin: base64ToUint8Array(message.update),
         timestamp: new Date(message.timestamp),
         editor: message.editor,
@@ -58,10 +58,13 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
     const response = await this.socket.emitWithAck('space:load-doc', {
       spaceType: this.spaceType,
       spaceId: this.spaceId,
-      docId,
+      docId: this.idConverter.newIdToOldId(docId),
     });
 
     if ('error' in response) {
+      if (response.error.name === 'DOC_NOT_FOUND') {
+        return null;
+      }
       // TODO: use [UserFriendlyError]
       throw new Error(response.error.message);
     }
@@ -77,11 +80,14 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
     const response = await this.socket.emitWithAck('space:load-doc', {
       spaceType: this.spaceType,
       spaceId: this.spaceId,
-      docId,
+      docId: this.idConverter.newIdToOldId(docId),
       stateVector: state ? await uint8ArrayToBase64(state) : void 0,
     });
 
     if ('error' in response) {
+      if (response.error.name === 'DOC_NOT_FOUND') {
+        return null;
+      }
       // TODO: use [UserFriendlyError]
       throw new Error(response.error.message);
     }
@@ -98,8 +104,8 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
     const response = await this.socket.emitWithAck('space:push-doc-update', {
       spaceType: this.spaceType,
       spaceId: this.spaceId,
-      docId: update.docId,
-      updates: await uint8ArrayToBase64(update.bin),
+      docId: this.idConverter.newIdToOldId(update.docId),
+      update: await uint8ArrayToBase64(update.bin),
     });
 
     if ('error' in response) {
@@ -120,7 +126,7 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
     const response = await this.socket.emitWithAck('space:load-doc', {
       spaceType: this.spaceType,
       spaceId: this.spaceId,
-      docId,
+      docId: this.idConverter.newIdToOldId(docId),
     });
 
     if ('error' in response) {
@@ -150,7 +156,7 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
     }
 
     return Object.entries(response.data).reduce((ret, [docId, timestamp]) => {
-      ret[docId] = new Date(timestamp);
+      ret[this.idConverter.oldIdToNewId(docId)] = new Date(timestamp);
       return ret;
     }, {} as DocClocks);
   }
@@ -159,7 +165,7 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
     this.socket.emit('space:delete-doc', {
       spaceType: this.spaceType,
       spaceId: this.spaceId,
-      docId,
+      docId: this.idConverter.newIdToOldId(docId),
     });
   }
 
@@ -174,83 +180,74 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
   }
 }
 
-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;
-  }
-
+class CloudDocStorageConnection extends SocketConnection {
   constructor(
     private readonly options: CloudDocStorageOptions,
     private readonly onServerUpdate: ServerEventsMap['space:broadcast-doc-update']
-  ) {}
-
-  get status() {
-    return this.connection.status;
+  ) {
+    super(`${options.serverBaseUrl}/`);
   }
 
-  get inner() {
-    return this.connection.inner;
-  }
+  idConverter: IdConverter | null = null;
 
-  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();
-  }
+  override async doConnect(signal?: AbortSignal) {
+    const socket = await super.doConnect(signal);
 
-  async join() {
     try {
-      const res = await this.socket.emitWithAck('space:join', {
+      const res = await 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));
+        throw new Error(res.error.message);
       }
+
+      if (!this.idConverter) {
+        this.idConverter = await this.getIdConverter(socket);
+      }
+
+      socket.on('space:broadcast-doc-update', this.onServerUpdate);
+
+      return socket;
     } catch (e) {
-      this.connection.setStatus('error', e as Error);
+      socket.close();
+      throw e;
     }
   }
 
-  disconnect() {
-    if (this.disposeConnectionStatusListener) {
-      this.disposeConnectionStatusListener();
-    }
-    this.socket.emit('space:leave', {
+  override doDisconnect(socket: Socket) {
+    socket.emit('space:leave', {
       spaceType: this.options.type,
       spaceId: this.options.id,
     });
-    this.socket.off('space:broadcast-doc-update', this.onServerUpdate);
-    this.connection.disconnect();
+    socket.off('space:broadcast-doc-update', this.onServerUpdate);
+    super.disconnect();
   }
 
-  waitForConnected(signal?: AbortSignal): Promise<void> {
-    return this.connection.waitForConnected(signal);
-  }
-  onStatusChanged(
-    cb: (status: ConnectionStatus, error?: Error) => void
-  ): () => void {
-    return this.connection.onStatusChanged(cb);
+  async getIdConverter(socket: Socket) {
+    return getIdConverter(
+      {
+        getDocBuffer: async id => {
+          const response = await socket.emitWithAck('space:load-doc', {
+            spaceType: this.options.type,
+            spaceId: this.options.id,
+            docId: id,
+          });
+
+          if ('error' in response) {
+            if (response.error.name === 'DOC_NOT_FOUND') {
+              return null;
+            }
+            // TODO: use [UserFriendlyError]
+            throw new Error(response.error.message);
+          }
+
+          return base64ToUint8Array(response.data.missing);
+        },
+      },
+      this.options.id
+    );
   }
 }

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

@@ -23,6 +23,7 @@ export class HttpConnection extends DummyConnection {
         ...init,
         signal: abortController.signal,
         headers: {
+          ...this.requestHeaders,
           ...init?.headers,
           'x-affine-version': BUILD_CONFIG.appVersion,
         },
@@ -35,7 +36,7 @@ export class HttpConnection extends DummyConnection {
       let reason: string | any = '';
       if (res.headers.get('Content-Type')?.includes('application/json')) {
         try {
-          reason = await res.json();
+          reason = JSON.stringify(await res.json());
         } catch {
           // ignore
         }
@@ -63,7 +64,10 @@ export class HttpConnection extends DummyConnection {
     this.fetch
   );
 
-  constructor(private readonly serverBaseUrl: string) {
+  constructor(
+    private readonly serverBaseUrl: string,
+    private readonly requestHeaders?: Record<string, string>
+  ) {
     super();
   }
 }

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

@@ -4,10 +4,8 @@ import {
   type SocketOptions,
 } from 'socket.io-client';
 
-import {
-  AutoReconnectConnection,
-  type ConnectionStatus,
-} from '../../connection';
+import { AutoReconnectConnection } from '../../connection';
+import { throwIfAborted } from '../../utils/throw-if-aborted';
 
 // TODO(@forehalo): use [UserFriendlyError]
 interface EventError {
@@ -82,7 +80,7 @@ interface ClientEvents {
   };
 
   'space:push-doc-update': [
-    { spaceType: string; spaceId: string; docId: string; updates: string },
+    { spaceType: string; spaceId: string; docId: string; update: string },
     { timestamp: number },
   ];
   'space:load-doc-timestamps': [
@@ -153,12 +151,24 @@ export function base64ToUint8Array(base64: string) {
   return new Uint8Array(binaryArray);
 }
 
+const SOCKET_IOMANAGER_CACHE = new Map<string, SocketIOManager>();
+function getSocketIOManager(endpoint: string) {
+  let manager = SOCKET_IOMANAGER_CACHE.get(endpoint);
+  if (!manager) {
+    manager = new SocketIOManager(endpoint, {
+      autoConnect: false,
+      transports: ['websocket'],
+      secure: new URL(endpoint).protocol === 'https:',
+      // we will handle reconnection by ourselves
+      reconnection: false,
+    });
+    SOCKET_IOMANAGER_CACHE.set(endpoint, manager);
+  }
+  return manager;
+}
+
 export class SocketConnection extends AutoReconnectConnection<Socket> {
-  manager = new SocketIOManager(this.endpoint, {
-    autoConnect: false,
-    transports: ['websocket'],
-    secure: new URL(this.endpoint).protocol === 'https:',
-  });
+  manager = getSocketIOManager(this.endpoint);
 
   constructor(
     private readonly endpoint: string,
@@ -171,32 +181,42 @@ export class SocketConnection extends AutoReconnectConnection<Socket> {
     return `socket:${this.endpoint}`;
   }
 
-  override async doConnect() {
-    const conn = this.manager.socket('/', this.socketOptions);
+  override async doConnect(signal?: AbortSignal) {
+    const socket = this.manager.socket('/', this.socketOptions);
+    try {
+      throwIfAborted(signal);
+      await Promise.race([
+        new Promise<void>((resolve, reject) => {
+          socket.once('connect', () => {
+            resolve();
+          });
+          socket.once('connect_error', err => {
+            reject(err);
+          });
+          socket.open();
+        }),
+        new Promise<void>((_resolve, reject) => {
+          signal?.addEventListener('abort', () => {
+            reject(new Error('Aborted'));
+          });
+        }),
+      ]);
+    } catch (err) {
+      socket.close();
+      throw err;
+    }
 
-    await new Promise<void>((resolve, reject) => {
-      conn.once('connect', () => {
-        resolve();
-      });
-      conn.once('connect_error', err => {
-        reject(err);
-      });
-      conn.open();
-    });
+    socket.on('disconnect', this.handleDisconnect);
 
-    return conn;
+    return socket;
   }
 
   override doDisconnect(conn: Socket) {
+    conn.off('disconnect', this.handleDisconnect);
     conn.close();
   }
 
-  /**
-   * Socket connection allow explicitly set status by user
-   *
-   * used when join space failed
-   */
-  override setStatus(status: ConnectionStatus, error?: Error) {
-    super.setStatus(status, error);
-  }
+  handleDisconnect = (reason: SocketIO.DisconnectReason) => {
+    this.error = new Error(reason);
+  };
 }