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

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

@@ -1,4 +1,3 @@
-import { share } from '../../connection';
 import {
   type AwarenessRecord,
   AwarenessStorageBase,
@@ -23,10 +22,10 @@ export class CloudAwarenessStorage extends AwarenessStorageBase {
     super();
   }
 
-  connection = share(new SocketConnection(`${this.options.serverBaseUrl}/`));
+  connection = new SocketConnection(`${this.options.serverBaseUrl}/`);
 
   private get socket() {
-    return this.connection.inner;
+    return this.connection.inner.socket;
   }
 
   override async update(record: AwarenessRecord): Promise<void> {

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

@@ -25,7 +25,7 @@ export class CloudDocStorage extends DocStorageBase<CloudDocStorageOptions> {
   static readonly identifier = 'CloudDocStorage';
 
   get socket() {
-    return this.connection.inner;
+    return this.connection.inner.socket;
   }
   get idConverter() {
     if (!this.connection.idConverter) {
@@ -191,7 +191,7 @@ class CloudDocStorageConnection extends SocketConnection {
   idConverter: IdConverter | null = null;
 
   override async doConnect(signal?: AbortSignal) {
-    const socket = await super.doConnect(signal);
+    const { socket, disconnect } = await super.doConnect(signal);
 
     try {
       const res = await socket.emitWithAck('space:join', {
@@ -210,20 +210,26 @@ class CloudDocStorageConnection extends SocketConnection {
 
       socket.on('space:broadcast-doc-update', this.onServerUpdate);
 
-      return socket;
+      return { socket, disconnect };
     } catch (e) {
-      socket.close();
+      disconnect();
       throw e;
     }
   }
 
-  override doDisconnect(socket: Socket) {
+  override doDisconnect({
+    socket,
+    disconnect,
+  }: {
+    socket: Socket;
+    disconnect: () => void;
+  }) {
     socket.emit('space:leave', {
       spaceType: this.options.type,
       spaceId: this.options.id,
     });
     socket.off('space:broadcast-doc-update', this.onServerUpdate);
-    super.disconnect();
+    super.doDisconnect({ socket, disconnect });
   }
 
   async getIdConverter(socket: Socket) {

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

@@ -1,7 +1,6 @@
 import {
   Manager as SocketIOManager,
   type Socket as SocketIO,
-  type SocketOptions,
 } from 'socket.io-client';
 
 import { AutoReconnectConnection } from '../../connection';
@@ -151,49 +150,78 @@ 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, {
+class SocketManager {
+  private readonly socketIOManager: SocketIOManager;
+  socket: Socket;
+  refCount = 0;
+
+  constructor(endpoint: string) {
+    this.socketIOManager = 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);
+    this.socket = this.socketIOManager.socket('/');
+  }
+
+  connect() {
+    let disconnected = false;
+    this.refCount++;
+    this.socket.connect();
+    return {
+      socket: this.socket,
+      disconnect: () => {
+        if (disconnected) {
+          return;
+        }
+        disconnected = true;
+        this.refCount--;
+        if (this.refCount === 0) {
+          this.socket.disconnect();
+        }
+      },
+    };
+  }
+}
+
+const SOCKET_MANAGER_CACHE = new Map<string, SocketManager>();
+function getSocketManager(endpoint: string) {
+  let manager = SOCKET_MANAGER_CACHE.get(endpoint);
+  if (!manager) {
+    manager = new SocketManager(endpoint);
+    SOCKET_MANAGER_CACHE.set(endpoint, manager);
   }
   return manager;
 }
 
-export class SocketConnection extends AutoReconnectConnection<Socket> {
-  manager = getSocketIOManager(this.endpoint);
+export class SocketConnection extends AutoReconnectConnection<{
+  socket: Socket;
+  disconnect: () => void;
+}> {
+  manager = getSocketManager(this.endpoint);
 
-  constructor(
-    private readonly endpoint: string,
-    private readonly socketOptions?: SocketOptions
-  ) {
+  constructor(private readonly endpoint: string) {
     super();
   }
 
-  override get shareId() {
-    return `socket:${this.endpoint}`;
-  }
-
   override async doConnect(signal?: AbortSignal) {
-    const socket = this.manager.socket('/', this.socketOptions);
+    const { socket, disconnect } = this.manager.connect();
     try {
       throwIfAborted(signal);
       await Promise.race([
         new Promise<void>((resolve, reject) => {
+          if (socket.connected) {
+            resolve();
+            return;
+          }
           socket.once('connect', () => {
             resolve();
           });
           socket.once('connect_error', err => {
             reject(err);
           });
-          socket.open();
         }),
         new Promise<void>((_resolve, reject) => {
           signal?.addEventListener('abort', () => {
@@ -202,18 +230,21 @@ export class SocketConnection extends AutoReconnectConnection<Socket> {
         }),
       ]);
     } catch (err) {
-      socket.close();
+      disconnect();
       throw err;
     }
 
     socket.on('disconnect', this.handleDisconnect);
 
-    return socket;
+    return {
+      socket,
+      disconnect,
+    };
   }
 
-  override doDisconnect(conn: Socket) {
-    conn.off('disconnect', this.handleDisconnect);
-    conn.close();
+  override doDisconnect(conn: { socket: Socket; disconnect: () => void }) {
+    conn.socket.off('disconnect', this.handleDisconnect);
+    conn.disconnect();
   }
 
   handleDisconnect = (reason: SocketIO.DisconnectReason) => {

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

@@ -15,7 +15,7 @@ export class SqliteDocStorage extends DocStorageBase<SqliteNativeDBOptions> {
     return this.connection.apis;
   }
 
-  override async pushDocUpdate(update: DocUpdate) {
+  override async pushDocUpdate(update: DocUpdate, origin?: string) {
     const timestamp = await this.db.pushUpdate(update.docId, update.bin);
 
     this.emit(