snapshot of https://github.com/toeverything/AFFiNE.git (canary, 2026-09-16; blobs > 131072B omitted)

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

@@ -0,0 +1,318 @@
+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;
+    retryDelayFor(retryCount: number) {
+      return this.getRetryDelay(retryCount);
+    }
+    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();
+    }
+    triggerError(error: Error) {
+      this.error = error;
+    }
+  }
+
+  const connection = new TestConnection();
+  expect([0, 1, 2, 8].map(count => connection.retryDelayFor(count))).toEqual([
+    300, 600, 1200, 60000,
+  ]);
+  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();
+  });
+
+  connection.triggerError(new Error('disconnected'));
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(4);
+    expect(connection.status).toBe('connected');
+  });
+});
+
+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();
+  });
+
+  // do not reconnect if the connection is closed
+  connection.disconnect();
+  connection.triggerError(new Error('test error2'));
+  await new Promise(resolve => setTimeout(resolve, 1000));
+  expect(connection.connectCount).toBe(2);
+  expect(connection.status).toBe('closed');
+});
+
+test('connecting timeout', async () => {
+  class TestConnection extends AutoReconnectConnection {
+    override connectingTimeout = 150;
+    override retryDelay = 150;
+    connectCount = 0;
+    disconnectCount = 0;
+    override async doConnect() {
+      this.connectCount++;
+      if (this.connectCount === 3) {
+        return { foo: 'bar' };
+      }
+      await new Promise(resolve => setTimeout(resolve, 300));
+      throw new Error('not connected, count: ' + this.connectCount);
+    }
+    override doDisconnect(conn: any) {
+      this.disconnectCount++;
+      expect(conn).toEqual({
+        foo: 'bar',
+      });
+    }
+    triggerError(error: Error) {
+      this.error = error;
+    }
+  }
+
+  const connection = new TestConnection();
+  connection.connect();
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(1);
+    expect(connection.disconnectCount).toBe(0);
+    expect(connection.status).toBe('error');
+    expect(connection.error?.message).toBe('connecting timeout');
+  });
+
+  // wait for reconnect
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(2);
+    expect(connection.disconnectCount).toBe(0);
+    expect(connection.status).toBe('connecting');
+    expect(connection.error?.message).toBe('connecting timeout');
+  });
+
+  // trigger error while connecting
+  connection.triggerError(new Error('test error2'));
+  connection.triggerError(new Error('test error2'));
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(2);
+    expect(connection.disconnectCount).toBe(0);
+    expect(connection.status).toBe('error');
+    expect(connection.error?.message).toBe('test error2');
+  });
+
+  // wait for reconnect
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(3);
+    expect(connection.disconnectCount).toBe(0);
+    expect(connection.status).toBe('connected');
+    expect(connection.error).toBeUndefined();
+  });
+
+  // trigger error after connected
+  connection.triggerError(new Error('test error3'));
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(3);
+    expect(connection.disconnectCount).toBe(1); // previous connect is disconnected
+    expect(connection.status).toBe('error');
+    expect(connection.error?.message).toBe('test error3');
+  });
+
+  // reconnect and timeout again
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(4);
+  });
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(4);
+    expect(connection.disconnectCount).toBe(1);
+    expect(connection.status).toBe('error');
+    expect(connection.error?.message).toBe('connecting timeout');
+  });
+
+  await vitest.waitFor(() => {
+    expect(connection.connectCount).toBe(5);
+  });
+
+  connection.disconnect();
+
+  await new Promise(resolve => setTimeout(resolve, 1000));
+
+  // no reconnect after disconnect
+  expect(connection.connectCount).toBe(5);
+  expect(connection.status).toBe('closed');
+});

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

@@ -0,0 +1,257 @@
+import EventEmitter2 from 'eventemitter2';
+
+import { MANUALLY_STOP } from '../utils/throw-if-aborted';
+
+export type ConnectionStatus =
+  | 'idle'
+  | 'connecting'
+  | 'connected'
+  | 'error'
+  | 'closed';
+
+export interface Connection<T = any> {
+  readonly status: ConnectionStatus;
+  readonly error?: Error;
+  readonly inner: T;
+  connect(): void;
+  disconnect(): void;
+  waitForConnected(signal?: AbortSignal): Promise<void>;
+  onStatusChanged(
+    cb: (status: ConnectionStatus, error?: Error) => void
+  ): () => void;
+}
+
+export abstract class AutoReconnectConnection<
+  T = any,
+> implements Connection<T> {
+  private readonly event = new EventEmitter2({
+    maxListeners: 100,
+  });
+  private _inner: T | undefined = undefined;
+  private _status: ConnectionStatus = 'idle';
+  private _error: Error | undefined = undefined;
+  retryDelay = 3000;
+  maxRetryDelay = 60000;
+  connectingTimeout = 15000;
+  private retryCount = 0;
+  private refCount = 0;
+  private connectingAbort?: AbortController;
+  private reconnectingAbort?: AbortController;
+
+  constructor() {}
+
+  get shareId(): string | undefined {
+    return undefined;
+  }
+
+  get maybeConnection() {
+    return this._inner;
+  }
+
+  get inner(): T {
+    if (this._inner === undefined) {
+      throw new Error(
+        `Connection ${this.constructor.name} has not been established.`
+      );
+    }
+
+    return this._inner;
+  }
+
+  private set inner(inner: T | undefined) {
+    this._inner = inner;
+  }
+
+  get status() {
+    return this._status;
+  }
+
+  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;
+    // we only clear-up error when status is connected
+    if (error || status === 'connected') {
+      this._error = error;
+    }
+    if (shouldEmit) {
+      this.emitStatusChanged(status, this._error);
+    }
+  }
+
+  protected abstract doConnect(signal?: AbortSignal): Promise<T>;
+  protected abstract doDisconnect(conn: T): void;
+
+  private innerConnect() {
+    if (this.status !== 'connecting') {
+      this.setStatus('connecting');
+      const connectingAbort = new AbortController();
+      this.connectingAbort = connectingAbort;
+      const signal = connectingAbort.signal;
+      const timeout = setTimeout(() => {
+        if (!signal.aborted) {
+          this.handleError(new Error('connecting timeout'));
+        }
+      }, this.connectingTimeout);
+      this.doConnect(signal)
+        .then(value => {
+          clearTimeout(timeout);
+          if (!signal.aborted) {
+            this._inner = value;
+            this.retryCount = 0;
+            this.setStatus('connected');
+          } else {
+            try {
+              this.doDisconnect(value);
+            } catch (error) {
+              console.error('failed to disconnect', error);
+            }
+          }
+        })
+        .catch(error => {
+          if (!signal.aborted) {
+            clearTimeout(timeout);
+            console.error('failed to connect', error);
+            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();
+    // if the connection is closed, do not reconnect
+    if (this.status === 'closed') {
+      return;
+    }
+    this.setStatus('error', reason);
+    // reconnect
+
+    this.reconnectingAbort = new AbortController();
+    const signal = this.reconnectingAbort.signal;
+    const retryDelay = this.getRetryDelay(this.retryCount++);
+    const timeout = setTimeout(() => {
+      if (!signal.aborted) {
+        this.innerConnect();
+      }
+    }, retryDelay);
+    signal.addEventListener('abort', () => {
+      clearTimeout(timeout);
+    });
+  }
+
+  protected getRetryDelay(retryCount: number) {
+    return Math.min(this.retryDelay * 2 ** retryCount, this.maxRetryDelay);
+  }
+
+  connect() {
+    this.refCount++;
+    if (this.refCount === 1) {
+      this.innerConnect();
+    }
+  }
+
+  disconnect(force?: boolean) {
+    if (force) {
+      this.refCount = 0;
+    } else {
+      this.refCount = Math.max(this.refCount - 1, 0);
+    }
+    if (this.refCount === 0) {
+      this.innerDisconnect();
+      this.retryCount = 0;
+      this.setStatus('closed');
+    }
+  }
+
+  waitForConnected(signal?: AbortSignal) {
+    return new Promise<void>((resolve, reject) => {
+      if (this.status === 'connected') {
+        resolve();
+        return;
+      }
+
+      if (signal?.aborted) {
+        reject(signal.reason);
+        return;
+      }
+
+      const off = this.onStatusChanged(status => {
+        if (status === 'connected') {
+          resolve();
+          cleanup();
+        }
+      });
+
+      const onAbort = () => {
+        reject(signal?.reason);
+        cleanup();
+      };
+      const cleanup = () => {
+        off();
+        signal?.removeEventListener('abort', onAbort);
+      };
+      signal?.addEventListener('abort', onAbort, { once: true });
+    });
+  }
+
+  onStatusChanged(
+    cb: (status: ConnectionStatus, error?: Error) => void
+  ): () => void {
+    this.event.on('statusChanged', cb);
+    return () => {
+      this.event.off('statusChanged', cb);
+    };
+  }
+
+  private readonly emitStatusChanged = (
+    status: ConnectionStatus,
+    error?: Error
+  ) => {
+    this.event.emit('statusChanged', status, error);
+  };
+}
+
+export class DummyConnection implements Connection<undefined> {
+  readonly status: ConnectionStatus = 'connected';
+  readonly inner: undefined;
+
+  connect(): void {
+    return;
+  }
+  disconnect(): void {
+    return;
+  }
+  waitForConnected(_signal?: AbortSignal): Promise<void> {
+    return Promise.resolve();
+  }
+  onStatusChanged(
+    _cb: (status: ConnectionStatus, error?: Error) => void
+  ): () => void {
+    return () => {};
+  }
+}

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

@@ -0,0 +1,2 @@
+export * from './connection';
+export * from './shared-connection';

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

@@ -0,0 +1,20 @@
+import type { AutoReconnectConnection } from './connection';
+
+const CONNECTIONS: Map<string, AutoReconnectConnection<any>> = new Map();
+export function share<T extends AutoReconnectConnection<any>>(conn: T): T {
+  if (!conn.shareId) {
+    throw new Error(
+      `Connection ${conn.constructor.name} is not shareable.\nIf you want to make it shareable, please override [shareId].`
+    );
+  }
+
+  const existing = CONNECTIONS.get(conn.shareId);
+
+  if (existing) {
+    return existing as T;
+  }
+
+  CONNECTIONS.set(conn.shareId, conn);
+
+  return conn;
+}