feat(nbstore): remove async on connection api (#9187)

We should not use async on `connect` and `disconnect`, for `WebSocketConnection` will never connect when offline.

We should handle the connection status of each storage in sync, using the `connection.waitForConnect`

This PR also puts the connection reference count on the `connect` and disconnect`

packages/common/nbstore/src/impls/broadcast-channel/channel.ts

@@ -12,7 +12,7 @@ export class BroadcastChannelConnection extends Connection<BroadcastChannel> {
     return new BroadcastChannel(this.channelName);
   }
 
-  override async doDisconnect() {
+  override doDisconnect() {
     this.close();
   }
 

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

@@ -25,10 +25,6 @@ export class CloudAwarenessStorage extends AwarenessStorage<CloudAwarenessStorag
     return this.connection.inner;
   }
 
-  override async connect(): Promise<void> {
-    await super.connect();
-  }
-
   override async update(record: AwarenessRecord): Promise<void> {
     const encodedUpdate = await uint8ArrayToBase64(record.bin);
     this.socket.emit('space:update-awareness', {
@@ -44,6 +40,7 @@ export class CloudAwarenessStorage extends AwarenessStorage<CloudAwarenessStorag
     onUpdate: (update: AwarenessRecord, origin?: string) => void,
     onCollect: () => AwarenessRecord
   ): () => void {
+    // TODO: handle disconnect
     // leave awareness
     const leave = () => {
       this.socket.emit('space:leave-awareness', {

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

@@ -24,29 +24,38 @@ export class CloudDocStorage extends DocStorage<CloudDocStorageOptions> {
     new SocketConnection(this.peer, this.options.socketOptions)
   );
 
+  private disposeConnectionStatusListener?: () => void;
+
   private get socket() {
     return this.connection.inner;
   }
 
-  override async connect(): Promise<void> {
-    await super.connect();
-    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);
-      }
-    });
+  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 async disconnect(): Promise<void> {
+  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);
-    await super.connect();
+    super.disconnect();
   }
 
   async join() {

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

@@ -184,7 +184,7 @@ export class SocketConnection extends Connection<Socket> {
     return conn;
   }
 
-  override async doDisconnect(conn: Socket) {
+  override doDisconnect(conn: Socket) {
     conn.close();
   }
 

packages/common/nbstore/src/impls/idb/db.ts

@@ -25,7 +25,8 @@ export class IDBConnection extends Connection<{
         blocking: () => {
           // if, for example, an tab with newer version is opened, this function will be called.
           // we should close current connection to allow the new version to upgrade the db.
-          this.close(
+          this.setStatus(
+            'closed',
             new Error('Blocking a new version. Closing the connection.')
           );
         },
@@ -38,13 +39,11 @@ export class IDBConnection extends Connection<{
     };
   }
 
-  override async doDisconnect() {
-    this.close();
-  }
-
-  private close(error?: Error) {
-    this.maybeConnection?.channel.close();
-    this.maybeConnection?.db.close();
-    this.setStatus('closed', error);
+  override doDisconnect(db: {
+    db: IDBPDatabase<DocStorageSchema>;
+    channel: BroadcastChannel;
+  }) {
+    db.channel.close();
+    db.db.close();
   }
 }

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

@@ -1,4 +1,3 @@
-import { share } from '../../connection';
 import {
   type DocClock,
   type DocClocks,
@@ -17,7 +16,7 @@ interface ChannelMessage {
 }
 
 export class IndexedDBDocStorage extends DocStorage {
-  readonly connection = share(new IDBConnection(this.options));
+  readonly connection = new IDBConnection(this.options);
 
   get db() {
     return this.connection.inner.db;

packages/common/nbstore/src/impls/idb/v1/db.ts

@@ -28,7 +28,7 @@ export class DocIDBConnection extends Connection<IDBPDatabase<DocDBSchema>> {
     });
   }
 
-  override async doDisconnect(conn: IDBPDatabase<DocDBSchema>) {
+  override doDisconnect(conn: IDBPDatabase<DocDBSchema>) {
     conn.close();
   }
 }
@@ -57,7 +57,7 @@ export class BlobIDBConnection extends Connection<IDBPDatabase<BlobDBSchema>> {
     });
   }
 
-  override async doDisconnect(conn: IDBPDatabase<BlobDBSchema>) {
+  override doDisconnect(conn: IDBPDatabase<BlobDBSchema>) {
     conn.close();
   }
 }

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

@@ -1,6 +1,6 @@
-import { apis, events } from '@affine/electron-api';
+import { apis } from '@affine/electron-api';
 
-import { Connection, type ConnectionStatus } from '../../connection';
+import { Connection } from '../../connection';
 import { type SpaceType, universalId } from '../../storage';
 
 type NativeDBApis = NonNullable<typeof apis>['nbstore'] extends infer APIs
@@ -27,7 +27,6 @@ export class NativeDBConnection extends Connection<void> {
     }
 
     this.apis = this.bindApis(apis.nbstore);
-    this.listenToConnectionEvents();
   }
 
   override get shareId(): string {
@@ -63,21 +62,9 @@ export class NativeDBConnection extends Connection<void> {
     await this.apis.connect();
   }
 
-  override async doDisconnect() {
-    await this.apis.close();
-  }
-
-  private listenToConnectionEvents() {
-    events?.nbstore.onConnectionStatusChanged(
-      ({ peer, spaceType, spaceId, status, error }) => {
-        if (
-          peer === this.peer &&
-          spaceType === this.type &&
-          spaceId === this.id
-        ) {
-          this.setStatus(status as ConnectionStatus, error);
-        }
-      }
-    );
+  override doDisconnect() {
+    this.apis.close().catch(err => {
+      console.error('NativeDBConnection close failed', err);
+    });
   }
 }