Add option to call initialize and end methods in workers

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - `[pretty-format]` Option to not escape strings in diff messages ([#5661](https://github.com/facebook/jest/pull/5661))
 - `[jest-haste-map]` Add `getFileIterator` to `HasteFS` for faster file iteration ([#7010](https://github.com/facebook/jest/pull/7010)).
+- `[jest-worker]` Add `initializeArgs` option to call an `initialize` method in the worker before the first call. Call `end` method in each worker when ending the farm ([#7014](https://github.com/facebook/jest/pull/7014)).
 
 ### Fixes
 

packages/jest-worker/src/__tests__/child.test.js

@@ -11,6 +11,7 @@ const mockError = new TypeError('Booo');
 const mockExtendedError = new ReferenceError('Booo extended');
 const processExit = process.exit;
 const processSend = process.send;
+const uninitializedParam = {};
 const sleep = ms => new Promise(resolve => setTimeout(resolve, ms));
 
 import {
@@ -21,17 +22,24 @@ import {
   PARENT_MESSAGE_ERROR,
 } from '../types';
 
+let ended;
 let mockCount;
+let initializeParm = uninitializedParam;
 
 beforeEach(() => {
   mockCount = 0;
+  ended = false;
 
   jest.mock(
     '../my-fancy-worker',
     () => {
       mockCount++;
 
       return {
+        end() {
+          ended = true;
+        },
+
         fooPromiseThrows() {
           return new Promise((resolve, reject) => {
             setTimeout(() => reject(mockError), 5);
@@ -68,6 +76,10 @@ beforeEach(() => {
         fooWorks() {
           return 1989;
         },
+
+        initialize(param) {
+          initializeParm = param;
+        },
       };
     },
     {virtual: true},
@@ -125,6 +137,50 @@ it('lazily requires the file', () => {
   ]);
 
   expect(mockCount).toBe(1);
+  expect(initializeParm).toBe(uninitializedParam); // Not called by default.
+});
+
+it('calls initialize before the first call', () => {
+  expect(mockCount).toBe(0);
+
+  process.emit('message', [
+    CHILD_MESSAGE_INITIALIZE,
+    true, // Not really used here, but for flow type purity.
+    './my-fancy-worker',
+    [], // Pass empty initialize params so the initialize method is called.
+  ]);
+
+  expect(initializeParm).toBe(uninitializedParam);
+
+  process.emit('message', [
+    CHILD_MESSAGE_CALL,
+    true, // Not really used here, but for flow type purity.
+    'fooWorks',
+    [],
+  ]);
+
+  expect(mockCount).toBe(1);
+  expect(initializeParm).toBe(undefined);
+});
+
+it('calls initialize with the correct arguments', () => {
+  expect(mockCount).toBe(0);
+
+  process.emit('message', [
+    CHILD_MESSAGE_INITIALIZE,
+    true, // Not really used here, but for flow type purity.
+    './my-fancy-worker',
+    ['foo'], // Pass empty initialize params so the initialize method is called.
+  ]);
+
+  process.emit('message', [
+    CHILD_MESSAGE_CALL,
+    true, // Not really used here, but for flow type purity.
+    'fooWorks',
+    [],
+  ]);
+
+  expect(initializeParm).toBe('foo');
 });
 
 it('returns results immediately when function is synchronous', () => {
@@ -295,6 +351,21 @@ it('finishes the process with exit code 0 if requested', () => {
   expect(process.exit.mock.calls[0]).toEqual([0]);
 });
 
+it('calls the end method ', () => {
+  process.emit('message', [
+    CHILD_MESSAGE_INITIALIZE,
+    true, // Not really used here, but for flow type purity.
+    './my-fancy-worker',
+  ]);
+
+  process.emit('message', [
+    CHILD_MESSAGE_END,
+    true, // Not really used here, but for flow type purity.
+  ]);
+
+  expect(ended).toBe(true);
+});
+
 it('throws if an invalid message is detected', () => {
   // Type 27 does not exist.
   expect(() => {

packages/jest-worker/src/__tests__/index.test.js

@@ -120,6 +120,7 @@ it('tries instantiating workers with the right options', () => {
   expect(Worker).toHaveBeenCalledTimes(4);
   expect(Worker.mock.calls[0][0]).toEqual({
     forkOptions: {execArgv: []},
+    initializeArgs: undefined,
     maxRetries: 6,
     workerId: 1,
     workerPath: '/tmp/baz.js',

packages/jest-worker/src/__tests__/worker.test.js

@@ -85,6 +85,7 @@ it('passes workerId to the child process and assign it to env.JEST_WORKER_ID', (
 it('initializes the child process with the given workerPath', () => {
   new Worker({
     forkOptions: {},
+    initializeArgs: ['foo', 'bar'],
     maxRetries: 3,
     workerPath: '/tmp/foo/bar/baz.js',
   });
@@ -93,6 +94,7 @@ it('initializes the child process with the given workerPath', () => {
     CHILD_MESSAGE_INITIALIZE,
     false,
     '/tmp/foo/bar/baz.js',
+    ['foo', 'bar'],
   ]);
 });
 

packages/jest-worker/src/child.js

@@ -14,10 +14,13 @@ import {
   CHILD_MESSAGE_END,
   CHILD_MESSAGE_INITIALIZE,
   PARENT_MESSAGE_ERROR,
+  PARENT_MESSAGE_INITIALIZE_ERROR,
   PARENT_MESSAGE_OK,
 } from './types';
 
 let file = null;
+let initializeArgs: ?Array<mixed> = null;
+let initialized = false;
 
 /**
  * This file is a small bootstrapper for workers. It sets up the communication
@@ -36,14 +39,15 @@ process.on('message', (request: any /* Should be ChildMessage */) => {
   switch (request[0]) {
     case CHILD_MESSAGE_INITIALIZE:
       file = request[2];
+      initializeArgs = request[3];
       break;
 
     case CHILD_MESSAGE_CALL:
       execMethod(request[2], request[3]);
       break;
 
     case CHILD_MESSAGE_END:
-      process.exit(0);
+      end();
       break;
 
     default:
@@ -61,7 +65,7 @@ function reportSuccess(result: any) {
   process.send([PARENT_MESSAGE_OK, result]);
 }
 
-function reportError(error: Error) {
+function reportError(error: Error, type?: number = PARENT_MESSAGE_ERROR) {
   if (!process || !process.send) {
     throw new Error('Child can only be used on a forked process');
   }
@@ -71,7 +75,7 @@ function reportError(error: Error) {
   }
 
   process.send([
-    PARENT_MESSAGE_ERROR,
+    type,
     error.constructor && error.constructor.name,
     error.message,
     error.stack,
@@ -80,25 +84,72 @@ function reportError(error: Error) {
   ]);
 }
 
+function end(): void {
+  // $FlowFixMe: This has to be a dynamic require.
+  const main = require(file);
+
+  if (!main['end']) {
+    process.exit(0);
+
+    return;
+  }
+
+  execFunction(main['end'], main, [], () => process.exit(0), () => {});
+}
+
 function execMethod(method: string, args: $ReadOnlyArray<any>): void {
   // $FlowFixMe: This has to be a dynamic require.
   const main = require(file);
+  const initializeArgsForFlow = initializeArgs;
+
+  let fn;
+  let ctx;
+
+  if (method === 'default') {
+    fn = main.__esModule ? main['default'] : main;
+    ctx = global;
+  } else {
+    fn = main[method];
+    ctx = main;
+  }
+
+  if (!initializeArgsForFlow || initialized || !main['initialize']) {
+    execFunction(fn, ctx, args, reportSuccess, reportError);
+
+    return;
+  }
+
+  initialized = true;
+
+  execFunction(
+    main['initialize'],
+    main,
+    initializeArgsForFlow,
+    () => execFunction(fn, ctx, args, reportSuccess, reportError),
+    error => reportError(error, PARENT_MESSAGE_INITIALIZE_ERROR),
+  );
+}
+
+function execFunction(
+  fn: Function,
+  ctx: any,
+  args: $ReadOnlyArray<any>,
+  onResult: (result: any) => void,
+  onError: (error: Error, type?: number) => void,
+): void {
   let result;
 
   try {
-    if (method === 'default') {
-      result = (main.__esModule ? main['default'] : main).apply(global, args);
-    } else {
-      result = main[method].apply(main, args);
-    }
+    result = fn.apply(ctx, args);
   } catch (err) {
-    reportError(err);
+    onError(err);
+
     return;
   }
 
   if (result && typeof result.then === 'function') {
-    result.then(reportSuccess, reportError);
+    result.then(onResult, onError);
   } else {
-    reportSuccess(result);
+    onResult(result);
   }
 }

packages/jest-worker/src/index.js

@@ -68,6 +68,7 @@ export default class {
 
     const sharedWorkerOptions = {
       forkOptions: options.forkOptions || {},
+      initializeArgs: options.initializeArgs,
       maxRetries: options.maxRetries || 3,
       workerPath,
     };

packages/jest-worker/src/types.js

@@ -21,6 +21,7 @@ export const CHILD_MESSAGE_END: 2 = 2;
 
 export const PARENT_MESSAGE_OK: 0 = 0;
 export const PARENT_MESSAGE_ERROR: 1 = 1;
+export const PARENT_MESSAGE_INITIALIZE_ERROR: 2 = 2;
 
 // Option objects.
 
@@ -41,12 +42,14 @@ export type FarmOptions = {
   computeWorkerKey?: (string, ...Array<any>) => ?string,
   exposedMethods?: $ReadOnlyArray<string>,
   forkOptions?: ForkOptions,
+  initializeArgs?: Array<mixed>,
   maxRetries?: number,
   numWorkers?: number,
 };
 
 export type WorkerOptions = {|
   forkOptions: ForkOptions,
+  initializeArgs: ?Array<mixed>,
   maxRetries: number,
   workerId: number,
   workerPath: string,
@@ -58,6 +61,7 @@ export type ChildMessageInitialize = [
   typeof CHILD_MESSAGE_INITIALIZE, // type
   boolean, // processed
   string, // file
+  ?Array<mixed>, // initializeArgs
 ];
 
 export type ChildMessageCall = [

packages/jest-worker/src/worker.js

@@ -14,6 +14,7 @@ import childProcess from 'child_process';
 import {
   CHILD_MESSAGE_INITIALIZE,
   PARENT_MESSAGE_ERROR,
+  PARENT_MESSAGE_INITIALIZE_ERROR,
   PARENT_MESSAGE_OK,
 } from './types';
 
@@ -108,7 +109,12 @@ export default class {
     child.on('exit', this._exit.bind(this));
 
     // $FlowFixMe: wrong "ChildProcess.send" signature.
-    child.send([CHILD_MESSAGE_INITIALIZE, false, this._options.workerPath]);
+    child.send([
+      CHILD_MESSAGE_INITIALIZE,
+      false,
+      this._options.workerPath,
+      this._options.initializeArgs,
+    ]);
 
     this._retries++;
     this._child = child;
@@ -175,13 +181,15 @@ export default class {
     this._busy = false;
     this._process();
 
+    let error;
+
     switch (response[0]) {
       case PARENT_MESSAGE_OK:
         onProcessEnd(null, response[1]);
         break;
 
       case PARENT_MESSAGE_ERROR:
-        let error = response[4];
+        error = response[4];
 
         if (error != null && typeof error === 'object') {
           const extra = error;
@@ -202,6 +210,16 @@ export default class {
         onProcessEnd(error, null);
         break;
 
+      case PARENT_MESSAGE_INITIALIZE_ERROR:
+        error = new Error('Error when calling initialize: ' + response[2]);
+
+        // $FlowFixMe: adding custom properties to errors.
+        error.type = response[1];
+        error.stack = response[3];
+
+        onProcessEnd(error, null);
+        break;
+
       default:
         throw new TypeError('Unexpected response from worker: ' + response[0]);
     }