fix(core): full close down of on premature exit

Author: blakebyrnes

client/connections/ConnectionToCore.ts

@@ -6,6 +6,7 @@ import IResolvablePromise from '@secret-agent/core-interfaces/IResolvablePromise
 import Log from '@secret-agent/commons/Logger';
 import ICreateSessionOptions from '@secret-agent/core-interfaces/ICreateSessionOptions';
 import ISessionMeta from '@secret-agent/core-interfaces/ISessionMeta';
+import { CanceledPromiseError } from '@secret-agent/commons/interfaces/IPendingWaitEvent';
 import IConnectionToCoreOptions from '../interfaces/IConnectionToCoreOptions';
 import CoreCommandQueue from '../lib/CoreCommandQueue';
 import CoreSession from '../lib/CoreSession';
@@ -17,11 +18,11 @@ const { log } = Log(module);
 
 export default abstract class ConnectionToCore {
   public readonly commandQueue: CoreCommandQueue;
+  public readonly hostOrError: Promise<string | Error>;
   public options: IConnectionToCoreOptions;
 
-  public hostOrError: Promise<string | Error>;
-
   private connectPromise: Promise<Error | null>;
+  private isClosing = false;
 
   private coreSessions: CoreSessions;
   private readonly pendingRequestsById = new Map<string, IResolvablePromiseWithId>();
@@ -34,31 +35,57 @@ export default abstract class ConnectionToCore {
       this.options.maxConcurrency,
       this.options.agentTimeoutMillis,
     );
+
+    if (this.options.host) {
+      this.hostOrError = Promise.resolve(this.options.host)
+        .then(x => {
+          if (!x.includes('://')) {
+            return `ws://${x}`;
+          }
+          return x;
+        })
+        .catch(err => err);
+    } else {
+      this.hostOrError = Promise.resolve(new Error('No host provided'));
+    }
   }
 
-  protected abstract internalSendRequest(payload: ICoreRequestPayload): void | Promise<void>;
+  protected abstract internalSendRequest(payload: ICoreRequestPayload): Promise<void>;
+  protected abstract createConnection(): Promise<Error | null>;
+  protected abstract destroyConnection(): Promise<any>;
 
   public connect(): Promise<Error | null> {
-    if (this.connectPromise) return this.connectPromise;
-
-    const { promise, id } = this.createPendingResult();
-    this.connectPromise = promise.then(x => this.onConnected(x.data)).catch(err => err);
-    this.internalSendRequest({
-      command: 'connect',
-      args: [this.options],
-      messageId: id,
-    });
+    this.connectPromise ??= this.createConnection()
+      .then(err => {
+        if (err) throw err;
+        return this.internalSendRequestAndWait({
+          command: 'connect',
+          args: [this.options],
+        });
+      })
+      .then(result => this.onConnected(result.data))
+      .catch(err => err);
 
     return this.connectPromise;
   }
 
   public async disconnect(fatalError?: Error): Promise<void> {
-    this.commandQueue.clearPending();
-    this.coreSessions.close();
-    if (this.connectPromise || fatalError) {
-      await this.commandQueue.run('disconnect', [fatalError]);
-      this.connectPromise = null;
+    if (this.isClosing) return;
+    this.isClosing = true;
+    const logid = log.stats('ConnectionToCore.Disconnecting');
+
+    this.cancelPendingRequests();
+    if (this.connectPromise) {
+      await this.internalSendRequestAndWait({
+        command: 'disconnect',
+        args: [fatalError],
+      });
     }
+    await this.destroyConnection();
+    log.stats('RemoteConnectionToCore.Disconnected', {
+      parentLogId: logid,
+      sessionId: null,
+    });
   }
 
   ///////  PIPE FUNCTIONS  /////////////////////////////////////////////////////////////////////////////////////////////
@@ -69,12 +96,7 @@ export default abstract class ConnectionToCore {
     const result = await this.connect();
     if (result) throw result;
 
-    const { promise, id } = this.createPendingResult();
-    await this.internalSendRequest({
-      messageId: id,
-      ...payload,
-    });
-    return promise;
+    return this.internalSendRequestAndWait(payload);
   }
 
   public onMessage(payload: ICoreResponsePayload | ICoreEventPayload): void {
@@ -122,6 +144,25 @@ export default abstract class ConnectionToCore {
     await this.commandQueue.run('logUnhandledError', error);
   }
 
+  protected async internalSendRequestAndWait(
+    payload: Omit<ICoreRequestPayload, 'messageId'>,
+  ): Promise<ICoreResponsePayload> {
+    const { promise, id } = this.createPendingResult();
+    try {
+      await this.internalSendRequest({
+        messageId: id,
+        ...payload,
+      });
+    } catch (error) {
+      if (error instanceof CanceledPromiseError) {
+        this.pendingRequestsById.delete(id);
+        return;
+      }
+      throw error;
+    }
+    return promise;
+  }
+
   protected onEvent(payload: ICoreEventPayload): void {
     const { meta, listenerId, eventArgs } = payload as ICoreEventPayload;
     const session = this.getSession(meta.sessionId);
@@ -136,13 +177,22 @@ export default abstract class ConnectionToCore {
     if (message.isError) {
       const error = new Error(message.data?.message);
       Object.assign(error, message.data);
-      error.stack += `\n${'------CONNECTION'.padEnd(50, '-')}\n${pending.stack}`;
-      pending.reject(error);
+      this.rejectPendingRequest(pending, error);
     } else {
       pending.resolve({ data: message.data, commandId: message.commandId });
     }
   }
 
+  protected cancelPendingRequests(): void {
+    this.commandQueue.clearPending();
+    this.coreSessions.close();
+    const pending = [...this.pendingRequestsById.values()];
+    this.pendingRequestsById.clear();
+    for (const entry of pending) {
+      this.rejectPendingRequest(entry, new CanceledPromiseError('Disconnecting from Core'));
+    }
+  }
+
   private createPendingResult(): IResolvablePromiseWithId {
     const resolvablePromise = createPromise<ICoreResponsePayload>() as IResolvablePromiseWithId;
     this.lastId += 1;
@@ -152,9 +202,15 @@ export default abstract class ConnectionToCore {
     return this.pendingRequestsById.get(id);
   }
 
+  private rejectPendingRequest(pending: IResolvablePromiseWithId, error: Error): void {
+    error.stack += `\n${'------CONNECTION'.padEnd(50, '-')}\n${pending.stack}`;
+    pending.reject(error);
+  }
+
   private onConnected(
     connectionParams: { maxConcurrency?: number; browserEmulatorIds?: string[] } = {},
   ): void {
+    this.isClosing = false;
     const { maxConcurrency, browserEmulatorIds } = connectionParams;
     if (!this.options.maxConcurrency || maxConcurrency < this.options.maxConcurrency) {
       log.info('Overriding max concurrency with Core value', {

client/connections/RemoteConnectionToCore.ts

@@ -1,76 +1,93 @@
 import ICoreRequestPayload from '@secret-agent/core-interfaces/ICoreRequestPayload';
 import WebSocket from 'ws';
 import TypeSerializer from '@secret-agent/commons/TypeSerializer';
+import { createPromise } from '@secret-agent/commons/utils';
+import IResolvablePromise from '@secret-agent/core-interfaces/IResolvablePromise';
+import { CanceledPromiseError } from '@secret-agent/commons/interfaces/IPendingWaitEvent';
 import ConnectionToCore from './ConnectionToCore';
 import IConnectionToCoreOptions from '../interfaces/IConnectionToCoreOptions';
 
 export default class RemoteConnectionToCore extends ConnectionToCore {
-  private wsConnectPromise: Promise<any>;
-  private webSocket: WebSocket;
+  private webSocketOrError: IResolvablePromise<WebSocket | Error>;
 
   constructor(options: IConnectionToCoreOptions) {
+    if (!options.host) throw new Error('A remote connection to core needs a host parameter!');
     super(options);
-    const host = options.host;
-    if (!host) throw new Error('A remote connection to core needs a host parameter!');
-
-    this.hostOrError = Promise.resolve(host)
-      .then(x => {
-        if (!x.includes('://')) {
-          return `ws://${x}`;
-        }
-        return x;
-      })
-      .catch(err => err);
+    this.disconnect = this.disconnect.bind(this);
   }
 
-  public internalSendRequest(payload: ICoreRequestPayload): Promise<void> {
+  public async internalSendRequest(payload: ICoreRequestPayload): Promise<void> {
+    if (!this.webSocketOrError) return;
     const message = TypeSerializer.stringify(payload);
+
+    const webSocket = await this.getWebsocket();
+
+    if (webSocket?.readyState !== WebSocket.OPEN) {
+      throw new CanceledPromiseError('Websocket was not open');
+    }
+
     return new Promise((resolve, reject) =>
-      this.webSocket.send(message, err => {
+      webSocket.send(message, err => {
         if (err) reject(err);
         else resolve();
       }),
     );
   }
 
-  public async disconnect(): Promise<void> {
-    if (this.wsConnectPromise && this.webSocket && this.webSocket.readyState === WebSocket.OPEN) {
-      this.wsConnectPromise = null;
-      await super.disconnect();
+  protected async destroyConnection(): Promise<any> {
+    const webSocket = await this.getWebsocket(false);
+    if (webSocket?.readyState === WebSocket.OPEN) {
       try {
-        this.webSocket.terminate();
+        webSocket.off('close', this.disconnect);
+        webSocket.terminate();
       } catch (_) {
         // ignore errors terminating
       }
     }
   }
 
-  public connect(): Promise<Error | null> {
-    if (!this.wsConnectPromise) {
-      this.wsConnectPromise = this.wsConnect().catch(err => err);
-    }
+  protected async createConnection(): Promise<Error | null> {
+    // do this first to see if we can resolve the host
+    const hostOrError = await this.hostOrError;
+    if (hostOrError instanceof Error) return hostOrError;
 
-    return this.wsConnectPromise;
+    if (!this.webSocketOrError) {
+      this.webSocketOrError = connectToWebsocketHost(hostOrError);
+      try {
+        const webSocket = await this.getWebsocket();
+        webSocket.once('close', this.disconnect);
+        webSocket.on('message', message => {
+          const payload = TypeSerializer.parse(message.toString());
+          this.onMessage(payload);
+        });
+      } catch (error) {
+        return error;
+      }
+    }
   }
 
-  private async wsConnect(): Promise<void> {
-    const hostOrError = await this.hostOrError;
-    if (hostOrError instanceof Error) throw hostOrError;
-
-    this.webSocket = new WebSocket(hostOrError);
-    await new Promise<void>((resolve, reject) => {
-      this.webSocket.on('error', reject);
-      this.webSocket.once('open', () => {
-        this.webSocket.off('error', reject);
-        resolve();
-      });
-    });
-    this.webSocket.once('close', this.disconnect.bind(this));
-    this.webSocket.on('message', message => {
-      const payload = TypeSerializer.parse(message.toString());
-      this.onMessage(payload);
-    });
+  private async getWebsocket(throwIfError = true): Promise<WebSocket> {
+    if (!this.webSocketOrError) return null;
+    const webSocketOrError = await this.webSocketOrError.promise;
+    if (webSocketOrError instanceof Error) {
+      if (throwIfError) throw webSocketOrError;
+      return null;
+    }
+    return webSocketOrError;
+  }
+}
 
-    await super.connect();
+function connectToWebsocketHost(host: string): IResolvablePromise<WebSocket | Error> {
+  const resolvable = createPromise<WebSocket | Error>(30e3);
+  const webSocket = new WebSocket(host);
+  function onError(error: Error): void {
+    if (error instanceof Error) resolvable.resolve(error);
+    else resolvable.resolve(new Error(`Error connecting to Websocket host -> ${error}`));
   }
+  webSocket.once('error', onError);
+  webSocket.once('open', () => {
+    webSocket.off('error', onError);
+    resolvable.resolve(webSocket);
+  });
+  return resolvable;
 }

client/lib/Agent.ts

@@ -17,7 +17,6 @@ import Request from 'awaited-dom/impl/official-klasses/Request';
 import IWaitForOptions from '@secret-agent/core-interfaces/IWaitForOptions';
 import { IElementIsolate } from 'awaited-dom/base/interfaces/isolate';
 import CSSStyleDeclaration from 'awaited-dom/impl/official-klasses/CSSStyleDeclaration';
-import { CanceledPromiseError } from '@secret-agent/commons/interfaces/IPendingWaitEvent';
 import IAgentMeta from '@secret-agent/core-interfaces/IAgentMeta';
 import IScreenshotOptions from '@secret-agent/core-interfaces/IScreenshotOptions';
 import WebsocketResource from './WebsocketResource';
@@ -380,10 +379,7 @@ class SessionConnection {
       connectionToCore ?? { isPersistent: false },
     );
 
-    this._coreSession = connection.createSession(options).catch(err => {
-      if (err instanceof CanceledPromiseError) return null;
-      return err;
-    });
+    this._coreSession = connection.createSession(options).catch(err => err);
 
     const defaultShowReplay = Boolean(JSON.parse(process.env.SA_SHOW_REPLAY ?? 'true'));
 

client/lib/CoreCommandQueue.ts

@@ -1,6 +1,7 @@
 import { createPromise } from '@secret-agent/commons/utils';
 import ISessionMeta from '@secret-agent/core-interfaces/ISessionMeta';
 import Log from '@secret-agent/commons/Logger';
+import { CanceledPromiseError } from '@secret-agent/commons/interfaces/IPendingWaitEvent';
 import ConnectionToCore from '../connections/ConnectionToCore';
 
 const { log } = Log(module);
@@ -49,7 +50,12 @@ export default class CoreCommandQueue {
   }
 
   public clearPending(): void {
-    this.items.length = 0;
+    while (this.items.length) {
+      const next = this.items.shift();
+      const cancel = new CanceledPromiseError(`Canceling pending ${next.command} command`);
+      cancel.stack = next.stack;
+      next.reject();
+    }
   }
 
   // PRIVATE

client/lib/Handler.ts

@@ -1,10 +1,10 @@
 import { createPromise, pickRandom } from '@secret-agent/commons/utils';
+import ShutdownHandler from '@secret-agent/commons/ShutdownHandler';
 import IAgentCreateOptions from '../interfaces/IAgentCreateOptions';
 import IConnectionToCoreOptions from '../interfaces/IConnectionToCoreOptions';
 import Agent from './Agent';
 import ConnectionToCore from '../connections/ConnectionToCore';
 import ConnectionFactory from '../connections/ConnectionFactory';
-import Signals = NodeJS.Signals;
 
 export default class Handler {
   public defaultAgentOptions: IAgentCreateOptions = {};
@@ -21,7 +21,7 @@ export default class Handler {
       this.connections.push(connection);
     }
 
-    this.registerShutdownHandlers();
+    ShutdownHandler.register(() => this.close());
     this.registerUnhandledExceptionHandlers();
   }
 
@@ -93,13 +93,7 @@ export default class Handler {
     await Promise.all(this.connections.map(x => x.disconnect(error)));
   }
 
-  public registerShutdownHandlers(): void {
-    for (const signal of ['exit', 'SIGTERM', 'SIGINT', 'SIGQUIT']) {
-      process.once(signal as Signals, this.close.bind(this));
-    }
-  }
-
-  public registerUnhandledExceptionHandlers(): void {
+  private registerUnhandledExceptionHandlers(): void {
     process.on('uncaughtExceptionMonitor', this.close.bind(this));
     process.on('unhandledRejection', this.logUnhandledError.bind(this));
   }