From fc3d0341a0a5d9b1b4b7ecb506f2698e6b15ccd2 Mon Sep 17 00:00:00 2001
From: IIIEII <al.iiieii@gmail.com>
Date: Fri, 25 Feb 2022 16:25:59 +0300
Subject: [PATCH] Eliminate Farm work results leakage by clearing callback
 reference (#12497)

---
 CHANGELOG.md                                  |  1 +
 packages/jest-worker/src/Farm.ts              |  7 +-
 .../src/__tests__/leak-integration.test.ts    | 81 ++++++++++++++-----
 3 files changed, 67 insertions(+), 22 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7af1fa88d4be..b2a3e63449ce 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -53,6 +53,7 @@
 - `[jest-matcher-utils]` Pass maxWidth to `pretty-format` to avoid printing every element in arrays by default ([#12402](https://github.com/facebook/jest/pull/12402))
 - `[jest-mock]` Fix function overloads for `spyOn` to allow more correct type inference in complex object ([#12442](https://github.com/facebook/jest/pull/12442))
 - `[jest-reporters]` Notifications generated by the `--notify` flag are no longer persistent in GNOME Shell. ([#11733](https://github.com/facebook/jest/pull/11733))
+- `[jest-worker]` Fix `Farm` execution results memory leak ([#12497](https://github.com/facebook/jest/pull/12497))
 
 ### Chore & Maintenance
 
diff --git a/packages/jest-worker/src/Farm.ts b/packages/jest-worker/src/Farm.ts
index 9a8dcb20f9b8..e3e1a84870fd 100644
--- a/packages/jest-worker/src/Farm.ts
+++ b/packages/jest-worker/src/Farm.ts
@@ -127,9 +127,12 @@ export default class Farm {
     // 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;
+    let taskOnEnd: OnEnd | null = task.onEnd;
     const onEnd: OnEnd = (error, result) => {
-      taskOnEnd(error, result);
+      if (taskOnEnd) {
+        taskOnEnd(error, result);
+      }
+      taskOnEnd = null;
 
       this._unlock(workerId);
       this._process(workerId);
diff --git a/packages/jest-worker/src/__tests__/leak-integration.test.ts b/packages/jest-worker/src/__tests__/leak-integration.test.ts
index f8e440e2d904..af5e827bc296 100644
--- a/packages/jest-worker/src/__tests__/leak-integration.test.ts
+++ b/packages/jest-worker/src/__tests__/leak-integration.test.ts
@@ -11,30 +11,71 @@ import {writeFileSync} from 'graceful-fs';
 import LeakDetector from 'jest-leak-detector';
 import {Worker} from '../../build/index';
 
-let workerFile!: string;
-beforeAll(() => {
-  workerFile = join(tmpdir(), 'baz.js');
-  writeFileSync(workerFile, 'module.exports.fn = () => {};');
-});
+describe('WorkerThreads leaks', () => {
+  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'],
+  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, reject) => {
+      const obj = {};
+      leakDetector = new LeakDetector(obj);
+      (worker as any).fn(obj).then(resolve, reject);
+    });
+
+    expect(await leakDetector.isLeaking()).toBe(false);
   });
 });
-afterEach(async () => {
-  await worker.end();
-});
 
-it('does not retain arguments after a task finished', async () => {
-  let leakDetector!: LeakDetector;
-  await new Promise((resolve, reject) => {
-    const obj = {};
-    leakDetector = new LeakDetector(obj);
-    (worker as any).fn(obj).then(resolve, reject);
+describe('Worker leaks', () => {
+  let workerFile!: string;
+  beforeAll(() => {
+    workerFile = join(tmpdir(), 'baz.js');
+    writeFileSync(workerFile, 'module.exports.fn = (obj) => [obj];');
+  });
+
+  let worker!: Worker;
+  beforeEach(() => {
+    worker = new Worker(workerFile, {
+      enableWorkerThreads: false,
+      exposedMethods: ['fn'],
+    });
   });
+  afterEach(async () => {
+    await worker.end();
+  });
+
+  it('does not retain result after next task call', async () => {
+    let leakDetector!: LeakDetector;
+    await new Promise((resolve, reject) => {
+      const obj = {};
+      (worker as any)
+        .fn(obj)
+        .then((result: unknown) => {
+          leakDetector = new LeakDetector(result);
+          return result;
+        })
+        .then(resolve, reject);
+    });
+    await new Promise((resolve, reject) => {
+      const obj = {};
+      (worker as any).fn(obj).then(resolve, reject);
+    });
 
-  expect(await leakDetector.isLeaking()).toBe(false);
+    expect(await leakDetector.isLeaking()).toBe(false);
+  });
 });