Fix memory leaks of jest-worker (#11187)

Co-authored-by: Tim Seckinger <[email protected]>

Author: shuding

CHANGELOG.md

@@ -103,6 +103,7 @@
 - `[jest-resolve]` Cache reading and parsing of `package.json`s ([#11076](https://github.com/facebook/jest/pull/11076))
 - `[jest-runtime, jest-transform]` share `cacheFS` between runtime and transformer ([#10901](https://github.com/facebook/jest/pull/10901))
 - `[jest-runtime]` Load `chalk` only once per worker ([#10864](https://github.com/facebook/jest/pull/10864))
+- `[jest-worker]` Fix memory leak of previous task arguments while no new task is scheduled ([#11187](https://github.com/facebook/jest/pull/11187))
 
 ## 26.6.3
 

packages/jest-worker/package.json

@@ -22,6 +22,7 @@
     "@types/merge-stream": "^1.1.2",
     "@types/supports-color": "^7.2.0",
     "get-stream": "^6.0.0",
+    "jest-leak-detector": "^27.0.0-next.3",
     "worker-farm": "^1.6.0"
   },
   "engines": {

packages/jest-worker/src/Farm.ts

@@ -66,7 +66,14 @@ export default class Farm {
     };
 
     const promise: PromiseWithCustomMessage<unknown> = new Promise(
-      (resolve, reject) => {
+      // Bind args to this function so it won't reference to the parent scope.
+      // This prevents a memory leak in v8, because otherwise the function will
+      // retaine args for the closure.
+      ((
+        args: Array<unknown>,
+        resolve: (value: unknown) => void,
+        reject: (reason?: any) => void,
+      ) => {
         const computeWorkerKey = this._computeWorkerKey;
         const request: ChildMessage = [CHILD_MESSAGE_CALL, false, method, args];
 
@@ -101,7 +108,7 @@ export default class Farm {
         } else {
           this._push(task);
         }
-      },
+      }).bind(null, args),
     );
 
     promise.UNSTABLE_onCustomMessage = addCustomMessageListener;
@@ -124,8 +131,12 @@ export default class Farm {
       throw new Error('Queue implementation returned processed task');
     }
 
+    // Reference the task object outside so it won't be retained by onEnd,
+    // and other properties of the task object, such as task.request can be
+    // garbage collected.
+    const taskOnEnd = task.onEnd;
     const onEnd = (error: Error | null, result: unknown) => {
-      task.onEnd(error, result);
+      taskOnEnd(error, result);
 
       this._unlock(workerId);
       this._process(workerId);

packages/jest-worker/src/__tests__/leak-integration.test.ts

@@ -0,0 +1,40 @@
+/**
+ * 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.
+ */
+
+import {tmpdir} from 'os';
+import {join} from 'path';
+import {writeFileSync} from 'graceful-fs';
+import LeakDetector from 'jest-leak-detector';
+import {Worker} from '../..';
+
+let workerFile!: string;
+beforeAll(() => {
+  workerFile = join(tmpdir(), 'baz.js');
+  writeFileSync(workerFile, `module.exports.fn = () => {};`);
+});
+
+let worker!: Worker;
+beforeEach(() => {
+  worker = new Worker(workerFile, {
+    enableWorkerThreads: true,
+    exposedMethods: ['fn'],
+  });
+});
+afterEach(async () => {
+  await worker.end();
+});
+
+it('does not retain arguments after a task finished', async () => {
+  let leakDetector!: LeakDetector;
+  await new Promise(resolve => {
+    const obj = {};
+    leakDetector = new LeakDetector(obj);
+    (worker as any).fn(obj).then(resolve);
+  });
+
+  expect(await leakDetector.isLeaking()).toBe(false);
+});

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

@@ -206,15 +206,21 @@ export default class ExperimentalWorker implements WorkerInterface {
   send(
     request: ChildMessage,
     onProcessStart: OnStart,
-    onProcessEnd: OnEnd,
+    onProcessEnd: OnEnd | null,
     onCustomMessage: OnCustomMessage,
   ): void {
     onProcessStart(this);
     this._onProcessEnd = (...args) => {
       // Clean the request to avoid sending past requests to workers that fail
       // while waiting for a new request (timers, unhandled rejections...)
       this._request = null;
-      return onProcessEnd(...args);
+
+      const res = onProcessEnd?.(...args);
+
+      // Clean up the reference so related closures can be garbage collected.
+      onProcessEnd = null;
+
+      return res;
     };
 
     this._onCustomMessage = (...arg) => onCustomMessage(...arg);

packages/jest-worker/tsconfig.json

@@ -3,5 +3,6 @@
   "compilerOptions": {
     "rootDir": "src",
     "outDir": "build"
-  }
+  },
+  "references": [{"path": "../jest-leak-detector"}]
 }

yarn.lock

@@ -14137,6 +14137,7 @@ fsevents@^1.2.7:
     "@types/node": "*"
     "@types/supports-color": ^7.2.0
     get-stream: ^6.0.0
+    jest-leak-detector: ^27.0.0-next.3
     merge-stream: ^2.0.0
     supports-color: ^8.0.0
     worker-farm: ^1.6.0