fix(client): handle ws closing during shutdown

Author: blakebyrnes

client/connections/ConnectionToCore.ts

@@ -30,6 +30,7 @@ export default abstract class ConnectionToCore extends TypedEventEmitter<{
   public readonly hostOrError: Promise<string | Error>;
   public options: IConnectionToCoreOptions;
   public isDisconnecting = false;
+  public isConnectionTerminated = false;
 
   protected resolvedHost: string;
 
@@ -287,11 +288,30 @@ export default abstract class ConnectionToCore extends TypedEventEmitter<{
     session?.onEvent(meta, listenerId, eventArgs);
   }
 
+  protected async onConnectionTerminated(): Promise<void> {
+    if (this.isConnectionTerminated) return;
+
+    this.isConnectionTerminated = true;
+    await this.internalDisconnect();
+    if (this.connectRequestId) {
+      this.onResponse(this.connectRequestId, {
+        data: new DisconnectedFromCoreError(this.resolvedHost),
+      });
+    }
+    if (this.disconnectRequestId) {
+      this.onResponse(this.disconnectRequestId, {
+        data: new DisconnectedFromCoreError(this.resolvedHost),
+      });
+    }
+  }
+
   protected onResponse(id: string, message: ICoreResponsePayload): void {
     const pending = this.pendingRequestsById.get(id);
     if (!pending) return;
     this.pendingRequestsById.delete(id);
     const isInternalRequest = this.connectRequestId === id || this.disconnectRequestId === id;
+    if (this.disconnectRequestId === id) this.disconnectRequestId = null;
+    if (this.connectRequestId === id) this.connectRequestId = null;
 
     if (message.data instanceof Error) {
       let responseError = message.data;

client/connections/RemoteConnectionToCore.ts

@@ -43,7 +43,7 @@ export default class RemoteConnectionToCore extends ConnectionToCore {
     const webSocket = await this.getWebsocket(false);
     if (webSocket?.readyState === WebSocket.OPEN) {
       try {
-        webSocket.off('close', this.internalDisconnect);
+        webSocket.off('close', this.onConnectionTerminated);
         webSocket.off('error', this.internalDisconnect);
         webSocket.terminate();
       } catch (_) {
@@ -61,7 +61,7 @@ export default class RemoteConnectionToCore extends ConnectionToCore {
       this.webSocketOrError = connectToWebsocketHost(hostOrError);
       try {
         const webSocket = await this.getWebsocket();
-        webSocket.once('close', this.internalDisconnect);
+        webSocket.once('close', this.onConnectionTerminated);
         webSocket.once('error', this.internalDisconnect);
         webSocket.on('message', message => {
           const payload = TypeSerializer.parse(message.toString(), 'REMOTE CORE');

commons/ShutdownHandler.ts

@@ -21,6 +21,11 @@ export default class ShutdownHandler {
     this.onShutdownFns.push({ fn: onShutdownFn, callsite });
   }
 
+  public static async shutdown(exitProcess: boolean): Promise<void> {
+    await this.onSignal('exit');
+    if (exitProcess) process.exit(0);
+  }
+
   private static registerSignals(): void {
     if (!this.isRegistered) {
       this.isRegistered = true;
@@ -39,7 +44,9 @@ export default class ShutdownHandler {
       sessionId: null,
     });
 
-    for (const entry of this.onShutdownFns) {
+    while (this.onShutdownFns.length) {
+      const entry = this.onShutdownFns.shift();
+
       log.stats('ShutdownHandler.execute', {
         signal,
         fn: entry.fn.toString(),

core/index.ts

@@ -3,9 +3,9 @@ import { LocationTrigger } from '@secret-agent/interfaces/Location';
 import Log, { hasBeenLoggedSymbol } from '@secret-agent/commons/Logger';
 import Resolvable from '@secret-agent/commons/Resolvable';
 import {
-  IHumanEmulatorClass,
   IBrowserEmulatorClass,
   ICorePluginClass,
+  IHumanEmulatorClass,
 } from '@secret-agent/interfaces/ICorePlugin';
 import { PluginTypes } from '@secret-agent/interfaces/IPluginTypes';
 import DefaultBrowserEmulator from '@secret-agent/default-browser-emulator';
@@ -20,7 +20,6 @@ import CoreProcess from './lib/CoreProcess';
 import Session from './lib/Session';
 import Tab from './lib/Tab';
 import GlobalPool from './lib/GlobalPool';
-import Signals = NodeJS.Signals;
 
 const { log } = Log(module);
 
@@ -94,6 +93,7 @@ export default class Core {
     this.isStarting = true;
     if (isExplicitlyStarted) this.wasManuallyStarted = true;
 
+    this.registerSignals();
     const { localProxyPortStart, sessionsDir, maxConcurrentAgentsCount } = options;
 
     if (maxConcurrentAgentsCount !== undefined)
@@ -175,21 +175,18 @@ export default class Core {
       });
     });
   }
-}
 
-['exit', 'SIGTERM', 'SIGINT', 'SIGQUIT'].forEach(name => {
-  process.once(name as Signals, async () => {
-    await Core.shutdown();
-    process.exit(0);
-  });
-});
-
-if (process.env.NODE_ENV !== 'test') {
-  process.on('uncaughtExceptionMonitor', async (error: Error) => {
-    await Core.logUnhandledError(error, true);
-    await Core.shutdown();
-  });
-  process.on('unhandledRejection', async (error: Error) => {
-    await Core.logUnhandledError(error, false);
-  });
+  private static registerSignals() {
+    ShutdownHandler.register(() => Core.shutdown());
+
+    if (process.env.NODE_ENV !== 'test') {
+      process.on('uncaughtExceptionMonitor', async (error: Error) => {
+        await Core.logUnhandledError(error, true);
+        await Core.shutdown();
+      });
+      process.on('unhandledRejection', async (error: Error) => {
+        await Core.logUnhandledError(error, false);
+      });
+    }
+  }
 }

core/start.ts

@@ -1,6 +1,7 @@
 import '@secret-agent/commons/SourceMapSupport';
 import ICoreConfigureOptions from '@secret-agent/interfaces/ICoreConfigureOptions';
 import Log from '@secret-agent/commons/Logger';
+import ShutdownHandler from '@secret-agent/commons/ShutdownHandler';
 import Core from '.';
 
 const { log } = Log(module);
@@ -11,7 +12,7 @@ const { log } = Log(module);
 
   Core.onShutdown = () => {
     log.stats('Exiting Core Process');
-    process.exit();
+    ShutdownHandler.shutdown(true);
   };
   await Core.start(startOptions, !process.env.SA_TEMPORARY_CORE);
 })().catch(error => {