From a124e2e7b1d3278a74cd5450b16d0450feb7214f Mon Sep 17 00:00:00 2001 From: Eyal Roth Date: Sat, 27 Jul 2024 19:29:45 +0300 Subject: [PATCH] #15215 Reduce memory leak in node env by shredding global properties in teardown --- packages/jest-circus/src/state.ts | 25 +++- packages/jest-environment-node/src/index.ts | 128 +++++++++++++++++++- packages/jest-repl/src/cli/runtime-cli.ts | 6 +- packages/jest-runner/src/runTest.ts | 2 +- packages/jest-util/src/index.ts | 1 + packages/jest-util/src/setGlobal.ts | 10 +- packages/jest-util/src/shredder.ts | 54 +++++++++ 7 files changed, 210 insertions(+), 16 deletions(-) create mode 100644 packages/jest-util/src/shredder.ts diff --git a/packages/jest-circus/src/state.ts b/packages/jest-circus/src/state.ts index deda31560871..053f251fff07 100644 --- a/packages/jest-circus/src/state.ts +++ b/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 => { diff --git a/packages/jest-environment-node/src/index.ts b/packages/jest-environment-node/src/index.ts index 11a36e97801b..a6a0afa6dccb 100644 --- a/packages/jest-environment-node/src/index.ts +++ b/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 { moduleMocker: ModuleMocker | null; customExportConditions = ['node', 'node-addons']; private readonly _configuredExportConditions?: Array; + 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 { 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 { this.context = null; this.fakeTimers = null; this.fakeTimersModern = null; + this._globalProxy.clear(); } exportConditions(): Array { @@ -221,3 +230,116 @@ export default class NodeEnvironment implements JestEnvironment { } 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 { + 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(); + 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); + } + } +} diff --git a/packages/jest-repl/src/cli/runtime-cli.ts b/packages/jest-repl/src/cli/runtime-cli.ts index 5997b040a769..394528f1963e 100644 --- a/packages/jest-repl/src/cli/runtime-cli.ts +++ b/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, diff --git a/packages/jest-runner/src/runTest.ts b/packages/jest-runner/src/runTest.ts index ed9a5ed951bd..3bac869fa096 100644 --- a/packages/jest-runner/src/runTest.ts +++ b/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, diff --git a/packages/jest-util/src/index.ts b/packages/jest-util/src/index.ts index dbfd9025175b..e82663821315 100644 --- a/packages/jest-util/src/index.ts +++ b/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'; diff --git a/packages/jest-util/src/setGlobal.ts b/packages/jest-util/src/setGlobal.ts index 4daa5132ec49..ece978243302 100644 --- a/packages/jest-util/src/setGlobal.ts +++ b/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); + } } diff --git a/packages/jest-util/src/shredder.ts b/packages/jest-util/src/shredder.ts new file mode 100644 index 000000000000..209c6c4a71c7 --- /dev/null +++ b/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( + value: T, + properties: Array = [], +): 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); +}