Dramatically improve watch mode performance. (#8201)

## Summary

Resolves #7341 

This PR dramatically improves watch mode performance, bringing it in line with single run mode performance. It accomplishes that by:
- Workers previously initialized a new `ModuleMap` and `Resolver` for every test in watch mode. Now, those objects are only initialized once when the worker is setup.
- In the main thread, caching the conversion of `ModuleMap` to a JSON-friendly object.
- Allowing watch mode to use the same number of CPUs as single run mode.

## Benchmarks

I benchmarked against Jest's own test suite, excluding e2e tests which don't provide good signal because they individually take a long time (so startup time for the test is marginalized). The numbers show that running in Watch mode previously added an extra 35%~ of runtime to the tests but that has now been reduced to almost nothing.

Watch mode should now just be paying a one-time initial cost for each worker when the haste map changes instead of paying that same cost for _every_ test run.

### branch: master

`yarn jest ./packages`
Run time: 15.091s

`yarn jest ./packages --watch`
Run time: 23.234s

### branch: watch-performance

`yarn jest ./packages`
Run time: 14.973s

`yarn jest ./packages --watch`
Run time: 15.196s


## Test plan

- All tests pass.
- Benchmarked to verify the performance wins.
- Verified that when the haste map is updated, the update is propagated out to all workers.

CHANGELOG.md

@@ -36,6 +36,7 @@
 - `[jest-haste-map]` Avoid persisting haste map or processing files when not changed ([#8153](https://github.com/facebook/jest/pull/8153))
 - `[jest-core]` Improve performance of SearchSource.findMatchingTests by 15% ([#8184](https://github.com/facebook/jest/pull/8184))
 - `[jest-resolve]` Optimize internal cache lookup performance ([#8183](https://github.com/facebook/jest/pull/8183))
+- `[jest-core]` Dramatically improve watch mode performance ([#8201](https://github.com/facebook/jest/pull/8201))
 
 ## 24.5.0
 

packages/jest-config/src/__tests__/getMaxWorkers.test.ts

@@ -32,7 +32,7 @@ describe('getMaxWorkers', () => {
 
   it('Returns based on the number of cpus', () => {
     expect(getMaxWorkers({})).toBe(3);
-    expect(getMaxWorkers({watch: true})).toBe(2);
+    expect(getMaxWorkers({watch: true})).toBe(3);
   });
 
   describe('% based', () => {

packages/jest-config/src/getMaxWorkers.ts

@@ -32,6 +32,6 @@ export default function getMaxWorkers(
     return parsed > 0 ? parsed : 1;
   } else {
     const cpus = os.cpus() ? os.cpus().length : 1;
-    return Math.max(argv.watch ? Math.floor(cpus / 2) : cpus - 1, 1);
+    return Math.max(cpus - 1, 1);
   }
 }

packages/jest-haste-map/src/ModuleMap.ts

@@ -32,8 +32,9 @@ export type SerializableModuleMap = {
 };
 
 export default class ModuleMap {
-  private readonly _raw: RawModuleMap;
   static DuplicateHasteCandidatesError: typeof DuplicateHasteCandidatesError;
+  private readonly _raw: RawModuleMap;
+  private json: SerializableModuleMap | undefined;
 
   constructor(raw: RawModuleMap) {
     this._raw = raw;
@@ -84,12 +85,15 @@ export default class ModuleMap {
   }
 
   toJSON(): SerializableModuleMap {
-    return {
-      duplicates: Array.from(this._raw.duplicates),
-      map: Array.from(this._raw.map),
-      mocks: Array.from(this._raw.mocks),
-      rootDir: this._raw.rootDir,
-    };
+    if (!this.json) {
+      this.json = {
+        duplicates: Array.from(this._raw.duplicates),
+        map: Array.from(this._raw.map),
+        mocks: Array.from(this._raw.mocks),
+        rootDir: this._raw.rootDir,
+      };
+    }
+    return this.json;
   }
 
   static fromJSON(serializableModuleMap: SerializableModuleMap) {

packages/jest-runner/src/__tests__/testRunner.test.js

@@ -51,7 +51,6 @@ test('injects the serializable module map into each worker in watch mode', () =>
             context: runContext,
             globalConfig,
             path: './file.test.js',
-            serializableModuleMap,
           },
         ],
         [
@@ -60,7 +59,6 @@ test('injects the serializable module map into each worker in watch mode', () =>
             context: runContext,
             globalConfig,
             path: './file2.test.js',
-            serializableModuleMap,
           },
         ],
       ]);
@@ -90,7 +88,6 @@ test('does not inject the serializable module map in serial mode', () => {
             context: runContext,
             globalConfig,
             path: './file.test.js',
-            serializableModuleMap: null,
           },
         ],
         [
@@ -99,7 +96,6 @@ test('does not inject the serializable module map in serial mode', () => {
             context: runContext,
             globalConfig,
             path: './file2.test.js',
-            serializableModuleMap: null,
           },
         ],
       ]);

packages/jest-runner/src/index.ts

@@ -11,7 +11,7 @@ import exit from 'exit';
 import throat from 'throat';
 import Worker from 'jest-worker';
 import runTest from './runTest';
-import {worker} from './testWorker';
+import {worker, SerializableResolver} from './testWorker';
 import {
   OnTestFailure,
   OnTestStart,
@@ -103,11 +103,31 @@ class TestRunner {
     onResult: OnTestSuccess,
     onFailure: OnTestFailure,
   ) {
+    let resolvers: Map<string, SerializableResolver> | undefined = undefined;
+    if (watcher.isWatchMode()) {
+      resolvers = new Map();
+      for (const test of tests) {
+        if (!resolvers.has(test.context.config.name)) {
+          resolvers.set(test.context.config.name, {
+            config: test.context.config,
+            serializableModuleMap: test.context.moduleMap.toJSON(),
+          });
+        }
+      }
+    }
+
     const worker = new Worker(TEST_WORKER_PATH, {
       exposedMethods: ['worker'],
       forkOptions: {stdio: 'pipe'},
       maxRetries: 3,
       numWorkers: this._globalConfig.maxWorkers,
+      setupArgs: resolvers
+        ? [
+            {
+              serializableResolvers: Array.from(resolvers.values()),
+            },
+          ]
+        : undefined,
     }) as WorkerInterface;
 
     if (worker.getStdout()) worker.getStdout().pipe(process.stdout);
@@ -135,9 +155,6 @@ class TestRunner {
           },
           globalConfig: this._globalConfig,
           path: test.path,
-          serializableModuleMap: watcher.isWatchMode()
-            ? test.context.moduleMap.toJSON()
-            : null,
         });
       });
 

packages/jest-runner/src/testWorker.ts

@@ -8,18 +8,23 @@
 
 import {Config} from '@jest/types';
 import {SerializableError, TestResult} from '@jest/test-result';
-import HasteMap, {SerializableModuleMap, ModuleMap} from 'jest-haste-map';
+import HasteMap, {ModuleMap, SerializableModuleMap} from 'jest-haste-map';
 import exit from 'exit';
 import {separateMessageFromStack} from 'jest-message-util';
 import Runtime from 'jest-runtime';
+import Resolver from 'jest-resolve';
 import {ErrorWithCode, TestRunnerSerializedContext} from './types';
 import runTest from './runTest';
 
+export type SerializableResolver = {
+  config: Config.ProjectConfig;
+  serializableModuleMap: SerializableModuleMap;
+};
+
 type WorkerData = {
   config: Config.ProjectConfig;
   globalConfig: Config.GlobalConfig;
   path: Config.Path;
-  serializableModuleMap: SerializableModuleMap | null;
   context?: TestRunnerSerializedContext;
 };
 
@@ -47,45 +52,50 @@ const formatError = (error: string | ErrorWithCode): SerializableError => {
   };
 };
 
-const resolvers = Object.create(null);
-const getResolver = (
-  config: Config.ProjectConfig,
-  moduleMap: ModuleMap | null,
-) => {
-  // In watch mode, the raw module map with all haste modules is passed from
-  // the test runner to the watch command. This is because jest-haste-map's
-  // watch mode does not persist the haste map on disk after every file change.
-  // To make this fast and consistent, we pass it from the TestRunner.
-  if (moduleMap) {
-    return Runtime.createResolver(config, moduleMap);
-  } else {
-    const name = config.name;
-    if (!resolvers[name]) {
-      resolvers[name] = Runtime.createResolver(
+const resolvers = new Map<string, Resolver>();
+const getResolver = (config: Config.ProjectConfig, moduleMap?: ModuleMap) => {
+  const name = config.name;
+  if (moduleMap || !resolvers.has(name)) {
+    resolvers.set(
+      name,
+      Runtime.createResolver(
         config,
-        Runtime.createHasteMap(config).readModuleMap(),
-      );
-    }
-    return resolvers[name];
+        moduleMap || Runtime.createHasteMap(config).readModuleMap(),
+      ),
+    );
   }
+  return resolvers.get(name)!;
 };
 
+export function setup(setupData?: {
+  serializableResolvers: Array<SerializableResolver>;
+}) {
+  // Setup data is only used in watch mode to pass the latest version of all
+  // module maps that will be used during the test runs. Otherwise, module maps
+  // are loaded from disk as needed.
+  if (setupData) {
+    for (const {
+      config,
+      serializableModuleMap,
+    } of setupData.serializableResolvers) {
+      const moduleMap = HasteMap.ModuleMap.fromJSON(serializableModuleMap);
+      getResolver(config, moduleMap);
+    }
+  }
+}
+
 export async function worker({
   config,
   globalConfig,
   path,
-  serializableModuleMap,
   context,
 }: WorkerData): Promise<TestResult> {
   try {
-    const moduleMap = serializableModuleMap
-      ? HasteMap.ModuleMap.fromJSON(serializableModuleMap)
-      : null;
     return await runTest(
       path,
       globalConfig,
       config,
-      getResolver(config, moduleMap),
+      getResolver(config),
       context && {
         ...context,
         changedFiles: context.changedFiles && new Set(context.changedFiles),