chore: improve benchmark CLI

packages/driver-bench/.gitignore

@@ -1,2 +1,3 @@
 node_modules/
 lib/
+results.json

packages/driver-bench/src/cli.mts

@@ -1,24 +1,54 @@
 #!/usr/bin/env node
 
-import child_process from 'node:child_process';
-import events from 'node:events';
-import module from 'node:module';
-import path from 'node:path';
+import os from 'node:os';
 import process from 'node:process';
+import util from 'node:util';
 
-let require;
+import { main } from './driverBench/index.mjs';
 
-let { MONGODB_DRIVER_PATH = '' } = process.env;
-const { DEBUG_BENCH = 'no' } = process.env;
+const argv = util.parseArgs({
+  args: process.argv.slice(2),
+  options: {
+    grep: { type: 'string', short: 'g', default: '' },
+    help: { type: 'boolean', short: 'h', default: false }
+  },
+  strict: true,
+  allowPositionals: true
+});
+
+const [driverPath] = argv.positionals;
 
-if (MONGODB_DRIVER_PATH === '') {
-  require ??= module.createRequire(import.meta.dirname);
-  MONGODB_DRIVER_PATH = require.resolve('mongodb');
+if (argv.values.help || driverPath == null || driverPath.length === 0) {
+  console.error('node cli.mjs DRIVER_PATH [--grep/-g BENCHMARK_NAME]');
+  process.exit(0);
 }
 
-const benchmark = child_process.fork(path.join(import.meta.dirname, './driverBench/index.mjs'), {
-  execArgv: DEBUG_BENCH === 'yes' ? ['--enable-source-maps'] : [],
-  stdio: 'inherit',
-  env: { MONGODB_DRIVER_PATH }
-});
-await events.once(benchmark, 'exit');
+const hw = os.cpus();
+const ram = os.totalmem() / 1024 ** 3;
+const platform = { name: hw[0].model, cores: hw.length, ram: `${ram}GB` };
+
+const systemInfo = () =>
+  [
+    `\n- cpu: ${platform.name}`,
+    `- cores: ${platform.cores}`,
+    `- arch: ${os.arch()}`,
+    `- os: ${process.platform} (${os.release()})`,
+    `- ram: ${platform.ram}\n`
+  ].join('\n');
+
+const mongodbDriver = await import(driverPath);
+
+let { MONGODB_URI = '' } = process.env;
+MONGODB_URI = MONGODB_URI.length === 0 ? 'mongodb://127.0.0.1:27017' : MONGODB_URI;
+
+const client = new mongodbDriver.MongoClient(MONGODB_URI, { serverSelectionTimeoutMS: 2000 });
+try {
+  console.log('server version:', (await client.db('admin').command({ buildInfo: 1 })).version);
+  await client.close();
+  console.log(`Benching driver at: ${driverPath}`);
+  console.log(systemInfo());
+  await main(argv, mongodbDriver);
+} catch (error) {
+  await client.close();
+  console.log('Unable to benchmark against:\n', MONGODB_URI, '\n', error);
+}

packages/driver-bench/src/driverBench/common.mts

@@ -4,10 +4,8 @@ import process from 'node:process';
 import { Readable } from 'node:stream';
 import { pipeline } from 'node:stream/promises';
 
-const { MONGODB_DRIVER_PATH = '' } = process.env;
-const { MongoClient, GridFSBucket } = await import(
-  MONGODB_DRIVER_PATH.length !== 0 ? MONGODB_DRIVER_PATH : 'mongodb'
-);
+let { MONGODB_URI = '' } = process.env;
+MONGODB_URI = MONGODB_URI.length === 0 ? 'mongodb://127.0.0.1:27017' : MONGODB_URI;
 
 const DB_NAME = 'perftest';
 const COLLECTION_NAME = 'corpus';
@@ -26,10 +24,10 @@ export function loadSpecString(filePath) {
   return loadSpecFile(filePath, 'utf8');
 }
 
-export function makeClient() {
-  this.client = new MongoClient(process.env.MONGODB_URI || 'mongodb://127.0.0.1:27017', {
-    serverSelectionTimeoutMS: 2000
-  });
+export function makeMakeClient({ MongoClient }) {
+  return function () {
+    this.client = new MongoClient(MONGODB_URI, { serverSelectionTimeoutMS: 2000 });
+  };
 }
 
 export function connectClient() {
@@ -64,8 +62,10 @@ export function dropCollection() {
   });
 }
 
-export function initBucket() {
-  this.bucket = new GridFSBucket(this.db);
+export function makeInitBucket({ GridFSBucket }) {
+  return function () {
+    this.bucket = new GridFSBucket(this.db);
+  };
 }
 
 export function dropBucket() {

packages/driver-bench/src/driverBench/index.mts

@@ -1,6 +1,5 @@
 import { writeFile } from 'node:fs/promises';
-import os from 'node:os';
-import { inspect } from 'node:util';
+import util from 'node:util';
 
 import { Runner } from '../mongoBench/index.mjs';
 import {
@@ -9,93 +8,74 @@ import {
   makeSingleBench
 } from '../mongoBench/suites/index.mjs';
 
-const bsonType = 'js-bson';
-// TODO(NODE-4606): test against different driver configurations in CI
-const hw = os.cpus();
-const ram = os.totalmem() / 1024 ** 3;
-const platform = { name: hw[0].model, cores: hw.length, ram: `${ram}GB` };
-
-const systemInfo = () =>
-  [
-    `\n- cpu: ${platform.name}`,
-    `- cores: ${platform.cores}`,
-    `- arch: ${os.arch()}`,
-    `- os: ${process.platform} (${os.release()})`,
-    `- ram: ${platform.ram}\n`
-  ].join('\n');
-console.log(systemInfo());
-
-function average(arr) {
+function average(arr: number[]) {
   return arr.reduce((x, y) => x + y, 0) / arr.length;
 }
 
-const benchmarkRunner = new Runner()
-  .suite('singleBench', suite => makeSingleBench(suite))
-  .suite('multiBench', suite => makeMultiBench(suite))
-  .suite('parallel', suite => makeParallelBenchmarks(suite));
+// TODO(NODE-4606): test against different driver configurations in CI
+export async function main(
+  argv: { values: { grep: string } },
+  mongodbDriver: { MongoClient; GridFSBucket }
+) {
+  const benchmarkRunner = new Runner(argv.values)
+    .suite('singleBench', suite => makeSingleBench(suite, mongodbDriver))
+    .suite('multiBench', suite => makeMultiBench(suite, mongodbDriver))
+    .suite('parallel', suite => makeParallelBenchmarks(suite, mongodbDriver));
 
-benchmarkRunner
-  .run()
-  .then(microBench => {
-    const singleBench = average([
-      microBench.singleBench.findOne,
-      microBench.singleBench.smallDocInsertOne,
-      microBench.singleBench.largeDocInsertOne
-    ]);
-    const multiBench = average(Object.values(microBench.multiBench));
+  const microBench = await benchmarkRunner.run();
 
-    const parallelBench = average([
-      microBench.parallel.ldjsonMultiFileUpload,
-      microBench.parallel.ldjsonMultiFileExport,
-      microBench.parallel.gridfsMultiFileUpload,
-      microBench.parallel.gridfsMultiFileDownload
-    ]);
+  const singleBench = average([
+    microBench.singleBench.findOne,
+    microBench.singleBench.smallDocInsertOne,
+    microBench.singleBench.largeDocInsertOne
+  ]);
+  const multiBench = average(Object.values(microBench.multiBench));
 
-    const readBench = average([
-      microBench.singleBench.findOne,
-      microBench.multiBench.findManyAndEmptyCursor,
-      microBench.multiBench.gridFsDownload,
-      microBench.parallel.gridfsMultiFileDownload,
-      microBench.parallel.ldjsonMultiFileExport
-    ]);
-    const writeBench = average([
-      microBench.singleBench.smallDocInsertOne,
-      microBench.singleBench.largeDocInsertOne,
-      microBench.multiBench.smallDocBulkInsert,
-      microBench.multiBench.largeDocBulkInsert,
-      microBench.multiBench.gridFsUpload,
-      microBench.parallel.ldjsonMultiFileUpload,
-      microBench.parallel.gridfsMultiFileUpload
-    ]);
+  const parallelBench = average([
+    microBench.parallel.ldjsonMultiFileUpload,
+    microBench.parallel.ldjsonMultiFileExport,
+    microBench.parallel.gridfsMultiFileUpload,
+    microBench.parallel.gridfsMultiFileDownload
+  ]);
 
-    const driverBench = average([readBench, writeBench]);
+  const readBench = average([
+    microBench.singleBench.findOne,
+    microBench.multiBench.findManyAndEmptyCursor,
+    microBench.multiBench.gridFsDownload,
+    microBench.parallel.gridfsMultiFileDownload,
+    microBench.parallel.ldjsonMultiFileExport
+  ]);
+  const writeBench = average([
+    microBench.singleBench.smallDocInsertOne,
+    microBench.singleBench.largeDocInsertOne,
+    microBench.multiBench.smallDocBulkInsert,
+    microBench.multiBench.largeDocBulkInsert,
+    microBench.multiBench.gridFsUpload,
+    microBench.parallel.ldjsonMultiFileUpload,
+    microBench.parallel.gridfsMultiFileUpload
+  ]);
 
-    const benchmarkResults = {
-      singleBench,
-      multiBench,
-      parallelBench,
-      readBench,
-      writeBench,
-      driverBench,
-      ...microBench.parallel,
-      ...microBench.bsonBench,
-      ...microBench.singleBench,
-      ...microBench.multiBench
-    };
+  const driverBench = average([readBench, writeBench]);
 
-    return Object.entries(benchmarkResults).map(([benchmarkName, result]) => {
-      return {
-        info: {
-          test_name: benchmarkName,
-          tags: [bsonType]
-        },
-        metrics: [{ name: 'megabytes_per_second', value: result }]
-      };
-    });
-  })
-  .then(data => {
-    const results = JSON.stringify(data, undefined, 2);
-    console.log(inspect(data, { depth: Infinity, colors: true }));
-    return writeFile('results.json', results);
-  })
-  .catch(err => console.error(err));
+  const benchmarkResults = {
+    singleBench,
+    multiBench,
+    parallelBench,
+    readBench,
+    writeBench,
+    driverBench,
+    ...microBench.parallel,
+    ...microBench.bsonBench,
+    ...microBench.singleBench,
+    ...microBench.multiBench
+  };
+
+  const data = Object.entries(benchmarkResults).map(([benchmarkName, result]) => ({
+    info: { test_name: benchmarkName, tags: ['js-bson'] },
+    metrics: [{ name: 'megabytes_per_second', value: result }]
+  }));
+
+  const results = JSON.stringify(data, undefined, 2);
+  console.log(util.inspect(data, { depth: Infinity, colors: true }));
+  await writeFile('results.json', results);
+}

packages/driver-bench/src/mod.mts

@@ -1,4 +1,4 @@
-// TODO: This module is not really meant to be imported the CLI is currently the focus of how to utilize this code
+// TODO: This module is not really meant to be imported. The CLI is currently the focus of how to utilize this code
 // See wishlist in readme
 export { Benchmark } from './mongoBench/benchmark.mjs';
 export { Runner } from './mongoBench/runner.mjs';

packages/driver-bench/src/mongoBench/runner.mts

@@ -53,6 +53,8 @@ export type RunnerOptions = {
   maxExecutionTime?: number;
   minExecutionCount?: number;
   reporter?: (message?: any, ...optionalParams: any[]) => void;
+  /** Only run TEST NAMES that match */
+  grep?: string;
 };
 
 export class Runner {
@@ -61,13 +63,14 @@ export class Runner {
   private minExecutionCount: number;
   private reporter: (message?: any, ...optionalParams: any[]) => void;
   private children: Record<string, Suite>;
+  private grep: RegExp | null;
 
-  constructor(options?: RunnerOptions) {
-    options = options || {};
+  constructor(options: RunnerOptions = {}) {
     this.minExecutionTime = options.minExecutionTime || CONSTANTS.DEFAULT_MIN_EXECUTION_TIME;
     this.maxExecutionTime = options.maxExecutionTime || CONSTANTS.DEFAULT_MAX_EXECUTION_TIME;
     this.minExecutionCount = options.minExecutionCount || CONSTANTS.DEFAULT_MIN_EXECUTION_COUNT;
     this.reporter = options.reporter ?? console.log.bind(console);
+    this.grep = options.grep ? new RegExp(options.grep, 'i') : null;
     this.children = {};
   }
 
@@ -121,6 +124,11 @@ export class Runner {
     const result: Record<string, number> = {};
 
     for (const [name, benchmark] of benchmarks) {
+      if (this.grep && !this.grep.test(name)) {
+        this.reporter(`    Skipping Benchmark "${name}"`);
+        result[name] = 0;
+        continue;
+      }
       this.reporter(`    Executing Benchmark "${name}"`);
       result[name] = await this._runBenchmark(benchmark);
       this.reporter(`    Result ${result[name].toFixed(4)} MB/s`);