Merge pull request #1107 from OpenFn/async-child-process-kill

Async child process kill

Author: josephjclark

.changeset/late-sloths-agree.md

@@ -0,0 +1,5 @@
+---
+'@openfn/engine-multi': patch
+---
+
+Wait for child processes to exit properly when shutting down (shouldn't affect prod environments)

packages/engine-multi/src/engine.ts

@@ -229,7 +229,7 @@ const createEngine = async (
     // How does this work if deferred?
   };
 
-  const destroy = (instant?: boolean) => closeWorkers(instant);
+  const destroy = async (instant?: boolean) => closeWorkers(instant);
 
   return Object.assign(engine, {
     options,

packages/engine-multi/src/types.ts

@@ -57,7 +57,7 @@ export type ExecutionContextOptions = ExecuteOptions & EngineOptions;
 
 export interface EngineAPI extends EventEmitter {
   callWorker: CallWorker;
-  closeWorkers: (instant?: boolean) => void;
+  closeWorkers: (instant?: boolean) => Promise<void>;
 }
 
 export interface RuntimeEngine {
@@ -74,7 +74,7 @@ export interface RuntimeEngine {
     options?: Partial<EngineOptions>
   ): Pick<EventEmitter, 'on' | 'off' | 'once'>;
 
-  destroy(): void;
+  destroy(instant?: boolean): Promise<void>;
 
   on: (evt: string, fn: (...args: any[]) => void) => void;
 }

packages/engine-multi/src/worker/pool.ts

@@ -265,18 +265,53 @@ function createPool(script: string, options: PoolOptions = {}, logger: Logger) {
     }
   };
 
-  const destroy = (immediate = false) => {
+  const waitForWorkerExit = (
+    worker: ChildProcess,
+    forceKillTimeout = 5000
+  ): Promise<void> => {
+    return new Promise((resolve) => {
+      if (!worker || worker.killed || !worker.connected) {
+        resolve();
+        return;
+      }
+
+      const timeout = setTimeout(() => {
+        logger.debug('pool: force killing worker', worker.pid);
+        worker.kill('SIGKILL');
+        resolve();
+      }, forceKillTimeout);
+
+      worker.once('exit', () => {
+        clearTimeout(timeout);
+        resolve();
+      });
+
+      worker.kill();
+    });
+  };
+
+  const destroy = async (immediate = false): Promise<void> => {
     destroyed = true;
 
+    const killPromises: Promise<void>[] = [];
+
     // Drain the pool
     while (pool.length) {
-      killWorker(pool.pop()!);
+      const worker = pool.pop();
+      if (worker) {
+        killPromises.push(waitForWorkerExit(worker));
+        delete allWorkers[worker.pid!];
+      }
     }
 
     if (immediate) {
-      Object.values(allWorkers).forEach(killWorker);
+      Object.values(allWorkers).forEach((worker) => {
+        killPromises.push(waitForWorkerExit(worker, 1000));
+        delete allWorkers[worker.pid!];
+      });
     }
-    // TODO set a timeout and force any outstanding workers to die
+
+    await Promise.all(killPromises);
   };
 
   const api = {

packages/engine-multi/test/api.test.ts

@@ -12,9 +12,9 @@ import loadVersions from '../src/util/load-versions';
 const logger = createMockLogger(undefined, { level: 'debug' });
 let api: RuntimeEngine;
 
-test.afterEach(() => {
+test.afterEach(async () => {
   logger._reset();
-  api?.destroy();
+  await api?.destroy();
 });
 
 test.serial('create a default engine api without throwing', async (t) => {

packages/engine-multi/test/integration.test.ts

@@ -11,9 +11,9 @@ let api: RuntimeEngine;
 
 const emptyState = {};
 
-test.afterEach(() => {
+test.afterEach(async () => {
   logger._reset();
-  api.destroy();
+  await api.destroy();
 });
 
 // this tests the full API with the actual runtime

packages/engine-multi/test/worker/pool.test.ts

@@ -187,7 +187,7 @@ test('throw if memory limit is exceeded', async (t) => {
 
 test('destroy should handle un-initialised workers', async (t) => {
   const pool = createPool(workerPath, { capacity: 10 }, logger);
-  pool.destroy();
+  await pool.destroy();
   t.is(pool._pool.length, 0);
 });
 
@@ -201,7 +201,7 @@ test('destroy should close all child processes', async (t) => {
   const workers = Object.values(pool._allWorkers);
 
   // now destroy it
-  pool.destroy();
+  await pool.destroy();
 
   // check that every child is disconnected
   t.true(workers.every((child) => child.killed));
@@ -210,27 +210,26 @@ test('destroy should close all child processes', async (t) => {
   t.is(pool._pool.length, 0);
 });
 
-test('destroy gracefully', (t) => {
-  return new Promise((done) => {
-    const pool = createPool(workerPath, {}, logger);
-    const workers = Object.values(pool._allWorkers);
+test('destroy gracefully', async (t) => {
+  const pool = createPool(workerPath, {}, logger);
+  const workers = Object.values(pool._allWorkers);
 
-    t.is(pool._pool.length, 5);
+  t.is(pool._pool.length, 5);
 
-    pool.exec('wait', [100]).then((result) => {
-      t.is(result, 1);
-      setTimeout(() => {
-        t.true(workers.every((child) => child.killed));
-        t.is(pool._pool.length, 0);
+  const taskPromise = pool.exec('wait', [100]);
 
-        done();
-      }, 1);
-    });
+  await pool.destroy();
 
-    pool.destroy();
+  t.is(pool._pool.length, 0);
 
-    t.is(pool._pool.length, 0);
-  });
+  // The task should complete gracefully
+  const result = await taskPromise;
+  t.is(result, 1);
+
+  // Give a moment for cleanup to finish
+  await new Promise((resolve) => setTimeout(resolve, 10));
+
+  t.true(workers.every((child) => child.killed));
 });
 
 // TODO should the worker throw on sigterm?