diff --git a/packages/common/nbstore/src/connection/__tests__/auto-reconnection.spec.ts b/packages/common/nbstore/src/connection/__tests__/auto-reconnection.spec.ts index f57d3886aa..e4121cf522 100644 --- a/packages/common/nbstore/src/connection/__tests__/auto-reconnection.spec.ts +++ b/packages/common/nbstore/src/connection/__tests__/auto-reconnection.spec.ts @@ -197,4 +197,107 @@ test('retry when error', async () => { 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'); }); diff --git a/packages/common/nbstore/src/connection/connection.ts b/packages/common/nbstore/src/connection/connection.ts index 1416a09b3d..e6e297aedf 100644 --- a/packages/common/nbstore/src/connection/connection.ts +++ b/packages/common/nbstore/src/connection/connection.ts @@ -31,6 +31,7 @@ export abstract class AutoReconnectConnection private _status: ConnectionStatus = 'idle'; private _error: Error | undefined = undefined; retryDelay = 3000; + connectingTimeout = 15000; private refCount = 0; private connectingAbort?: AbortController; private reconnectingAbort?: AbortController; @@ -89,10 +90,17 @@ export abstract class AutoReconnectConnection private innerConnect() { if (this.status !== 'connecting') { this.setStatus('connecting'); - this.connectingAbort = new AbortController(); - const signal = this.connectingAbort.signal; + 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.setStatus('connected'); @@ -106,6 +114,7 @@ export abstract class AutoReconnectConnection }) .catch(error => { if (!signal.aborted) { + clearTimeout(timeout); console.error('failed to connect', error); this.handleError(error as any); } @@ -132,16 +141,23 @@ export abstract class AutoReconnectConnection // 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; - setTimeout(() => { + const timeout = setTimeout(() => { if (!signal.aborted) { this.innerConnect(); } }, this.retryDelay); + signal.addEventListener('abort', () => { + clearTimeout(timeout); + }); } connect() {