Fix memory leak (all solutions combined)

Author: eyalroth

packages/jest-circus/src/state.ts

@@ -6,6 +6,7 @@
  */
 
 import type {Circus, Global} from '@jest/types';
+import {setGlobal, setNotShreddable} from 'jest-util';
 import eventHandler from './eventHandler';
 import formatNodeAssertErrors from './formatNodeAssertErrors';
 import {STATE_SYM} from './types';
@@ -39,16 +40,28 @@ const createState = (): Circus.State => {
 };
 
 /* eslint-disable no-restricted-globals */
+export const getState = (): Circus.State =>
+  (global as Global.Global)[STATE_SYM] as Circus.State;
+export const setState = (state: Circus.State): Circus.State => {
+  setGlobal(global, STATE_SYM, state);
+  setNotShreddable(state, [
+    'hasFocusedTests',
+    'hasStarted',
+    'includeTestLocationInResult',
+    'maxConcurrency',
+    'seed',
+    'testNamePattern',
+    'testTimeout',
+    'unhandledErrors',
+    'unhandledRejectionErrorByPromise',
+  ]);
+  return state;
+};
 export const resetState = (): void => {
-  (global as Global.Global)[STATE_SYM] = createState();
+  setState(createState());
 };
 
 resetState();
-
-export const getState = (): Circus.State =>
-  (global as Global.Global)[STATE_SYM] as Circus.State;
-export const setState = (state: Circus.State): Circus.State =>
-  ((global as Global.Global)[STATE_SYM] = state);
 /* eslint-enable */
 
 export const dispatch = async (event: Circus.AsyncEvent): Promise<void> => {

packages/jest-environment-node/src/index.ts

@@ -14,7 +14,12 @@ import type {
 import {LegacyFakeTimers, ModernFakeTimers} from '@jest/fake-timers';
 import type {Global} from '@jest/types';
 import {ModuleMocker} from 'jest-mock';
-import {installCommonGlobals} from 'jest-util';
+import {
+  installCommonGlobals,
+  isShreddable,
+  setNotShreddable,
+  shred,
+} from 'jest-util';
 
 type Timer = {
   id: number;
@@ -80,12 +85,13 @@ export default class NodeEnvironment implements JestEnvironment<Timer> {
   moduleMocker: ModuleMocker | null;
   customExportConditions = ['node', 'node-addons'];
   private readonly _configuredExportConditions?: Array<string>;
+  private _globalProxy: GlobalProxy;
 
   // while `context` is unused, it should always be passed
   constructor(config: JestEnvironmentConfig, _context: EnvironmentContext) {
     const {projectConfig} = config;
-    this.context = createContext();
-
+    this._globalProxy = new GlobalProxy();
+    this.context = createContext(this._globalProxy.proxy());
     const global = runInContext(
       'this',
       Object.assign(this.context, projectConfig.testEnvironmentOptions),
@@ -194,6 +200,8 @@ export default class NodeEnvironment implements JestEnvironment<Timer> {
       config: projectConfig,
       global,
     });
+
+    this._globalProxy.envSetupCompleted();
   }
 
   // eslint-disable-next-line @typescript-eslint/no-empty-function
@@ -206,9 +214,29 @@ export default class NodeEnvironment implements JestEnvironment<Timer> {
     if (this.fakeTimersModern) {
       this.fakeTimersModern.dispose();
     }
-    this.context = null;
+
+    if (this.context) {
+      // source-map-support keeps memory leftovers in `Error.prepareStackTrace`
+      runInContext("Error.prepareStackTrace = () => '';", this.context);
+
+      // remove any leftover listeners that may hold references to sizable memory
+      this.context.process.removeAllListeners();
+      const cluster = runInContext(
+        "require('node:cluster')",
+        Object.assign(this.context, {
+          require:
+            // get native require instead of webpack's
+            // @ts-expect-error https://webpack.js.org/api/module-variables/#__non_webpack_require__-webpack-specific
+            __non_webpack_require__,
+        }),
+      );
+      cluster.removeAllListeners();
+
+      this.context = null;
+    }
     this.fakeTimers = null;
     this.fakeTimersModern = null;
+    this._globalProxy.clear();
   }
 
   exportConditions(): Array<string> {
@@ -221,3 +249,116 @@ export default class NodeEnvironment implements JestEnvironment<Timer> {
 }
 
 export const TestEnvironment = NodeEnvironment;
+
+/**
+ * Creates a new empty global object and wraps it with a {@link Proxy}.
+ *
+ * The purpose is to register any property set on the global object,
+ * and {@link #shred} them at environment teardown, to clean up memory and
+ * prevent leaks.
+ */
+class GlobalProxy implements ProxyHandler<typeof globalThis> {
+  private global: typeof globalThis = Object.create(
+    Object.getPrototypeOf(globalThis),
+  );
+  private globalProxy: typeof globalThis = new Proxy(this.global, this);
+  private isEnvSetup = false;
+  private propertyToValue = new Map<string | symbol, unknown>();
+  private leftovers: Array<{property: string | symbol; value: unknown}> = [];
+
+  constructor() {
+    this.register = this.register.bind(this);
+  }
+
+  proxy(): typeof globalThis {
+    return this.globalProxy;
+  }
+
+  /**
+   * Marks that the environment setup has completed, and properties set on
+   * the global object from now on should be shredded at teardown.
+   */
+  envSetupCompleted(): void {
+    this.isEnvSetup = true;
+  }
+
+  /**
+   * Shreds any property that was set on the global object, except for:
+   * 1. Properties that were set before {@link #envSetupCompleted} was invoked.
+   * 2. Properties protected by {@link #setNotShreddable}.
+   */
+  clear(): void {
+    for (const {property, value} of [
+      ...[...this.propertyToValue.entries()].map(([property, value]) => ({
+        property,
+        value,
+      })),
+      ...this.leftovers,
+    ]) {
+      /*
+       * react-native invoke its custom `performance` property after env teardown.
+       * its setup file should use `setNotShreddable` to prevent this.
+       */
+      if (property !== 'performance') {
+        shred(value);
+      }
+    }
+    this.propertyToValue.clear();
+    this.leftovers = [];
+    this.global = {} as typeof globalThis;
+    this.globalProxy = {} as typeof globalThis;
+  }
+
+  defineProperty(
+    target: typeof globalThis,
+    property: string | symbol,
+    attributes: PropertyDescriptor,
+  ): boolean {
+    const newAttributes = {...attributes};
+
+    if ('set' in newAttributes && newAttributes.set !== undefined) {
+      const originalSet = newAttributes.set;
+      const register = this.register;
+      newAttributes.set = value => {
+        originalSet(value);
+        const newValue = Reflect.get(target, property);
+        register(property, newValue);
+      };
+    }
+
+    const result = Reflect.defineProperty(target, property, newAttributes);
+
+    if ('value' in newAttributes) {
+      this.register(property, newAttributes.value);
+    }
+
+    return result;
+  }
+
+  deleteProperty(
+    target: typeof globalThis,
+    property: string | symbol,
+  ): boolean {
+    const result = Reflect.deleteProperty(target, property);
+    const value = this.propertyToValue.get(property);
+    if (value) {
+      this.leftovers.push({property, value});
+      this.propertyToValue.delete(property);
+    }
+    return result;
+  }
+
+  private register(property: string | symbol, value: unknown) {
+    const currentValue = this.propertyToValue.get(property);
+    if (value !== currentValue) {
+      if (!this.isEnvSetup && isShreddable(value)) {
+        setNotShreddable(value);
+      }
+      if (currentValue) {
+        this.leftovers.push({property, value: currentValue});
+      }
+
+      this.propertyToValue.set(property, value);
+    }
+  }
+}

packages/jest-repl/src/cli/runtime-cli.ts

@@ -99,9 +99,9 @@ export async function run(
       },
       {console: customConsole, docblockPragmas: {}, testPath: filePath},
     );
-    setGlobal(environment.global, 'console', customConsole);
-    setGlobal(environment.global, 'jestProjectConfig', projectConfig);
-    setGlobal(environment.global, 'jestGlobalConfig', globalConfig);
+    setGlobal(environment.global, 'console', customConsole, false);
+    setGlobal(environment.global, 'jestProjectConfig', projectConfig, false);
+    setGlobal(environment.global, 'jestGlobalConfig', globalConfig, false);
 
     const runtime = new Runtime(
       projectConfig,

packages/jest-runner/src/runTest.ts

@@ -183,7 +183,7 @@ async function runTestInternal(
     ? new LeakDetector(environment)
     : null;
 
-  setGlobal(environment.global, 'console', testConsole);
+  setGlobal(environment.global, 'console', testConsole, false);
 
   const runtime = new Runtime(
     projectConfig,
@@ -312,8 +312,12 @@ async function runTestInternal(
         sendMessageToJest,
       );
     } catch (error: any) {
-      // Access stack before uninstalling sourcemaps
-      error.stack;
+      // Access all stacks before uninstalling sourcemaps
+      let e = error;
+      while (typeof e === 'object' && 'stack' in e) {
+        e.stack;
+        e = e?.cause ?? {};
+      }
 
       throw error;
     } finally {

packages/jest-util/src/index.ts

@@ -29,3 +29,4 @@ export {default as tryRealpath} from './tryRealpath';
 export {default as requireOrImportModule} from './requireOrImportModule';
 export {default as invariant} from './invariant';
 export {default as isNonNullable} from './isNonNullable';
+export {isShreddable, setNotShreddable, shred} from './shredder';

packages/jest-util/src/setGlobal.ts

@@ -6,12 +6,16 @@
  */
 
 import type {Global} from '@jest/types';
+import {isShreddable, setNotShreddable} from './shredder';
 
 export default function setGlobal(
   globalToMutate: typeof globalThis | Global.Global,
-  key: string,
+  key: string | symbol,
   value: unknown,
+  shredAfterTeardown = true,
 ): void {
-  // @ts-expect-error: no index
-  globalToMutate[key] = value;
+  Reflect.set(globalToMutate, key, value);
+  if (!shredAfterTeardown && isShreddable(value)) {
+    setNotShreddable(value);
+  }
 }

packages/jest-util/src/shredder.ts

@@ -0,0 +1,54 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+const NO_SHRED_AFTER_TEARDOWN = Symbol.for('$$jest-no-shred');
+
+/**
+ * Deletes all the properties from the given value (if it's an object),
+ * unless the value was protected via {@link #setNotShreddable}.
+ *
+ * @param value the given value.
+ */
+export function shred(value: unknown): void {
+  if (isShreddable(value)) {
+    const protectedProperties = Reflect.get(value, NO_SHRED_AFTER_TEARDOWN);
+    if (!Array.isArray(protectedProperties) || protectedProperties.length > 0) {
+      for (const key of Reflect.ownKeys(value)) {
+        if (!protectedProperties?.includes(key)) {
+          Reflect.deleteProperty(value, key);
+        }
+      }
+    }
+  }
+}
+
+/**
+ * Protects the given value from being shredded by {@link #shred}.
+ *
+ * @param value The given value.
+ * @param properties If the array contains any property,
+ * then only these properties will not be deleted; otherwise if the array is empty,
+ * all properties will not be deleted.
+ */
+export function setNotShreddable<T extends object>(
+  value: T,
+  properties: Array<keyof T> = [],
+): boolean {
+  if (isShreddable(value)) {
+    return Reflect.set(value, NO_SHRED_AFTER_TEARDOWN, properties);
+  }
+  return false;
+}
+
+/**
+ * Whether the given value is possible to be shredded.
+ *
+ * @param value The given value.
+ */
+export function isShreddable(value: unknown): value is object {
+  return value !== null && ['object', 'function'].includes(typeof value);
+}