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

packages/common/nbstore/src/connection/connection.ts

@@ -1,5 +1,6 @@
 import EventEmitter2 from 'eventemitter2';
-import { throttle } from 'lodash-es';
+
+import { MANUALLY_STOP } from '../utils/throw-if-aborted';
 
 export type ConnectionStatus =
   | 'idle'
@@ -10,6 +11,7 @@ export type ConnectionStatus =
 
 export interface Connection<T = any> {
   readonly status: ConnectionStatus;
+  readonly error?: Error;
   readonly inner: T;
   connect(): void;
   disconnect(): void;
@@ -23,16 +25,15 @@ export abstract class AutoReconnectConnection<T = any>
   implements Connection<T>
 {
   private readonly event = new EventEmitter2();
-  private _inner: T | null = null;
+  private _inner: T | undefined = undefined;
   private _status: ConnectionStatus = 'idle';
-  protected error?: Error;
+  private _error: Error | undefined = undefined;
+  retryDelay = 3000;
   private refCount = 0;
-  private _enableAutoReconnect = false;
   private connectingAbort?: AbortController;
+  private reconnectingAbort?: AbortController;
 
-  constructor() {
-    this.autoReconnect();
-  }
+  constructor() {}
 
   get shareId(): string | undefined {
     return undefined;
@@ -43,7 +44,7 @@ export abstract class AutoReconnectConnection<T = any>
   }
 
   get inner(): T {
-    if (!this._inner) {
+    if (this._inner === undefined) {
       throw new Error(
         `Connection ${this.constructor.name} has not been established.`
       );
@@ -52,7 +53,7 @@ export abstract class AutoReconnectConnection<T = any>
     return this._inner;
   }
 
-  protected set inner(inner: T | null) {
+  private set inner(inner: T | undefined) {
     this._inner = inner;
   }
 
@@ -60,12 +61,23 @@ export abstract class AutoReconnectConnection<T = any>
     return this._status;
   }
 
-  protected setStatus(status: ConnectionStatus, error?: Error) {
-    const shouldEmit = status !== this._status || error !== this.error;
+  get error() {
+    return this._error;
+  }
+
+  protected set error(error: Error | undefined) {
+    this.handleError(error);
+  }
+
+  private setStatus(status: ConnectionStatus, error?: Error) {
+    const shouldEmit = status !== this._status || error !== this._error;
     this._status = status;
-    this.error = error;
+    // we only clear-up error when status is connected
+    if (error || status === 'connected') {
+      this._error = error;
+    }
     if (shouldEmit) {
-      this.emitStatusChanged(status, error);
+      this.emitStatusChanged(status, this._error);
     }
   }
 
@@ -73,15 +85,15 @@ export abstract class AutoReconnectConnection<T = any>
   protected abstract doDisconnect(conn: T): void;
 
   private innerConnect() {
-    if (this.status === 'idle' || this.status === 'error') {
-      this._enableAutoReconnect = true;
+    if (this.status !== 'connecting') {
       this.setStatus('connecting');
       this.connectingAbort = new AbortController();
-      this.doConnect(this.connectingAbort.signal)
+      const signal = this.connectingAbort.signal;
+      this.doConnect(signal)
         .then(value => {
-          if (!this.connectingAbort?.signal.aborted) {
-            this.setStatus('connected');
+          if (!signal.aborted) {
             this._inner = value;
+            this.setStatus('connected');
           } else {
             try {
               this.doDisconnect(value);
@@ -91,14 +103,45 @@ export abstract class AutoReconnectConnection<T = any>
           }
         })
         .catch(error => {
-          if (!this.connectingAbort?.signal.aborted) {
+          if (!signal.aborted) {
             console.error('failed to connect', error);
-            this.setStatus('error', error as any);
+            this.handleError(error as any);
           }
         });
     }
   }
 
+  private innerDisconnect() {
+    this.connectingAbort?.abort(MANUALLY_STOP);
+    this.reconnectingAbort?.abort(MANUALLY_STOP);
+    try {
+      if (this._inner) {
+        this.doDisconnect(this._inner);
+      }
+    } catch (error) {
+      console.error('failed to disconnect', error);
+    }
+    this.reconnectingAbort = undefined;
+    this.connectingAbort = undefined;
+    this._inner = undefined;
+  }
+
+  private handleError(reason?: Error) {
+    // on error
+    console.error('connection error, will reconnect', reason);
+    this.innerDisconnect();
+    this.setStatus('error', reason);
+    // reconnect
+
+    this.reconnectingAbort = new AbortController();
+    const signal = this.reconnectingAbort.signal;
+    setTimeout(() => {
+      if (!signal.aborted) {
+        this.innerConnect();
+      }
+    }, this.retryDelay);
+  }
+
   connect() {
     this.refCount++;
     if (this.refCount === 1) {
@@ -106,36 +149,16 @@ export abstract class AutoReconnectConnection<T = any>
     }
   }
 
-  disconnect() {
-    this.refCount--;
-    if (this.refCount === 0) {
-      this._enableAutoReconnect = false;
-      this.connectingAbort?.abort();
-      try {
-        if (this._inner) {
-          this.doDisconnect(this._inner);
-        }
-      } catch (error) {
-        console.error('failed to disconnect', error);
-      }
-      this.setStatus('closed');
-      this._inner = null;
+  disconnect(force?: boolean) {
+    if (force) {
+      this.refCount = 0;
+    } else {
+      this.refCount = Math.max(this.refCount - 1, 0);
+    }
+    if (this.refCount === 0) {
+      this.innerDisconnect();
+      this.setStatus('closed');
     }
-  }
-
-  private autoReconnect() {
-    // TODO:
-    //   - maximum retry count
-    //   - dynamic sleep time (attempt < 3 ? 1s : 1min)?
-    this.onStatusChanged(
-      throttle(() => {
-        () => {
-          if (this._enableAutoReconnect) {
-            this.innerConnect();
-          }
-        };
-      }, 1000)
-    );
   }
 
   waitForConnected(signal?: AbortSignal) {