fix(nbstore): connect before do operation (#11569)

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

@@ -172,28 +172,14 @@ class WorkerDocConnection extends DummyConnection {
     super();
   }
 
-  override waitForConnected(signal?: AbortSignal): Promise<void> {
-    return new Promise((resolve, reject) => {
-      const abortListener = () => {
-        reject(signal?.reason);
-        subscription.unsubscribe();
-      };
+  promise: Promise<void> | undefined;
 
-      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);
-          },
-        });
-    });
+  override waitForConnected(): Promise<void> {
+    if (this.promise) {
+      return this.promise;
+    }
+    this.promise = this.client.call('docStorage.waitForConnected');
+    return this.promise;
   }
 }
 
@@ -226,7 +212,23 @@ class WorkerBlobStorage implements BlobStorage {
     return this.client.call('blobStorage.listBlobs');
   }
 
-  connection = new DummyConnection();
+  connection = new WorkerBlobConnection(this.client);
+}
+
+class WorkerBlobConnection extends DummyConnection {
+  constructor(private readonly client: OpClient<WorkerOps>) {
+    super();
+  }
+
+  promise: Promise<void> | undefined;
+
+  override waitForConnected(): Promise<void> {
+    if (this.promise) {
+      return this.promise;
+    }
+    this.promise = this.client.call('blobStorage.waitForConnected');
+    return this.promise;
+  }
 }
 
 class WorkerDocSync implements DocSync {
@@ -346,7 +348,7 @@ class WorkerIndexerStorage implements IndexerStorage {
   constructor(private readonly client: OpClient<WorkerOps>) {}
   readonly storageType = 'indexer';
   readonly isReadonly = true;
-  connection = new DummyConnection();
+  connection = new WorkerIndexerConnection(this.client);
 
   search<T extends keyof IndexerSchema, const O extends SearchOptions<T>>(
     table: T,
@@ -421,6 +423,22 @@ class WorkerIndexerStorage implements IndexerStorage {
   }
 }
 
+class WorkerIndexerConnection extends DummyConnection {
+  constructor(private readonly client: OpClient<WorkerOps>) {
+    super();
+  }
+
+  promise: Promise<void> | undefined;
+
+  override waitForConnected(): Promise<void> {
+    if (this.promise) {
+      return this.promise;
+    }
+    this.promise = this.client.call('indexerStorage.waitForConnected');
+    return this.promise;
+  }
+}
+
 class WorkerIndexerSync implements IndexerSync {
   constructor(private readonly client: OpClient<WorkerOps>) {}
   waitForCompleted(signal?: AbortSignal): Promise<void> {

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

@@ -158,26 +158,16 @@ class StoreConsumer {
             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: any) => {
-              subscriber.error(error);
-            });
-          return () => abortController.abort(MANUALLY_STOP);
-        }),
+      'docStorage.waitForConnected': (_, ctx) =>
+        this.docStorage.connection.waitForConnected(ctx.signal),
       '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(),
+      'blobStorage.waitForConnected': (_, ctx) =>
+        this.blobStorage.connection.waitForConnected(ctx.signal),
       'awarenessStorage.update': ({ awareness, origin }) =>
         this.awarenessStorage.update(awareness, origin),
       'awarenessStorage.subscribeUpdate': docId =>
@@ -205,6 +195,8 @@ class StoreConsumer {
         }),
       'awarenessStorage.collect': ({ collectId, awareness }) =>
         collectJobs.get(collectId)?.(awareness),
+      'awarenessStorage.waitForConnected': (_, ctx) =>
+        this.awarenessStorage.connection.waitForConnected(ctx.signal),
       'docSync.state': () => this.docSync.state$,
       'docSync.docState': docId =>
         new Observable(subscriber => {
@@ -278,6 +270,8 @@ class StoreConsumer {
         this.indexerStorage.search$(table, query, options),
       'indexerStorage.subscribeAggregate': ({ table, query, field, options }) =>
         this.indexerStorage.aggregate$(table, query, field, options),
+      'indexerStorage.waitForConnected': (_, ctx) =>
+        this.indexerStorage.connection.waitForConnected(ctx.signal),
       'indexerSync.state': () => this.indexerSync.state$,
       'indexerSync.docState': (docId: string) =>
         this.indexerSync.docState$(docId),
@@ -286,32 +280,10 @@ class StoreConsumer {
           const undo = this.indexerSync.addPriority(docId, priority);
           return () => undo();
         }),
-      'indexerSync.waitForCompleted': () =>
-        new Observable(subscriber => {
-          this.indexerSync
-            .waitForCompleted()
-            .then(() => {
-              subscriber.next();
-              subscriber.complete();
-            })
-            .catch(error => {
-              subscriber.error(error);
-            });
-        }),
-      'indexerSync.waitForDocCompleted': (docId: string) =>
-        new Observable(subscriber => {
-          const abortController = new AbortController();
-          this.indexerSync
-            .waitForDocCompleted(docId, abortController.signal)
-            .then(() => {
-              subscriber.next();
-              subscriber.complete();
-            })
-            .catch(error => {
-              subscriber.error(error);
-            });
-          return () => abortController.abort(MANUALLY_STOP);
-        }),
+      'indexerSync.waitForCompleted': (_, ctx) =>
+        this.indexerSync.waitForCompleted(ctx.signal),
+      'indexerSync.waitForDocCompleted': (docId: string, ctx) =>
+        this.indexerSync.waitForDocCompleted(docId, ctx.signal),
     });
   }
 }

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

@@ -40,7 +40,7 @@ interface GroupedWorkerOps {
     getDocTimestamp: [string, DocClock | null];
     deleteDoc: [string, void];
     subscribeDocUpdate: [void, { update: DocRecord; origin?: string }];
-    waitForConnected: [void, boolean];
+    waitForConnected: [void, void];
   };
 
   blobStorage: {
@@ -49,6 +49,7 @@ interface GroupedWorkerOps {
     deleteBlob: [{ key: string; permanently: boolean }, void];
     releaseBlobs: [void, void];
     listBlobs: [void, ListedBlobRecord[]];
+    waitForConnected: [void, void];
   };
 
   awarenessStorage: {
@@ -65,6 +66,7 @@ interface GroupedWorkerOps {
       ),
     ];
     collect: [{ collectId: string; awareness: AwarenessRecord }, void];
+    waitForConnected: [void, void];
   };
 
   indexerStorage: {
@@ -94,6 +96,7 @@ interface GroupedWorkerOps {
       },
       AggregateResult<any, any>,
     ];
+    waitForConnected: [void, void];
   };
 
   docSync: {