feat(handler): re-q unstarted agents on disconnect

Author: blakebyrnes

client/lib/Handler.ts

@@ -9,12 +9,20 @@ import ConnectionToCore from '../connections/ConnectionToCore';
 import ConnectionFactory from '../connections/ConnectionFactory';
 import DisconnectedFromCoreError from '../connections/DisconnectedFromCoreError';
 
-type SettledDispatchesBySessionId = { [sessionId: string]: { args: any; error?: Error } };
-type PendingDispatch = { resolution: Promise<Error | void>; sessionId?: string; args: any };
+type SettledDispatchesBySessionId = {
+  [sessionId: string]: { args: any; error?: Error; retries: number };
+};
+type PendingDispatch = {
+  resolution: Promise<Error | void>;
+  sessionId?: string;
+  args: any;
+  retries: number;
+};
 
 const { log } = Log(module);
 
 export default class Handler {
+  public disconnectedDispatchRetries = 3;
   public defaultAgentOptions: IAgentCreateOptions = {};
   public get coreHosts(): Promise<string[]> {
     return Promise.all(this.connections.map(x => x.hostOrError)).then(x => {
@@ -30,6 +38,8 @@ export default class Handler {
   private readonly connections: ConnectionToCore[] = [];
   private readonly dispatches: PendingDispatch[] = [];
 
+  private isClosing = false;
+
   constructor(...connectionOptions: (IConnectionToCoreOptions | ConnectionToCore)[]) {
     if (!connectionOptions.length) {
       connectionOptions.push({});
@@ -68,14 +78,25 @@ export default class Handler {
     runFn: (agent: Agent, args?: T) => Promise<void>,
     args?: T,
     createAgentOptions?: IAgentCreateOptions,
+    pendingDispatch?: PendingDispatch,
   ): void {
     const options = {
       ...this.defaultAgentOptions,
       ...createAgentOptions,
     };
+
+    const dispatched: PendingDispatch = pendingDispatch ?? { args, resolution: null, retries: 0 };
+
+    // if no available connection, return
     const connection = this.getConnection();
+    if (!connection) {
+      dispatched.resolution = Promise.resolve(
+        new Error("There aren't any connections available to dispatch this agent"),
+      );
+      this.dispatches.push(dispatched);
+      return;
+    }
 
-    const dispatched: PendingDispatch = { args, resolution: null };
     dispatched.resolution = connection
       .useAgent(options, async agent => {
         try {
@@ -85,7 +106,16 @@ export default class Handler {
           await agent.close();
         }
       })
-      .catch((err: Error) => err);
+      .catch(err => {
+        const canRetry =
+          !dispatched.sessionId && dispatched.retries < this.disconnectedDispatchRetries;
+        if (canRetry && !this.isClosing && this.connections.length) {
+          dispatched.retries += 1;
+          return this.dispatchAgent(runFn, args, createAgentOptions, dispatched);
+        }
+
+        return err;
+      });
 
     this.dispatches.push(dispatched);
   }
@@ -144,9 +174,9 @@ export default class Handler {
       this.dispatches.length = 0;
 
       await Promise.all(dispatches.map(x => x.resolution));
-      for (const { sessionId, resolution, args } of dispatches) {
+      for (const { sessionId, resolution, args, retries } of dispatches) {
         const error = <Error>await resolution;
-        result[sessionId] = { args, error };
+        result[sessionId] = { args, error, retries };
       }
 
       await new Promise(setImmediate);
@@ -156,15 +186,21 @@ export default class Handler {
   }
 
   public async close(error?: Error): Promise<void> {
+    if (this.isClosing) return;
+    this.isClosing = true;
     // eslint-disable-next-line promise/no-promise-in-callback
     await Promise.all(this.connections.map(x => x.disconnect(error)));
   }
 
-  private getConnection(): ConnectionToCore {
+  private getAvailableConnections(): ConnectionToCore[] {
     // prefer a connection that can create a session right now
     let connections = this.connections.filter(x => x.canCreateSessionNow());
     if (!connections.length) connections = this.connections.filter(x => !x.isDisconnecting);
-    return pickRandom(connections);
+    return connections;
+  }
+
+  private getConnection(): ConnectionToCore {
+    return pickRandom(this.getAvailableConnections());
   }
 
   private registerUnhandledExceptionHandlers(): void {

core/server/ConnectionToClient.ts

@@ -34,6 +34,7 @@ export default class ConnectionToClient extends TypedEventEmitter<{
 
   private autoShutdownTimer: NodeJS.Timer;
   private readonly sessionIds = new Set<string>();
+  private hasActiveCommand = false;
 
   private clientExposedMethods = new Map<string, keyof this & string>([
     ['Core.connect', 'connect'],
@@ -62,6 +63,7 @@ export default class ConnectionToClient extends TypedEventEmitter<{
 
     let data: any;
     try {
+      this.hasActiveCommand = true;
       data = await this.executeCommand(command, args, meta);
     } catch (error) {
       // if we're closing, don't emit errors
@@ -77,6 +79,8 @@ export default class ConnectionToClient extends TypedEventEmitter<{
       }
       data = this.serializeError(error);
       data.isDisconnecting = this.isClosing || session?.isClosing;
+    } finally {
+      this.hasActiveCommand = false;
     }
 
     const commandId = session?.sessionState?.lastCommand?.id;
@@ -132,7 +136,7 @@ export default class ConnectionToClient extends TypedEventEmitter<{
   }
 
   public isActive(): boolean {
-    return this.sessionIds.size > 0 || this.isPersistent;
+    return this.sessionIds.size > 0 || this.isPersistent || this.hasActiveCommand;
   }
 
   ///////  SESSION /////////////////////////////////////////////////////////////////////////////////////////////////////

full-client/test/handler.test.ts

@@ -1,6 +1,6 @@
 import { Helpers } from '@secret-agent/testing';
 import { ITestKoaServer } from '@secret-agent/testing/helpers';
-import Core, { Session, CoreProcess } from '@secret-agent/core/index';
+import Core, { CoreProcess, Session } from '@secret-agent/core/index';
 import DisconnectedFromCoreError from '@secret-agent/client/connections/DisconnectedFromCoreError';
 import { Agent, RemoteConnectionToCore } from '@secret-agent/client/index';
 import { createPromise } from '@secret-agent/commons/utils';
@@ -12,6 +12,7 @@ beforeAll(async () => {
   await Core.start();
   koaServer = await Helpers.runKoaServer(true);
 });
+afterEach(Helpers.afterEach);
 afterAll(Helpers.afterAll);
 
 describe('Full client Handler', () => {
@@ -68,6 +69,8 @@ describe('Full client Handler', () => {
     await expect(agent2.sessionId).resolves.toBeTruthy();
     const agent3 = handler.createAgent();
 
+    Helpers.needsClosing.push(agent2);
+
     async function isAgent3Available(millis = 100): Promise<boolean> {
       const result = await Promise.race([
         agent3,
@@ -81,6 +84,7 @@ describe('Full client Handler', () => {
     await agent1.close();
 
     await expect(isAgent3Available(5e3)).resolves.toBe(true);
+    await (await agent3).close();
   });
 });
 
@@ -374,4 +378,63 @@ describe('connectionToCore', () => {
 
     await handler.close();
   });
+
+  it('can re-queue dispatched agents that never started', async () => {
+    const coreHost = await CoreProcess.spawn({});
+    Helpers.onClose(() => CoreProcess.kill('SIGINT'));
+    const connection1 = new RemoteConnectionToCore({
+      maxConcurrency: 1,
+      host: coreHost,
+    });
+    await connection1.connect();
+
+    const handler = new Handler(connection1);
+    Helpers.needsClosing.push(handler);
+
+    const waitForGoto = createPromise();
+    const dispatchErrorPromise = createPromise<Error>();
+    handler.dispatchAgent(async agent => {
+      try {
+        await agent.goto(koaServer.baseUrl);
+        // create a command we can disconnect from (don't await yet)
+        const promise = agent.waitForMillis(5e3);
+        await new Promise(resolve => setTimeout(resolve, 50));
+        waitForGoto.resolve();
+        await promise;
+      } catch (error) {
+        dispatchErrorPromise.resolve(error);
+        throw error;
+      }
+    });
+
+    let counter = 0;
+    const incr = async agent => {
+      await agent.goto(koaServer.baseUrl);
+      counter += 1;
+    };
+    handler.dispatchAgent(incr);
+    handler.dispatchAgent(incr);
+
+    // first 2 will be queued against the first connection
+    const coreHost2 = await Core.server.address;
+    await handler.addConnectionToCore({ maxConcurrency: 2, host: coreHost2 });
+    handler.dispatchAgent(incr);
+    handler.dispatchAgent(incr);
+    await waitForGoto.promise;
+
+    // disconnect the first connection. the first two handlers should get re-queued
+    await connection1.disconnect();
+    await new Promise(setImmediate);
+
+    // should have an error thrown if it actually the process. this one should NOT get re-queued
+    await expect(dispatchErrorPromise).resolves.toBeTruthy();
+    const dispatchError = await dispatchErrorPromise;
+    expect(dispatchError).toBeInstanceOf(DisconnectedFromCoreError);
+
+    const allDispatches = await handler.waitForAllDispatchesSettled();
+
+    expect(counter).toBe(4);
+    expect(Object.keys(allDispatches)).toHaveLength(5);
+    expect(Object.values(allDispatches).filter(x => !!x.error)).toHaveLength(1);
+  });
 });