Add option to call setup and teardown 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]` [**BREAKING**] Add functionality to call a `setup` method in the worker before the first call and a `teardown` method when ending the farm ([#7014](https://github.com/facebook/jest/pull/7014)).
 
 ### Fixes
 

packages/jest-worker/README.md

@@ -73,6 +73,10 @@ The callback you provide is called with the method name, plus all the rest of th
 
 By default, no process is bound to any worker.
 
+#### `setupArgs: Array<mixed>` (optional)
+
+The arguments that will be passed to the `setup` method during initialization.
+
 ## Worker
 
 The returned `Worker` instance has all the exposed methods, plus some additional ones to interact with the workers itself:
@@ -91,6 +95,13 @@ Finishes the workers by killing all workers. No further calls can be done to the
 
 **Note:** Each worker has a unique id (index that starts with `1`) which is available on `process.env.JEST_WORKER_ID`
 
+## Setting up and tearing down the child process
+
+The child process can define two special methods (both of them can be asynchronous):
+
+* `setup()`: If defined, it's executed before the first call to any method in the child.
+* `teardown()`: If defined, it's executed when the farm ends.
+
 # More examples
 
 ## Standard usage

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

@@ -11,20 +11,24 @@ 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 {
   CHILD_MESSAGE_INITIALIZE,
   CHILD_MESSAGE_CALL,
   CHILD_MESSAGE_END,
   PARENT_MESSAGE_OK,
-  PARENT_MESSAGE_ERROR,
+  PARENT_MESSAGE_CLIENT_ERROR,
 } from '../types';
 
+let ended;
 let mockCount;
+let initializeParm = uninitializedParam;
 
 beforeEach(() => {
   mockCount = 0;
+  ended = false;
 
   jest.mock(
     '../my-fancy-worker',
@@ -68,6 +72,14 @@ beforeEach(() => {
         fooWorks() {
           return 1989;
         },
+
+        setup(param) {
+          initializeParm = param;
+        },
+
+        teardown() {
+          ended = true;
+        },
       };
     },
     {virtual: true},
@@ -116,6 +128,7 @@ it('lazily requires the file', () => {
   ]);
 
   expect(mockCount).toBe(0);
+  expect(initializeParm).toBe(uninitializedParam); // Not called yet.
 
   process.emit('message', [
     CHILD_MESSAGE_CALL,
@@ -125,6 +138,27 @@ it('lazily requires the file', () => {
   ]);
 
   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', () => {
@@ -153,7 +187,7 @@ it('returns results immediately when function is synchronous', () => {
   ]);
 
   expect(process.send.mock.calls[1][0]).toEqual([
-    PARENT_MESSAGE_ERROR,
+    PARENT_MESSAGE_CLIENT_ERROR,
     'TypeError',
     'Booo',
     mockError.stack,
@@ -168,7 +202,7 @@ it('returns results immediately when function is synchronous', () => {
   ]);
 
   expect(process.send.mock.calls[2][0]).toEqual([
-    PARENT_MESSAGE_ERROR,
+    PARENT_MESSAGE_CLIENT_ERROR,
     'Number',
     void 0,
     void 0,
@@ -183,7 +217,7 @@ it('returns results immediately when function is synchronous', () => {
   ]);
 
   expect(process.send.mock.calls[3][0]).toEqual([
-    PARENT_MESSAGE_ERROR,
+    PARENT_MESSAGE_CLIENT_ERROR,
     'ReferenceError',
     'Booo extended',
     mockExtendedError.stack,
@@ -197,7 +231,7 @@ it('returns results immediately when function is synchronous', () => {
     [],
   ]);
 
-  expect(process.send.mock.calls[4][0][0]).toBe(PARENT_MESSAGE_ERROR);
+  expect(process.send.mock.calls[4][0][0]).toBe(PARENT_MESSAGE_CLIENT_ERROR);
   expect(process.send.mock.calls[4][0][1]).toBe('Error');
   expect(process.send.mock.calls[4][0][2]).toEqual(
     '"null" or "undefined" thrown',
@@ -236,7 +270,7 @@ it('returns results when it gets resolved if function is asynchronous', async ()
   await sleep(10);
 
   expect(process.send.mock.calls[1][0]).toEqual([
-    PARENT_MESSAGE_ERROR,
+    PARENT_MESSAGE_CLIENT_ERROR,
     'TypeError',
     'Booo',
     mockError.stack,
@@ -295,6 +329,21 @@ it('finishes the process with exit code 0 if requested', () => {
   expect(process.exit.mock.calls[0]).toEqual([0]);
 });
 
+it('calls the teardown 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

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

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

@@ -14,7 +14,7 @@ import EventEmitter from 'events';
 import {
   CHILD_MESSAGE_CALL,
   CHILD_MESSAGE_INITIALIZE,
-  PARENT_MESSAGE_ERROR,
+  PARENT_MESSAGE_CLIENT_ERROR,
   PARENT_MESSAGE_OK,
 } from '../types';
 
@@ -86,13 +86,15 @@ it('initializes the child process with the given workerPath', () => {
   new Worker({
     forkOptions: {},
     maxRetries: 3,
+    setupArgs: ['foo', 'bar'],
     workerPath: '/tmp/foo/bar/baz.js',
   });
 
   expect(forkInterface.send.mock.calls[0][0]).toEqual([
     CHILD_MESSAGE_INITIALIZE,
     false,
     '/tmp/foo/bar/baz.js',
+    ['foo', 'bar'],
   ]);
 });
 
@@ -201,7 +203,7 @@ it('relates replies to requests, in order', () => {
 
   // and then the second call replies...
   forkInterface.emit('message', [
-    PARENT_MESSAGE_ERROR,
+    PARENT_MESSAGE_CLIENT_ERROR,
     'TypeError',
     'foo',
     'TypeError: foo',
@@ -287,7 +289,7 @@ it('creates error instances for known errors', () => {
   worker.send([CHILD_MESSAGE_CALL, false, 'method', []], () => {}, callback1);
 
   forkInterface.emit('message', [
-    PARENT_MESSAGE_ERROR,
+    PARENT_MESSAGE_CLIENT_ERROR,
     'TypeError',
     'bar',
     'TypeError: bar',
@@ -303,7 +305,7 @@ it('creates error instances for known errors', () => {
   worker.send([CHILD_MESSAGE_CALL, false, 'method', []], () => {}, callback2);
 
   forkInterface.emit('message', [
-    PARENT_MESSAGE_ERROR,
+    PARENT_MESSAGE_CLIENT_ERROR,
     'RandomCustomError',
     'bar',
     'RandomCustomError: bar',
@@ -320,7 +322,7 @@ it('creates error instances for known errors', () => {
   worker.send([CHILD_MESSAGE_CALL, false, 'method', []], () => {}, callback3);
 
   forkInterface.emit('message', [
-    PARENT_MESSAGE_ERROR,
+    PARENT_MESSAGE_CLIENT_ERROR,
     'Number',
     null,
     null,

packages/jest-worker/src/child.js

@@ -13,11 +13,16 @@ import {
   CHILD_MESSAGE_CALL,
   CHILD_MESSAGE_END,
   CHILD_MESSAGE_INITIALIZE,
-  PARENT_MESSAGE_ERROR,
+  PARENT_MESSAGE_CLIENT_ERROR,
+  PARENT_MESSAGE_SETUP_ERROR,
   PARENT_MESSAGE_OK,
 } from './types';
 
+import type {PARENT_MESSAGE_ERROR} from './types';
+
 let file = null;
+let setupArgs: Array<mixed> = [];
+let initialized = false;
 
 /**
  * This file is a small bootstrapper for workers. It sets up the communication
@@ -36,14 +41,15 @@ process.on('message', (request: any /* Should be ChildMessage */) => {
   switch (request[0]) {
     case CHILD_MESSAGE_INITIALIZE:
       file = request[2];
+      setupArgs = 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 +67,15 @@ function reportSuccess(result: any) {
   process.send([PARENT_MESSAGE_OK, result]);
 }
 
-function reportError(error: Error) {
+function reportClientError(error: Error) {
+  return reportError(error, PARENT_MESSAGE_CLIENT_ERROR);
+}
+
+function reportInitializeError(error: Error) {
+  return reportError(error, PARENT_MESSAGE_SETUP_ERROR);
+}
+
+function reportError(error: Error, type: PARENT_MESSAGE_ERROR) {
   if (!process || !process.send) {
     throw new Error('Child can only be used on a forked process');
   }
@@ -71,7 +85,7 @@ function reportError(error: Error) {
   }
 
   process.send([
-    PARENT_MESSAGE_ERROR,
+    type,
     error.constructor && error.constructor.name,
     error.message,
     error.stack,
@@ -80,25 +94,70 @@ function reportError(error: Error) {
   ]);
 }
 
+function end(): void {
+  // $FlowFixMe: This has to be a dynamic require.
+  const main = require(file);
+
+  if (!main.teardown) {
+    exitProcess();
+
+    return;
+  }
+
+  execFunction(main.teardown, main, [], exitProcess, exitProcess);
+}
+
+function exitProcess(): void {
+  process.exit(0);
+}
+
 function execMethod(method: string, args: $ReadOnlyArray<any>): void {
   // $FlowFixMe: This has to be a dynamic require.
   const main = require(file);
+
+  let fn;
+
+  if (method === 'default') {
+    fn = main.__esModule ? main['default'] : main;
+  } else {
+    fn = main[method];
+  }
+
+  function execHelper() {
+    execFunction(fn, main, args, reportSuccess, reportClientError);
+  }
+
+  if (initialized || !main.setup) {
+    execHelper();
+
+    return;
+  }
+
+  initialized = true;
+
+  execFunction(main.setup, main, setupArgs, execHelper, reportInitializeError);
+}
+
+function execFunction(
+  fn: (...args: $ReadOnlyArray<mixed>) => mixed,
+  ctx: mixed,
+  args: $ReadOnlyArray<mixed>,
+  onResult: (result: mixed) => void,
+  onError: (error: Error) => 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

@@ -69,6 +69,7 @@ export default class {
     const sharedWorkerOptions = {
       forkOptions: options.forkOptions || {},
       maxRetries: options.maxRetries || 3,
+      setupArgs: options.setupArgs || [],
       workerPath,
     };