cli(feat): support diffing memory leaks from single heap snapshots

Summary:
Add support for diffing memory leak clusters based on a single heap snapshot from the control or treatment group:

Example: diff heap snapshots from a treatment working directory vs a single heap snapshot file location for a control run
```
memlab diff-leaks --treatment-work-dir /tmp/memlab-test/ --control-snapshot /tmp/memlab-test/data/cur/memlab-plugin.heapsnapshot
```

Or we can use snapshot options only:

```
memlab diff-leaks --treatment-snapshot /tmp/treatment.heapsnapshot --control-snapshot /tmp/control.heapsnapshot
```

Reviewed By: tulga1970

Differential Revision: D52713530

fbshipit-source-id: 03ece7de7ecc8a3f5fb22bc2c5327a234a89eb2c

packages/cli/src/Dispatcher.ts

@@ -9,7 +9,7 @@
  */
 
 import type {ParsedArgs} from 'minimist';
-import type {AnyRecord, MemLabConfig} from '@memlab/core';
+import type {AnyRecord, AnyValue, MemLabConfig, Nullable} from '@memlab/core';
 
 import {config, info} from '@memlab/core';
 import BaseCommand from './BaseCommand';
@@ -83,16 +83,45 @@ export class CommandDispatcher {
     args: ParsedArgs,
   ): Promise<AnyRecord> {
     const options = [...universalOptions, ...command.getOptions()];
-    const configFromOptions = Object.create(null);
+    const configFromOptions: AnyRecord = Object.create(null);
     for (const option of options) {
       const ret = await option.run(config, args);
       if (ret) {
-        Object.assign(configFromOptions, ret);
+        this.mergeConfigFromOptions(configFromOptions, ret);
       }
     }
     return configFromOptions;
   }
 
+  private mergeConfigFromOptions(to: AnyRecord, from: AnyRecord): AnyRecord {
+    for (const key in from) {
+      if (Array.isArray(to[key]) && Array.isArray(from[key])) {
+        // both are arrays, merge them
+        this.mergeArrays(to[key] as AnyValue[], from[key] as AnyValue[]);
+      } else if (from[key] == null || to[key] == null) {
+        // one of them is null, use the other one
+        to[key] = to[key] || from[key];
+      } else {
+        // both have existing values, first one wins
+        info.warning(`Merge conflict CLI options key: ${key}`);
+      }
+    }
+    return to;
+  }
+
+  private mergeArrays(arr1: Nullable<AnyValue[]>, arr2: Nullable<AnyValue[]>) {
+    if (arr1 == null) {
+      return arr2;
+    }
+    if (arr2 == null) {
+      return arr1;
+    }
+    for (const v of arr2) {
+      arr1.push(v);
+    }
+    return arr1;
+  }
+
   private async runCommand(
     command: BaseCommand,
     args: ParsedArgs,

packages/cli/src/commands/heap/DiffLeakCommand.ts

@@ -28,6 +28,8 @@ import SetControlWorkDirOption from '../../options/experiment/SetControlWorkDirO
 import SetTreatmentWorkDirOption from '../../options/experiment/SetTreatmentWorkDirOption';
 import SetMaxClusterSampleSizeOption from '../../options/SetMaxClusterSampleSizeOption';
 import SetTraceContainsFilterOption from '../../options/heap/SetTraceContainsFilterOption';
+import SetControlSnapshotOption from '../../options/experiment/SetControlSnapshotOption';
+import SetTreatmentSnapshotOption from '../../options/experiment/SetTreatmentSnapshotOption';
 
 export type WorkDirSettings = {
   controlWorkDirs: Array<string>;
@@ -74,7 +76,9 @@ export default class CheckLeakCommand extends BaseCommand {
 
   getOptions(): BaseOption[] {
     return [
+      new SetControlSnapshotOption(),
       new SetControlWorkDirOption(),
+      new SetTreatmentSnapshotOption(),
       new SetTreatmentWorkDirOption(),
       new JSEngineOption(),
       new LeakFilterFileOption(),
@@ -90,14 +94,25 @@ export default class CheckLeakCommand extends BaseCommand {
     ];
   }
 
+  protected showWorkingDirErrorMessageAndHalt(): never {
+    throw utils.haltOrThrow(
+      'No control or test working directory or snapshot location specified, ' +
+        'please specify them via: \n' +
+        ` --${new SetControlSnapshotOption().getOptionName()} with ` +
+        ` --${new SetTreatmentSnapshotOption().getOptionName()}\n` +
+        'alternatively, you can also specify them via: \n' +
+        ` --${new SetControlWorkDirOption().getOptionName()} with ` +
+        ` --${new SetTreatmentWorkDirOption().getOptionName()}`,
+    );
+  }
+
   protected getWorkDirs(options: CLIOptions): WorkDirSettings {
     // double check parameters
     if (
       !options.configFromOptions?.controlWorkDirs ||
       !options.configFromOptions?.treatmentWorkDirs
     ) {
-      info.error('Please specify control and test working directory');
-      throw utils.haltOrThrow('No control or test working directory specified');
+      this.showWorkingDirErrorMessageAndHalt();
     }
     // get parameters
     const controlWorkDirs = options.configFromOptions[
@@ -112,6 +127,21 @@ export default class CheckLeakCommand extends BaseCommand {
     };
   }
 
+  private dumpVerboseInfo(
+    controlWorkDirs: string[],
+    treatmentWorkDirs: string[],
+  ): void {
+    if (config.verbose) {
+      info.lowLevel(
+        `control working directories: ${controlWorkDirs.join(', ')}`,
+      );
+      info.lowLevel(
+        `treatment working directories: ${treatmentWorkDirs.join(', ')}`,
+      );
+      info.lowLevel(`diffing working directory: ${treatmentWorkDirs[0]}`);
+    }
+  }
+
   async run(options: CLIOptions): Promise<void> {
     config.chaseWeakMapEdge = false;
     const {controlWorkDirs, treatmentWorkDirs} = this.getWorkDirs(options);
@@ -120,6 +150,7 @@ export default class CheckLeakCommand extends BaseCommand {
       workDir: treatmentWorkDirs[0],
       silentFail: true,
     });
+    this.dumpVerboseInfo(controlWorkDirs, treatmentWorkDirs);
     // diff memory leaks
     this.useDefaultMLClusteringSetting(options.cliArgs);
     await analysis.diffLeakByWorkDir({controlWorkDirs, treatmentWorkDirs});

packages/cli/src/options/experiment/SetControlSnapshotOption.ts

@@ -0,0 +1,45 @@
+/**
+ * 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
+ */
+
+import type {ParsedArgs} from 'minimist';
+import {MemLabConfig, Nullable} from '@memlab/core';
+
+import {BaseOption, info} from '@memlab/core';
+import optionConstants from '../lib/OptionConstant';
+import {
+  createTransientWorkDirFromSingleHeapSnapshot,
+  validateHeapSnapshotFileOrThrow,
+} from './utils/ExperimentOptionsUtils';
+
+export default class SetControlWorkDirOption extends BaseOption {
+  getOptionName(): string {
+    return optionConstants.optionNames.CONTROL_SNAPSHOT;
+  }
+
+  getDescription(): string {
+    return 'set the single (target) snapshot of control run';
+  }
+
+  async parse(
+    config: MemLabConfig,
+    args: ParsedArgs,
+  ): Promise<{controlWorkDirs?: Nullable<string[]>}> {
+    let dirs = null;
+    const optionName = this.getOptionName();
+    if (optionName in args) {
+      const snapshotFile = validateHeapSnapshotFileOrThrow(args[optionName]);
+      dirs = [createTransientWorkDirFromSingleHeapSnapshot(snapshotFile)];
+    }
+    if (config.verbose && dirs && dirs[0]) {
+      info.lowLevel(`creating control working directory: ${dirs[0]}`);
+    }
+    return {controlWorkDirs: dirs};
+  }
+}

packages/cli/src/options/experiment/SetTreatmentSnapshotOption.ts

@@ -0,0 +1,45 @@
+/**
+ * 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
+ */
+
+import type {ParsedArgs} from 'minimist';
+import {MemLabConfig, Nullable} from '@memlab/core';
+
+import {BaseOption, info} from '@memlab/core';
+import optionConstants from '../lib/OptionConstant';
+import {
+  createTransientWorkDirFromSingleHeapSnapshot,
+  validateHeapSnapshotFileOrThrow,
+} from './utils/ExperimentOptionsUtils';
+
+export default class SetTreatmentWorkDirOption extends BaseOption {
+  getOptionName(): string {
+    return optionConstants.optionNames.TREATMENT_SNAPSHOT;
+  }
+
+  getDescription(): string {
+    return 'set the single (target) snapshot of treatment run';
+  }
+
+  async parse(
+    config: MemLabConfig,
+    args: ParsedArgs,
+  ): Promise<{treatmentWorkDirs?: Nullable<string[]>}> {
+    let dirs = null;
+    const optionName = this.getOptionName();
+    if (optionName in args) {
+      const snapshotFile = validateHeapSnapshotFileOrThrow(args[optionName]);
+      dirs = [createTransientWorkDirFromSingleHeapSnapshot(snapshotFile)];
+    }
+    if (config.verbose && dirs && dirs[0]) {
+      info.lowLevel(`creating treatment working directory: ${dirs[0]}`);
+    }
+    return {treatmentWorkDirs: dirs};
+  }
+}

packages/cli/src/options/experiment/utils/ExperimentOptionsUtils.ts

@@ -0,0 +1,36 @@
+/**
+ * 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
+ */
+import {AnyValue, MemLabConfig, fileManager, utils} from '@memlab/core';
+import {existsSync} from 'fs-extra';
+
+export function validateHeapSnapshotFileOrThrow(file: AnyValue): string {
+  if (typeof file !== 'string') {
+    throw utils.haltOrThrow(
+      `Heap snapshot file must be a string, but got ${typeof file}`,
+    );
+  }
+  if (!file.endsWith('.heapsnapshot')) {
+    throw utils.haltOrThrow(
+      `Heap snapshot file must end with .heapsnapshot, but got ${file}`,
+    );
+  }
+  if (!existsSync(file)) {
+    throw utils.haltOrThrow(`Heap snapshot file ${file} does not exist`);
+  }
+  return file;
+}
+
+export function createTransientWorkDirFromSingleHeapSnapshot(
+  file: string,
+): string {
+  const config = MemLabConfig.resetConfigWithTransientDir();
+  fileManager.createOrOverrideVisitOrderMetaFileForExternalSnapshot(file);
+  return config.workDir;
+}

packages/cli/src/options/lib/OptionConstant.ts

@@ -15,6 +15,7 @@ const optionNames = {
   BASELINE: 'baseline',
   CLEAN_UP_SNAPSHOT: 'clean-up-snapshot',
   CONTINUS_TEST: 'ContinuousTest',
+  CONTROL_SNAPSHOT: 'control-snapshot',
   CONTROL_WORK_DIR: 'control-work-dir',
   CHROMIUM_BINARY: 'chromium-binary',
   DEVICE: 'device',
@@ -55,6 +56,7 @@ const optionNames = {
   TRACE_ALL_OBJECTS: 'trace-all-objects',
   TRACE_CONTAINS: 'trace-contains',
   TRACE_OBJECT_SIZE_ABOVE: 'trace-object-size-above',
+  TREATMENT_SNAPSHOT: 'treatment-snapshot',
   TREATMENT_WORK_DIR: 'treatment-work-dir',
   USER_AGENT: 'user-agent',
   VERBOSE: 'verbose',

packages/core/src/lib/FileManager.ts

@@ -501,6 +501,27 @@ export class FileManager {
     );
   }
 
+  // make sure the visit order meta file exists and points to a single
+  // heap snapshot which may be outside of the working directory
+  public createOrOverrideVisitOrderMetaFileForExternalSnapshot(
+    snapshotFile: string,
+    options: FileOption = FileManager.defaultFileOption,
+  ): void {
+    // if memlab/data/cur doesn't exist, return
+    const curDataDir = this.getCurDataDir(options);
+    if (!fs.existsSync(curDataDir)) {
+      return;
+    }
+    // TODO: maybe remove the existing heap snapshot files
+
+    // If there is at least one snapshot, create a snap-seq.json file.
+    // First, get the meta file for leak detection in a single heap snapshot
+    this.createDefaultVisitOrderMetaFileWithSingleSnapshot(
+      options,
+      snapshotFile,
+    );
+  }
+
   public createDefaultVisitOrderMetaFileWithSingleSnapshot(
     options: FileOption = FileManager.defaultFileOption,
     snapshotFile: string,

packages/core/src/lib/Utils.ts

@@ -1433,7 +1433,7 @@ function checkSnapshots(
 
   if (options.snapshotDir) {
     const snapshots = getSnapshotFilesInDir(snapshotDir);
-    const min = options.minSnapshots || 1;
+    const min = options.minSnapshots || 0;
     if (snapshots.length < min) {
       utils.haltOrThrow(
         `Directory has < ${min} snapshot files: ${options.snapshotDir}`,

website/docs/cli/CLI-commands.md

@@ -118,7 +118,9 @@ memlab diff-leaks
 ```
 
 **Options**:
+ * **`--control-snapshot`**: set the single (target) snapshot of control run
  * **`--control-work-dir`**: set the working directory of the control run
+ * **`--treatment-snapshot`**: set the single (target) snapshot of treatment run
  * **`--treatment-work-dir`**: set the working directory of the treatment run
  * **`--engine`**: set the JavaScript engine (default to V8)
  * **`--leak-filter`**: specify a definition JS file for leak filter