Support tagging memory leaks in js-under-test

Summary: **How it works:** 1. In web app, we know at certain points some objects are no longer needed to be kept alive in memory. For those objects at the right time, we can "tag" or "mark" those objects that should not be kept alive. Those targged objects will be added to a special group of `WeakSet` in browser memory after their React components are released. 2. MemLab will examine the heap snapshots and mark all such objects as memory leaks. NOTE: In MemLab, we control when GC will be triggered, so all objects in the heap snapshots are guaranteed to survive the GC, which means all tagged objects tracked by the logic added in this diff are memory leaks. Differential Revision: D50312753 Privacy Context Container: L1212261 fbshipit-source-id: fc85bde625c06ef4a5a3626a3d7e0ca0ce12ebd3
facebook · Nov 16, 2023 · 040315e · 040315e
1 parent 80fd71b
commit 040315e
Show file tree

Hide file tree

Showing 15 changed files with 255 additions and 17 deletions.
diff --git a/packages/api/src/__tests__/API/E2EFindTaggedMemoryLeaks.test.ts b/packages/api/src/__tests__/API/E2EFindTaggedMemoryLeaks.test.ts
@@ -0,0 +1,101 @@
+/**
+ * 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.
+ *
+ * @format
+ * @oncall web_perf_infra
+ */
+
+/* eslint-disable @typescript-eslint/ban-ts-comment */
+
+import type {Page} from 'puppeteer';
+import type {IScenario} from '@memlab/core';
+
+import os from 'os';
+import path from 'path';
+import fs from 'fs-extra';
+import {run} from '../../index';
+import {testSetup, testTimeout} from './lib/E2ETestSettings';
+
+beforeEach(testSetup);
+
+// eslint-disable-next-line @typescript-eslint/ban-types
+type Objectish = object | Function;
+
+// The structure of classes and objects should be fixed
+// so that MemLab can analyze them correctly in heap.
+type TrackedItem = {
+  useCaseId: string;
+  taggedObjects: WeakSet<Objectish>;
+};
+
+function tagObjectsAsLeaks() {
+  // this class definition must be defined within `tagObjectsAsLeaks`
+  // since this function will be serialized and executed in browser context
+  class MemLabTracker {
+    memlabIdentifier: string;
+    tagToTrackedObjectsMap: Map<string, TrackedItem>;
+
+    constructor() {
+      this.memlabIdentifier = 'MemLabObjectTracker';
+      this.tagToTrackedObjectsMap = new Map();
+    }
+
+    tag(target: Objectish, useCaseId = 'MEMLAB_TAGGED'): void {
+      let trackedItems = this.tagToTrackedObjectsMap.get(useCaseId);
+      if (!trackedItems) {
+        trackedItems = {
+          useCaseId,
+          taggedObjects: new WeakSet(),
+        };
+        this.tagToTrackedObjectsMap.set(useCaseId, trackedItems);
+      }
+      trackedItems.taggedObjects.add(target);
+    }
+  }
+
+  // @ts-ignore
+  window.injectHookForLink4 = () => {
+    // @ts-ignore
+    const tracker = (window._memlabTrack = new MemLabTracker());
+    const leaks: Array<{x: number}> = [];
+    // @ts-ignore
+    window._taggedLeaks = leaks;
+    for (let i = 0; i < 11; ++i) {
+      leaks.push({x: i});
+    }
+    leaks.forEach(item => {
+      tracker.tag(item);
+    });
+  };
+}
+
+test(
+  'tagged leak objects can be identified as leaks',
+  async () => {
+    const selfDefinedScenario: IScenario = {
+      app: (): string => 'test-spa',
+      url: (): string => '',
+      action: async (page: Page): Promise<void> =>
+        await page.click('[data-testid="link-4"]'),
+    };
+
+    const workDir = path.join(os.tmpdir(), 'memlab-api-test', `${process.pid}`);
+    fs.mkdirsSync(workDir);
+
+    const result = await run({
+      scenario: selfDefinedScenario,
+      evalInBrowserAfterInitLoad: tagObjectsAsLeaks,
+      workDir,
+    });
+    // tagged objects should be detected as leaks
+    expect(result.leaks.length).toBe(1);
+    // expect all traces are found
+    expect(
+      result.leaks.some(leak => JSON.stringify(leak).includes('_taggedLeaks')),
+    );
+  },
+  testTimeout,
+);
diff --git a/packages/core/src/lib/HeapAnalyzer.ts b/packages/core/src/lib/HeapAnalyzer.ts
@@ -433,6 +433,7 @@ class MemoryAnalyst {
     }
 
     const leakFilter = new LeakObjectFilter();
+    leakFilter.beforeFiltering(config, snapshot, leakedNodeIds);
     // start filtering memory leaks
     utils.filterNodesInPlace(leakedNodeIds, snapshot, node =>
       leakFilter.filter(config, node, snapshot, leakedNodeIds),

diff --git a/packages/core/src/lib/Utils.ts b/packages/core/src/lib/Utils.ts
@@ -1675,7 +1675,10 @@ function isMeaningfulEdge(
     const index = edgeNameOrIndex;
     // only elements at particular indexes of (map descriptors) are holding
     // representative references to objects.
-    if (index >= 2 || (typeof index === 'number' && index % 3 === 1)) {
+    if (parseInt(index.toString(), 10) >= 2) {
+      return false;
+    }
+    if (typeof index === 'number' && index % 3 === 1) {
       return false;
     }
   }
@@ -2111,6 +2114,10 @@ function tryToMutePuppeteerWarning() {
   process.env['PUPPETEER_DISABLE_HEADLESS_WARNING'] = '1';
 }
 
+function isStandardNumberToString(input: string): boolean {
+  return parseInt(input, 10).toString() === input;
+}
+
 export default {
   aggregateDominatorMetrics,
   applyToNodes,
@@ -2191,6 +2198,7 @@ export default {
   isRootNode,
   isSlicedStringNode,
   isStackTraceFrame,
+  isStandardNumberToString,
   isStringNode,
   isURLEqual,
   isWeakMapEdge,

diff --git a/packages/core/src/lib/leak-filters/BaseLeakFilter.rule.ts b/packages/core/src/lib/leak-filters/BaseLeakFilter.rule.ts
@@ -22,10 +22,38 @@ export enum LeakDecision {
 }
 
 export interface ILeakObjectFilterRule {
+  beforeFiltering(
+    config: MemLabConfig,
+    snapshot: IHeapSnapshot,
+    leakedNodeIds: HeapNodeIdSet,
+  ): void;
+
   filter(
     config: MemLabConfig,
     node: IHeapNode,
     snapshot: IHeapSnapshot,
     leakedNodeIds: HeapNodeIdSet,
   ): LeakDecision;
 }
+
+export abstract class LeakObjectFilterRuleBase
+  implements ILeakObjectFilterRule
+{
+  beforeFiltering(
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _config: MemLabConfig,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _snapshot: IHeapSnapshot,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _leakedNodeIds: HeapNodeIdSet,
+  ): void {
+    // do nothing by default
+  }
+
+  abstract filter(
+    config: MemLabConfig,
+    node: IHeapNode,
+    snapshot: IHeapSnapshot,
+    leakedNodeIds: HeapNodeIdSet,
+  ): LeakDecision;
+}
diff --git a/packages/core/src/lib/leak-filters/LeakFilterRuleList.ts b/packages/core/src/lib/leak-filters/LeakFilterRuleList.ts
@@ -19,8 +19,10 @@ import {FilterStackTraceFrameRule} from './rules/FilterStackTraceFrame.rule';
 import {FilterTrivialNodeRule} from './rules/FilterTrivialNode.rule';
 import {FilterUnmountedFiberNodeRule} from './rules/FilterUnmountedFiberNode.rule';
 import {FilterXMLHTTPRequestRule} from './rules/FilterXMLHTTPRequest.rule';
+import {FilterUserTaggedLeaksRule} from './rules/FilterUserTaggedLeaks.rule';
 
 const list: ILeakObjectFilterRule[] = [
+  new FilterUserTaggedLeaksRule(),
   new FilterByExternalFilterRule(),
   new FilterTrivialNodeRule(),
   new FilterHermesNodeRule(),

diff --git a/packages/core/src/lib/leak-filters/LeakObjectFilter.ts b/packages/core/src/lib/leak-filters/LeakObjectFilter.ts
@@ -18,6 +18,15 @@ import rules from './LeakFilterRuleList';
  * if an object is a memory leak or not
  */
 export class LeakObjectFilter {
+  beforeFiltering(
+    config: MemLabConfig,
+    snapshot: IHeapSnapshot,
+    leakedNodeIds: HeapNodeIdSet,
+  ): void {
+    for (const rule of rules) {
+      rule.beforeFiltering(config, snapshot, leakedNodeIds);
+    }
+  }
   public filter(
     config: MemLabConfig,
     node: IHeapNode,

diff --git a/packages/core/src/lib/leak-filters/rules/FilterByExternalFilter.rule.ts b/packages/core/src/lib/leak-filters/rules/FilterByExternalFilter.rule.ts
@@ -11,12 +11,12 @@
 import type {MemLabConfig} from '../../Config';
 import type {HeapNodeIdSet, IHeapNode, IHeapSnapshot} from '../../Types';
 
-import {ILeakObjectFilterRule, LeakDecision} from '../BaseLeakFilter.rule';
+import {LeakDecision, LeakObjectFilterRuleBase} from '../BaseLeakFilter.rule';
 
 /**
  * filter memory leaks defined by external leak filter
  */
-export class FilterByExternalFilterRule implements ILeakObjectFilterRule {
+export class FilterByExternalFilterRule extends LeakObjectFilterRuleBase {
   filter(
     config: MemLabConfig,
     node: IHeapNode,

diff --git a/packages/core/src/lib/leak-filters/rules/FilterDetachedDOMElement.rule.ts b/packages/core/src/lib/leak-filters/rules/FilterDetachedDOMElement.rule.ts
@@ -10,13 +10,13 @@
 
 import type {MemLabConfig} from '../../Config';
 import type {IHeapNode} from '../../Types';
-import {ILeakObjectFilterRule, LeakDecision} from '../BaseLeakFilter.rule';
+import {LeakDecision, LeakObjectFilterRuleBase} from '../BaseLeakFilter.rule';
 import utils from '../../Utils';
 
 /**
  * mark detached DOM elements as memory leaks
  */
-export class FilterDetachedDOMElementRule implements ILeakObjectFilterRule {
+export class FilterDetachedDOMElementRule extends LeakObjectFilterRuleBase {
   filter(_config: MemLabConfig, node: IHeapNode): LeakDecision {
     const isDetached = utils.isDetachedDOMNode(node, {
       ignoreInternalNode: true,

diff --git a/packages/core/src/lib/leak-filters/rules/FilterHermesNode.rule.ts b/packages/core/src/lib/leak-filters/rules/FilterHermesNode.rule.ts
@@ -10,10 +10,10 @@
 
 import type {MemLabConfig} from '../../Config';
 import type {IHeapNode} from '../../Types';
-import {ILeakObjectFilterRule, LeakDecision} from '../BaseLeakFilter.rule';
+import {LeakDecision, LeakObjectFilterRuleBase} from '../BaseLeakFilter.rule';
 import utils from '../../Utils';
 
-export class FilterHermesNodeRule implements ILeakObjectFilterRule {
+export class FilterHermesNodeRule extends LeakObjectFilterRuleBase {
   public filter(config: MemLabConfig, node: IHeapNode): LeakDecision {
     // when analyzing hermes heap snapshots, filter Hermes internal objects
     if (config.jsEngine === 'hermes' && utils.isHermesInternalObject(node)) {

diff --git a/packages/core/src/lib/leak-filters/rules/FilterOverSizedNodeAsLeak.rule.ts b/packages/core/src/lib/leak-filters/rules/FilterOverSizedNodeAsLeak.rule.ts
@@ -13,12 +13,12 @@ import type {IHeapNode} from '../../Types';
 
 import utils from '../../Utils';
 import {TraceObjectMode} from '../../Config';
-import {ILeakObjectFilterRule, LeakDecision} from '../BaseLeakFilter.rule';
+import {LeakDecision, LeakObjectFilterRuleBase} from '../BaseLeakFilter.rule';
 
 /**
  * trivial nodes are not reported as memory leaks
  */
-export class FilterOverSizedNodeAsLeakRule implements ILeakObjectFilterRule {
+export class FilterOverSizedNodeAsLeakRule extends LeakObjectFilterRuleBase {
   filter(config: MemLabConfig, node: IHeapNode): LeakDecision {
     if (config.oversizeObjectAsLeak) {
       // TODO: add support to skip this check

diff --git a/packages/core/src/lib/leak-filters/rules/FilterStackTraceFrame.rule.ts b/packages/core/src/lib/leak-filters/rules/FilterStackTraceFrame.rule.ts
@@ -11,13 +11,13 @@
 import type {MemLabConfig} from '../../Config';
 import type {IHeapNode} from '../../Types';
 
-import {ILeakObjectFilterRule, LeakDecision} from '../BaseLeakFilter.rule';
+import {LeakDecision, LeakObjectFilterRuleBase} from '../BaseLeakFilter.rule';
 import utils from '../../Utils';
 
 /**
  * stack trace frames as memory leaks
  */
-export class FilterStackTraceFrameRule implements ILeakObjectFilterRule {
+export class FilterStackTraceFrameRule extends LeakObjectFilterRuleBase {
   filter(_config: MemLabConfig, node: IHeapNode): LeakDecision {
     return utils.isStackTraceFrame(node)
       ? LeakDecision.LEAK

diff --git a/packages/core/src/lib/leak-filters/rules/FilterTrivialNode.rule.ts b/packages/core/src/lib/leak-filters/rules/FilterTrivialNode.rule.ts
@@ -10,13 +10,13 @@
 
 import type {MemLabConfig} from '../../Config';
 import type {IHeapNode} from '../../Types';
-import {ILeakObjectFilterRule, LeakDecision} from '../BaseLeakFilter.rule';
+import {LeakDecision, LeakObjectFilterRuleBase} from '../BaseLeakFilter.rule';
 import utils from '../../Utils';
 
 /**
  * trivial nodes are not reported as memory leaks
  */
-export class FilterTrivialNodeRule implements ILeakObjectFilterRule {
+export class FilterTrivialNodeRule extends LeakObjectFilterRuleBase {
   filter(config: MemLabConfig, node: IHeapNode): LeakDecision {
     return this.isTrivialNode(node)
       ? LeakDecision.NOT_LEAK

diff --git a/packages/core/src/lib/leak-filters/rules/FilterUnmountedFiberNode.rule.ts b/packages/core/src/lib/leak-filters/rules/FilterUnmountedFiberNode.rule.ts
@@ -10,13 +10,13 @@
 
 import type {MemLabConfig} from '../../Config';
 import type {IHeapNode} from '../../Types';
-import {ILeakObjectFilterRule, LeakDecision} from '../BaseLeakFilter.rule';
+import {LeakDecision, LeakObjectFilterRuleBase} from '../BaseLeakFilter.rule';
 import utils from '../../Utils';
 
 /**
  * mark React FiberNodes without a React Fiber Root as memory leaks
  */
-export class FilterUnmountedFiberNodeRule implements ILeakObjectFilterRule {
+export class FilterUnmountedFiberNodeRule extends LeakObjectFilterRuleBase {
   filter(config: MemLabConfig, node: IHeapNode): LeakDecision {
     if (this.checkDetachedFiberNode(config, node)) {
       return LeakDecision.LEAK;