Mitigate memory leaks in jest-environment-node (#15215)

Author: eyalroth

packages/jest-circus/src/state.ts

@@ -6,16 +6,16 @@
  */
 
 import type {Circus, Global} from '@jest/types';
+import {protectProperties, setGlobal} from 'jest-util';
 import eventHandler from './eventHandler';
 import formatNodeAssertErrors from './formatNodeAssertErrors';
 import {EVENT_HANDLERS, STATE_SYM} from './types';
 import {makeDescribe} from './utils';
 
 const handlers: Array<Circus.EventHandler> = ((globalThis as Global.Global)[
   EVENT_HANDLERS
-] = ((globalThis as Global.Global)[
-  EVENT_HANDLERS
-] as Array<Circus.EventHandler>) || [eventHandler, formatNodeAssertErrors]);
+] as Array<Circus.EventHandler>) || [eventHandler, formatNodeAssertErrors];
+setGlobal(globalThis, EVENT_HANDLERS, handlers, 'retain');
 
 export const ROOT_DESCRIBE_BLOCK_NAME = 'ROOT_DESCRIBE_BLOCK';
 
@@ -39,17 +39,29 @@ const createState = (): Circus.State => {
   };
 };
 
+export const getState = (): Circus.State =>
+  (globalThis as Global.Global)[STATE_SYM] as Circus.State;
+export const setState = (state: Circus.State): Circus.State => {
+  setGlobal(globalThis, STATE_SYM, state);
+  protectProperties(state, [
+    'hasFocusedTests',
+    'hasStarted',
+    'includeTestLocationInResult',
+    'maxConcurrency',
+    'seed',
+    'testNamePattern',
+    'testTimeout',
+    'unhandledErrors',
+    'unhandledRejectionErrorByPromise',
+  ]);
+  return state;
+};
 export const resetState = (): void => {
-  (globalThis as Global.Global)[STATE_SYM] = createState();
+  setState(createState());
 };
 
 resetState();
 
-export const getState = (): Circus.State =>
-  (globalThis as Global.Global)[STATE_SYM] as Circus.State;
-export const setState = (state: Circus.State): Circus.State =>
-  ((globalThis as Global.Global)[STATE_SYM] = state);
-
 export const dispatch = async (event: Circus.AsyncEvent): Promise<void> => {
   for (const handler of handlers) {
     await handler(event, getState());

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 {
+  canDeleteProperties,
+  deleteProperties,
+  installCommonGlobals,
+  protectProperties,
+} 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
@@ -209,6 +217,7 @@ export default class NodeEnvironment implements JestEnvironment<Timer> {
     this.context = null;
     this.fakeTimers = null;
     this.fakeTimersModern = null;
+    this._globalProxy.clear();
   }
 
   exportConditions(): Array<string> {
@@ -221,3 +230,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 #deleteProperties} on 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 deleted at teardown.
+   */
+  envSetupCompleted(): void {
+    this.isEnvSetup = true;
+  }
+
+  /**
+   * Deletes 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 #protectProperties}.
+   */
+  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 `protectProperties` to prevent this.
+       */
+      if (property !== 'performance') {
+        deleteProperties(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 && canDeleteProperties(value)) {
+        protectProperties(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, 'retain');
+    setGlobal(environment.global, 'jestProjectConfig', projectConfig, 'retain');
+    setGlobal(environment.global, 'jestGlobalConfig', globalConfig, 'retain');
 
     const runtime = new Runtime(
       projectConfig,

packages/jest-runner/src/runTest.ts

@@ -184,7 +184,7 @@ async function runTestInternal(
     ? new LeakDetector(environment)
     : null;
 
-  setGlobal(environment.global, 'console', testConsole);
+  setGlobal(environment.global, 'console', testConsole, 'retain');
 
   const runtime = new Runtime(
     projectConfig,

packages/jest-util/src/garbage-collection-utils.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 PROTECT_PROPERTY = Symbol.for('$$jest-protect-from-deletion');
+
+/**
+ * Deletes all the properties from the given value (if it's an object),
+ * unless the value was protected via {@link #protectProperties}.
+ *
+ * @param value the given value.
+ */
+export function deleteProperties(value: unknown): void {
+  if (canDeleteProperties(value)) {
+    const protectedProperties = Reflect.get(value, PROTECT_PROPERTY);
+    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 deleted by {@link #deleteProperties}.
+ *
+ * @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 protectProperties<T extends object>(
+  value: T,
+  properties: Array<keyof T> = [],
+): boolean {
+  if (canDeleteProperties(value)) {
+    return Reflect.set(value, PROTECT_PROPERTY, properties);
+  }
+  return false;
+}
+
+/**
+ * Whether the given value has properties that can be deleted (regardless of protection).
+ *
+ * @param value The given value.
+ */
+export function canDeleteProperties(value: unknown): value is object {
+  return value !== null && ['object', 'function'].includes(typeof value);
+}

packages/jest-util/src/index.ts

@@ -29,3 +29,8 @@ 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 {
+  canDeleteProperties,
+  protectProperties,
+  deleteProperties,
+} from './garbage-collection-utils';

packages/jest-util/src/setGlobal.ts

@@ -6,12 +6,19 @@
  */
 
 import type {Global} from '@jest/types';
+import {
+  canDeleteProperties,
+  protectProperties,
+} from './garbage-collection-utils';
 
 export default function setGlobal(
   globalToMutate: typeof globalThis | Global.Global,
-  key: string,
+  key: string | symbol,
   value: unknown,
+  afterTeardown: 'clean' | 'retain' = 'clean',
 ): void {
-  // @ts-expect-error: no index
-  globalToMutate[key] = value;
+  Reflect.set(globalToMutate, key, value);
+  if (afterTeardown === 'retain' && canDeleteProperties(value)) {
+    protectProperties(value);
+  }
 }