WIP: t31135461 / Refine jest-worker API

packages/jest-resolve/src/is_builtin_module.js

@@ -15,11 +15,13 @@ declare var process: {
   binding(type: string): {},
 };
 
+const EXPERIMENTAL_MODULES = ['worker_threads'];
+
 const BUILTIN_MODULES =
-  builtinModules ||
-  Object.keys(process.binding('natives')).filter(
-    (module: string) => !/^internal\//.test(module),
-  );
+  builtinModules.concat(EXPERIMENTAL_MODULES) ||
+  Object.keys(process.binding('natives'))
+    .filter((module: string) => !/^internal\//.test(module))
+    .concat([EXPERIMENTAL_MODULES]);
 
 export default function isBuiltinModule(module: string): boolean {
   return BUILTIN_MODULES.indexOf(module) !== -1;

packages/jest-worker/src/WorkerPool.js

@@ -0,0 +1,27 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc. 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.
+ *
+ * @flow
+ */
+
+'use strict';
+
+import BaseWorkerPool from './base/BaseWorkerPool';
+
+import type {ChildMessage, OnStart, OnEnd, WorkerPoolInterface} from './types';
+
+class WorkerPool extends BaseWorkerPool implements WorkerPoolInterface {
+  send(
+    workerId: number,
+    request: ChildMessage,
+    onStart: OnStart,
+    onEnd: OnEnd,
+  ): void {
+    this.getWorkerById(workerId).send(request, onStart, onEnd);
+  }
+}
+
+export default WorkerPool;

packages/jest-worker/src/WorkerQueueManager.js

@@ -0,0 +1,131 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc. 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.
+ *
+ * @flow
+ */
+
+'use strict';
+
+import type {
+  QueueChildMessage,
+  WorkerPoolInterface,
+  WorkerInterface,
+} from './types';
+
+export default class WorkerQueueManager {
+  _workerPool: WorkerPoolInterface;
+  _queue: Array<?QueueChildMessage>;
+  _last: Array<QueueChildMessage>;
+  _locks: Array<boolean>;
+  _offset: number;
+
+  constructor(workerPool: WorkerPoolInterface) {
+    this._workerPool = workerPool;
+    this._queue = [];
+    this._last = [];
+    this._locks = [];
+    this._offset = 0;
+
+    // If we exceeded the amount of retries, we will emulate an error reply
+    // coming from the child. This avoids code duplication related with cleaning
+    // the queue, and scheduling the next call.
+
+    // if (this._retries > this._options.maxRetries) {
+    //   const error = new Error('Call retries were exceeded');
+
+    //   this.onMessage([
+    //     PARENT_MESSAGE_ERROR,
+    //     error.name,
+    //     error.message,
+    //     error.stack,
+    //     {type: 'WorkerError'},
+    //   ]);
+    // }
+  }
+
+  _process(workerId: number): WorkerQueueManager {
+    if (this.isLocked(workerId)) {
+      return this;
+    }
+
+    const job = this.getNextJob(workerId);
+
+    if (!job) {
+      return this;
+    }
+
+    const onEnd = (error: ?Error, result: mixed, worker: WorkerInterface) => {
+      this.unlock(workerId);
+      job.onEnd(error, result, worker);
+      this._process(workerId);
+    };
+
+    this.lock(workerId);
+
+    this._workerPool.send(workerId, job.request, job.onStart, onEnd);
+
+    job.request[1] = true;
+
+    return this;
+  }
+
+  getNextJob(workerId: number): ?QueueChildMessage {
+    if (!this._queue[workerId]) {
+      return null;
+    }
+
+    let job;
+
+    while (this._queue[workerId]) {
+      if (!this._queue[workerId].request[1]) {
+        job = this._queue[workerId];
+        break;
+      }
+      this._queue[workerId] = this._queue[workerId].next;
+    }
+
+    return job;
+  }
+
+  enqueue(task: QueueChildMessage, workerId?: number): WorkerQueueManager {
+    if (workerId != null) {
+      if (task.request[1]) {
+        return this;
+      }
+
+      if (this._queue[workerId]) {
+        this._last[workerId].next = task;
+      } else {
+        this._queue[workerId] = task;
+      }
+
+      this._last[workerId] = task;
+
+      this._process(workerId);
+    } else {
+      const numOfWorkers = this._workerPool.getWorkers().length;
+      for (let i = 0; i < numOfWorkers; i++) {
+        const workerIdx = (this._offset + i) % numOfWorkers;
+        this.enqueue(task, workerIdx);
+      }
+      this._offset++;
+    }
+
+    return this;
+  }
+
+  lock(workerId: number): void {
+    this._locks[workerId] = true;
+  }
+
+  unlock(workerId: number): void {
+    this._locks[workerId] = false;
+  }
+
+  isLocked(workerId: number): boolean {
+    return this._locks[workerId];
+  }
+}

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

@@ -38,7 +38,7 @@ beforeEach(() => {
     return forkInterface;
   });
 
-  Worker = require('../worker').default;
+  Worker = require('../workers/ChildProcessWorker').default;
 });
 
 afterEach(() => {
@@ -47,7 +47,7 @@ afterEach(() => {
 });
 
 it('passes fork options down to child_process.fork, adding the defaults', () => {
-  const child = require.resolve('../child');
+  const child = require.resolve('../workers/processChild');
 
   process.execArgv = ['--inspect', '-p'];
 

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

@@ -27,7 +27,7 @@ beforeEach(() => {
   mockCount = 0;
 
   jest.mock(
-    '../my-fancy-worker',
+    '../workers/my-fancy-worker',
     () => {
       mockCount++;
 
@@ -74,15 +74,15 @@ beforeEach(() => {
   );
 
   jest.mock(
-    '../my-fancy-standalone-worker',
+    '../workers/my-fancy-standalone-worker',
     () => jest.fn().mockImplementation(() => 12345),
     {virtual: true},
   );
 
   // This mock emulates a transpiled Babel module that carries a default export
   // that corresponds to a method.
   jest.mock(
-    '../my-fancy-babel-worker',
+    '../workers/my-fancy-babel-worker',
     () => ({
       __esModule: true,
       default: jest.fn().mockImplementation(() => 67890),
@@ -94,7 +94,7 @@ beforeEach(() => {
   process.send = jest.fn();
 
   // Require the child!
-  require('../child');
+  require('../workers/processChild');
 });
 
 afterEach(() => {

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

@@ -30,7 +30,7 @@ beforeEach(() => {
   // The worker mock returns a worker with custom methods, plus it stores them
   // in a global list, so that they can be accessed later. This list is reset in
   // every test.
-  jest.mock('../worker', () => {
+  jest.mock('../workers/ChildProcessWorker', () => {
     const fakeClass = jest.fn(() => {
       const fakeWorker = {
         getStderr: () => ({once() {}, pipe() {}}),
@@ -63,8 +63,8 @@ beforeEach(() => {
     virtual: true,
   });
 
-  Worker = require('../worker').default;
-  Farm = require('../index').default;
+  Worker = require('../workers/ChildProcessWorker').default;
+  Farm = require('..').default;
 });
 
 afterEach(() => {

packages/jest-worker/src/base/BaseWorkerPool.js

@@ -0,0 +1,115 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc. 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.
+ *
+ * @flow
+ */
+
+'use strict';
+
+import mergeStream from 'merge-stream';
+import os from 'os';
+import path from 'path';
+
+import ChildProcessWorker from '../workers/ChildProcessWorker';
+import NodeThreadsWorker from '../workers/NodeThreadsWorker';
+import {CHILD_MESSAGE_END} from '../types';
+
+import type {Readable} from 'stream';
+import type {FarmOptions, WorkerOptions, WorkerInterface} from '../types';
+
+/* istanbul ignore next */
+const emptyMethod = () => {};
+
+export default class BaseWorkerPool {
+  _stderr: Readable;
+  _stdout: Readable;
+  _options: FarmOptions;
+  _workers: Array<WorkerInterface>;
+
+  constructor(workerPath: string, options: FarmOptions) {
+    this._options = options;
+
+    const numWorkers = options.numWorkers || os.cpus().length - 1;
+    this._workers = new Array(numWorkers);
+
+    if (!path.isAbsolute(workerPath)) {
+      workerPath = require.resolve(workerPath);
+    }
+
+    const stdout = mergeStream();
+    const stderr = mergeStream();
+
+    for (let i = 0; i < numWorkers; i++) {
+      const workerOptions: WorkerOptions = {
+        forkOptions: options.forkOptions || {},
+        maxRetries: options.maxRetries || 3,
+        useNodeWorkersIfPossible: options.useNodeWorkersIfPossible,
+        workerId: i,
+        workerPath,
+      };
+
+      const worker = this.createWorker(workerOptions);
+      const workerStdout = worker.getStdout();
+      const workerStderr = worker.getStderr();
+
+      if (workerStdout) {
+        stdout.add(workerStdout);
+      }
+
+      if (workerStderr) {
+        stderr.add(workerStderr);
+      }
+
+      this._workers[i] = worker;
+
+      // $FlowFixMe
+      this.getStderr = this.getStderr.bind(this);
+      // $FlowFixMe
+      this.getStdout = this.getStdout.bind(this);
+      // $FlowFixMe
+      this.getWorkers = this.getWorkers.bind(this);
+      // $FlowFixMe
+      this.end = this.end.bind(this);
+    }
+
+    this._stdout = stdout;
+    this._stderr = stderr;
+  }
+
+  getStderr(): Readable {
+    return this._stderr;
+  }
+
+  getStdout(): Readable {
+    return this._stdout;
+  }
+
+  getWorkers(): Array<WorkerInterface> {
+    return this._workers;
+  }
+
+  getWorkerById(workerId: number): WorkerInterface {
+    return this._workers[workerId];
+  }
+
+  createWorker(workerOptions: WorkerOptions): WorkerInterface {
+    return workerOptions.useNodeWorkersIfPossible
+      ? new NodeThreadsWorker(workerOptions)
+      : new ChildProcessWorker(workerOptions);
+  }
+
+  end(): void {
+    // We do not cache the request object here. If so, it would only be only
+    // processed by one of the workers, and we want them all to close.
+    for (let i = 0; i < this._workers.length; i++) {
+      this._workers[i].send(
+        [CHILD_MESSAGE_END, false],
+        emptyMethod,
+        emptyMethod,
+      );
+    }
+  }
+}