fix: worker being killed after being spawned and other worker bugs

e2e/worker-restarting/__tests__/test1.js

@@ -0,0 +1,10 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+test('basic test', () => {
+  expect(true).toBeTruthy();
+});

e2e/worker-restarting/__tests__/test2.js

@@ -0,0 +1,10 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+test('basic test', () => {
+  expect(true).toBeTruthy();
+});

e2e/worker-restarting/__tests__/test3.js

@@ -0,0 +1,10 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+test('basic test', () => {
+  expect(true).toBeTruthy();
+});

e2e/worker-restarting/package.json

@@ -0,0 +1,6 @@
+{
+  "jest": {
+    "maxWorkers": 2,
+    "workerIdleMemoryLimit": "1MB"
+  }
+}

packages/jest-worker/src/types.ts

@@ -66,12 +66,15 @@ export interface WorkerPoolInterface {
 }
 
 export interface WorkerInterface {
+  get state(): WorkerStates;
+
   send(
     request: ChildMessage,
     onProcessStart: OnStart,
     onProcessEnd: OnEnd,
     onCustomMessage: OnCustomMessage,
   ): void;
+
   waitForExit(): Promise<void>;
   forceExit(): void;
 
@@ -83,6 +86,18 @@ export interface WorkerInterface {
    */
   getWorkerSystemId(): number;
   getMemoryUsage(): Promise<number | null>;
+  /**
+   * Checks to see if the child worker is actually running.
+   */
+  isWorkerRunning(): boolean;
+  /**
+   * When the worker child is started and ready to start handling requests.
+   *
+   * @remarks
+   * This mostly exists to help with testing so that you don't check the status
+   * of things like isWorkerRunning before it actually is.
+   */
+  waitForWorkerReady(): Promise<void>;
 }
 
 export type PoolExitResult = {
@@ -170,8 +185,21 @@ export type WorkerOptions = {
    * the raw output of the worker.
    */
   silent?: boolean;
+  /**
+   * Used to immediately bind event handlers.
+   */
+  on?: {
+    [WorkerEvents.STATE_CHANGE]:
+      | onStateChangeHandler
+      | Array<onStateChangeHandler>;
+  };
 };
 
+export type onStateChangeHandler = (
+  state: WorkerStates,
+  oldState: WorkerStates,
+) => void;
+
 // Messages passed from the parent to the children.
 
 export type MessagePort = typeof EventEmitter & {
@@ -265,3 +293,7 @@ export enum WorkerStates {
   SHUTTING_DOWN = 'shutting-down',
   SHUT_DOWN = 'shut-down',
 }
+
+export enum WorkerEvents {
+  STATE_CHANGE = 'state-change',
+}

packages/jest-worker/src/workers/ChildProcessWorker.ts

@@ -7,7 +7,6 @@
 
 import {ChildProcess, ForkOptions, fork} from 'child_process';
 import {totalmem} from 'os';
-import {PassThrough} from 'stream';
 import mergeStream = require('merge-stream');
 import {stdout as stdoutSupportsColor} from 'supports-color';
 import {
@@ -27,13 +26,14 @@ import {
   WorkerOptions,
   WorkerStates,
 } from '../types';
+import WorkerAbstract from './WorkerAbstract';
 
 const SIGNAL_BASE_EXIT_CODE = 128;
 const SIGKILL_EXIT_CODE = SIGNAL_BASE_EXIT_CODE + 9;
 const SIGTERM_EXIT_CODE = SIGNAL_BASE_EXIT_CODE + 15;
 
 // How long to wait after SIGTERM before sending SIGKILL
-const SIGKILL_DELAY = 500;
+export const SIGKILL_DELAY = 500;
 
 /**
  * This class wraps the child process and provides a nice interface to
@@ -53,7 +53,10 @@ const SIGKILL_DELAY = 500;
  * field is changed to "true", so that other workers which might encounter the
  * same call skip it.
  */
-export default class ChildProcessWorker implements WorkerInterface {
+export default class ChildProcessWorker
+  extends WorkerAbstract
+  implements WorkerInterface
+{
   private _child!: ChildProcess;
   private _options: WorkerOptions;
 
@@ -62,50 +65,42 @@ export default class ChildProcessWorker implements WorkerInterface {
   private _onProcessEnd!: OnEnd;
   private _onCustomMessage!: OnCustomMessage;
 
-  private _fakeStream: PassThrough | null;
   private _stdout: ReturnType<typeof mergeStream> | null;
   private _stderr: ReturnType<typeof mergeStream> | null;
 
-  private _exitPromise: Promise<void>;
-  private _resolveExitPromise!: () => void;
-
   private _memoryUsagePromise: Promise<number> | undefined;
   private _resolveMemoryUsage: ((arg0: number) => void) | undefined;
 
   private _childIdleMemoryUsage: number | null;
   private _childIdleMemoryUsageLimit: number | null;
   private _memoryUsageCheck = false;
-  private _state: WorkerStates;
 
   private _childWorkerPath: string;
 
   constructor(options: WorkerOptions) {
+    super(options);
+
     this._options = options;
 
     this._request = null;
 
-    this._fakeStream = null;
     this._stdout = null;
     this._stderr = null;
     this._childIdleMemoryUsage = null;
     this._childIdleMemoryUsageLimit = options.idleMemoryLimit || null;
 
-    this._exitPromise = new Promise(resolve => {
-      this._resolveExitPromise = resolve;
-    });
-
     this._childWorkerPath =
       options.childWorkerPath || require.resolve('./processChild');
 
-    this._state = WorkerStates.STARTING;
+    this.state = WorkerStates.STARTING;
     this.initialize();
   }
 
   initialize(): void {
     if (
-      this._state === WorkerStates.OUT_OF_MEMORY ||
-      this._state === WorkerStates.SHUTTING_DOWN ||
-      this._state === WorkerStates.SHUT_DOWN
+      this.state === WorkerStates.OUT_OF_MEMORY ||
+      this.state === WorkerStates.SHUTTING_DOWN ||
+      this.state === WorkerStates.SHUT_DOWN
     ) {
       return;
     }
@@ -114,7 +109,7 @@ export default class ChildProcessWorker implements WorkerInterface {
       this._child.kill('SIGKILL');
     }
 
-    this._state = WorkerStates.STARTING;
+    this.state = WorkerStates.STARTING;
 
     const forceColor = stdoutSupportsColor ? {FORCE_COLOR: '1'} : {};
     const options: ForkOptions = {
@@ -132,42 +127,39 @@ export default class ChildProcessWorker implements WorkerInterface {
       ...this._options.forkOptions,
     };
 
-    const child = fork(this._childWorkerPath, [], options);
+    this._child = fork(this._childWorkerPath, [], options);
 
-    if (child.stdout) {
+    if (this._child.stdout) {
       if (!this._stdout) {
         // We need to add a permanent stream to the merged stream to prevent it
         // from ending when the subprocess stream ends
         this._stdout = mergeStream(this._getFakeStream());
       }
 
-      this._stdout.add(child.stdout);
+      this._stdout.add(this._child.stdout);
     }
 
-    if (child.stderr) {
+    if (this._child.stderr) {
       if (!this._stderr) {
         // We need to add a permanent stream to the merged stream to prevent it
         // from ending when the subprocess stream ends
         this._stderr = mergeStream(this._getFakeStream());
       }
 
-      this._stderr.add(child.stderr);
+      this._stderr.add(this._child.stderr);
     }
 
-    this._detectOutOfMemoryCrash(child);
-    child.on('message', this._onMessage.bind(this));
-    child.on('exit', this._onExit.bind(this));
+    this._detectOutOfMemoryCrash(this._child);
+    this._child.on('message', this._onMessage.bind(this));
+    this._child.on('exit', this._onExit.bind(this));
 
-    child.send([
+    this._child.send([
       CHILD_MESSAGE_INITIALIZE,
       false,
       this._options.workerPath,
       this._options.setupArgs,
     ]);
 
-    this._child = child;
-    this._state = WorkerStates.OK;
-
     this._retries++;
 
     // If we exceeded the amount of retries, we will emulate an error reply
@@ -189,13 +181,18 @@ export default class ChildProcessWorker implements WorkerInterface {
       // Clear the request so we don't keep executing it.
       this._request = null;
     }
+
+    this.state = WorkerStates.OK;
+    if (this._resolveWorkerReady) {
+      this._resolveWorkerReady();
+    }
   }
 
   private _detectOutOfMemoryCrash(child: ChildProcess): void {
     let stderrStr = '';
 
     const handler = (chunk: any) => {
-      if (this._state !== WorkerStates.OUT_OF_MEMORY) {
+      if (this.state === WorkerStates.OK) {
         let str: string | undefined = undefined;
 
         if (chunk instanceof Buffer) {
@@ -209,26 +206,14 @@ export default class ChildProcessWorker implements WorkerInterface {
         }
 
         if (stderrStr.includes('heap out of memory')) {
-          this._state = WorkerStates.OUT_OF_MEMORY;
+          this.state = WorkerStates.OUT_OF_MEMORY;
         }
       }
     };
 
     child.stderr?.on('data', handler);
   }
 
-  private _shutdown() {
-    this._state = WorkerStates.SHUTTING_DOWN;
-
-    // End the temporary streams so the merged streams end too
-    if (this._fakeStream) {
-      this._fakeStream.end();
-      this._fakeStream = null;
-    }
-
-    this._resolveExitPromise();
-  }
-
   private _onMessage(response: ParentMessage) {
     // TODO: Add appropriate type check
     let error: any;
@@ -311,15 +296,18 @@ export default class ChildProcessWorker implements WorkerInterface {
         this._childIdleMemoryUsage &&
         this._childIdleMemoryUsage > limit
       ) {
-        this._state = WorkerStates.RESTARTING;
+        this.state = WorkerStates.RESTARTING;
 
         this.killChild();
       }
     }
   }
 
   private _onExit(exitCode: number | null) {
-    if (exitCode !== 0 && this._state === WorkerStates.OUT_OF_MEMORY) {
+    this._workerReadyPromise = undefined;
+    this._resolveWorkerReady = undefined;
+
+    if (exitCode !== 0 && this.state === WorkerStates.OUT_OF_MEMORY) {
       this._onProcessEnd(
         new Error('Jest worker ran out of memory and crashed'),
         null,
@@ -331,8 +319,8 @@ export default class ChildProcessWorker implements WorkerInterface {
         exitCode !== null &&
         exitCode !== SIGTERM_EXIT_CODE &&
         exitCode !== SIGKILL_EXIT_CODE &&
-        this._state !== WorkerStates.SHUTTING_DOWN) ||
-      this._state === WorkerStates.RESTARTING
+        this.state !== WorkerStates.SHUTTING_DOWN) ||
+      this.state === WorkerStates.RESTARTING
     ) {
       this.initialize();
 
@@ -353,11 +341,17 @@ export default class ChildProcessWorker implements WorkerInterface {
     onProcessStart(this);
 
     this._onProcessEnd = (...args) => {
+      const hasRequest = !!this._request;
+
       // Clean the request to avoid sending past requests to workers that fail
       // while waiting for a new request (timers, unhandled rejections...)
       this._request = null;
 
-      if (this._childIdleMemoryUsageLimit && this._child.connected) {
+      if (
+        this._childIdleMemoryUsageLimit &&
+        this._child.connected &&
+        hasRequest
+      ) {
         this.checkMemoryUsage();
       }
 
@@ -377,12 +371,16 @@ export default class ChildProcessWorker implements WorkerInterface {
   }
 
   killChild(): NodeJS.Timeout {
-    this._child.kill('SIGTERM');
-    return setTimeout(() => this._child.kill('SIGKILL'), SIGKILL_DELAY);
+    // We store a reference so that there's no way we can accidentally
+    // kill a new worker that has been spawned.
+    const childToKill = this._child;
+
+    childToKill.kill('SIGTERM');
+    return setTimeout(() => childToKill.kill('SIGKILL'), SIGKILL_DELAY);
   }
 
   forceExit(): void {
-    this._state = WorkerStates.SHUTTING_DOWN;
+    this.state = WorkerStates.SHUTTING_DOWN;
 
     const sigkillTimeout = this.killChild();
     this._exitPromise.then(() => clearTimeout(sigkillTimeout));
@@ -466,10 +464,7 @@ export default class ChildProcessWorker implements WorkerInterface {
     }
   }
 
-  private _getFakeStream() {
-    if (!this._fakeStream) {
-      this._fakeStream = new PassThrough();
-    }
-    return this._fakeStream;
+  isWorkerRunning(): boolean {
+    return this._child.connected && !this._child.killed;
   }
 }