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

packages/common/nbstore/src/connection/__tests__/auto-reconnection.spec.ts

@@ -0,0 +1,200 @@
+import { expect, test, vitest } from 'vitest';
+
+import { AutoReconnectConnection } from '../connection';
+
+test('connect and disconnect', async () => {
+  class TestConnection extends AutoReconnectConnection<{
+    disconnect: () => void;
+  }> {
+    connectCount = 0;
+    abortCount = 0;
+    disconnectCount = 0;
+    notListenAbort = false;
+    override async doConnect(signal?: AbortSignal) {
+      this.connectCount++;
+      return new Promise<{ disconnect: () => void }>((resolve, reject) => {
+        setTimeout(() => {
+          resolve({
+            disconnect: () => {
+              this.disconnectCount++;
+            },
+          });
+        }, 300);
+        if (!this.notListenAbort) {
+          signal?.addEventListener('abort', reason => {
+            reject(reason);
+          });
+        }
+      }).catch(err => {
+        this.abortCount++;
+        throw err;
+      });
+    }
+    override doDisconnect(t: { disconnect: () => void }) {
+      return t.disconnect();
+    }
+  }
+
+  const connection = new TestConnection();
+  connection.connect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(1);
+    expect(connection.disconnectCount).toBe(0);
+    expect(connection.abortCount).toBe(0);
+    expect(connection.status).toBe('connected');
+  });
+
+  connection.disconnect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(1);
+    expect(connection.disconnectCount).toBe(1);
+    expect(connection.abortCount).toBe(0);
+    expect(connection.status).toBe('closed');
+  });
+
+  // connect twice
+  connection.connect();
+  connection.connect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(2);
+    expect(connection.disconnectCount).toBe(1);
+    expect(connection.abortCount).toBe(0);
+    expect(connection.status).toBe('connected');
+  });
+
+  connection.disconnect();
+  connection.disconnect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(2);
+    expect(connection.disconnectCount).toBe(2);
+    expect(connection.abortCount).toBe(0);
+    expect(connection.status).toBe('closed');
+  });
+
+  // calling connect disconnect consecutively, the previous connect call will be aborted.
+  connection.connect();
+  connection.disconnect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(3);
+    expect(connection.disconnectCount).toBe(2);
+    expect(connection.abortCount).toBe(1);
+    expect(connection.status).toBe('closed');
+  });
+
+  connection.connect();
+  connection.disconnect();
+  connection.connect();
+  connection.disconnect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(5);
+    expect(connection.disconnectCount).toBe(2);
+    expect(connection.abortCount).toBe(3);
+    expect(connection.status).toBe('closed');
+  });
+
+  // if connection is not listening to abort event, disconnect will be called
+  connection.notListenAbort = true;
+  connection.connect();
+  connection.disconnect();
+  connection.connect();
+  connection.disconnect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(7);
+    expect(connection.disconnectCount).toBe(4);
+    expect(connection.abortCount).toBe(3);
+    expect(connection.status).toBe('closed');
+  });
+});
+
+test('retry when connect failed', async () => {
+  class TestConnection extends AutoReconnectConnection {
+    override retryDelay = 300;
+    connectCount = 0;
+    override async doConnect() {
+      this.connectCount++;
+      if (this.connectCount === 3) {
+        return { hello: 'world' };
+      }
+      throw new Error('not connected, count: ' + this.connectCount);
+    }
+    override doDisconnect() {
+      return Promise.resolve();
+    }
+  }
+
+  const connection = new TestConnection();
+  connection.connect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(1);
+    expect(connection.status).toBe('error');
+    expect(connection.error?.message).toContain('not connected, count: 1');
+  });
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(2);
+    expect(connection.status).toBe('error');
+    expect(connection.error?.message).toBe('not connected, count: 2');
+  });
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(3);
+    expect(connection.status).toBe('connected');
+    expect(connection.error).toBeUndefined();
+  });
+});
+
+test('retry when error', async () => {
+  class TestConnection extends AutoReconnectConnection {
+    override retryDelay = 300;
+    connectCount = 0;
+    disconnectCount = 0;
+    override async doConnect() {
+      this.connectCount++;
+      return {
+        hello: 'world',
+      };
+    }
+    override doDisconnect(conn: any) {
+      this.disconnectCount++;
+      expect(conn).toEqual({
+        hello: 'world',
+      });
+    }
+    triggerError(error: Error) {
+      this.error = error;
+    }
+  }
+
+  const connection = new TestConnection();
+  connection.connect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(1);
+    expect(connection.status).toBe('connected');
+  });
+
+  connection.triggerError(new Error('test error'));
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(1);
+    expect(connection.disconnectCount).toBe(1);
+    expect(connection.status).toBe('error');
+    expect(connection.error?.message).toBe('test error');
+  });
+
+  // waitfor reconnect
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(2);
+    expect(connection.disconnectCount).toBe(1);
+    expect(connection.status).toBe('connected');
+    expect(connection.error).toBeUndefined();
+  });
+});

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) {