chore: re-enable startup snapshot support MONGOSH-1605 (#1842)

This involves adding nodejs/node@e54ddf898f6c as a (cleanly
applying) patch but otherwise seems to work just fine.

.gitignore

@@ -25,3 +25,4 @@ compiled-ts.tgz
 mongocryptd.pid
 .sbom
 .nvm
+snapshot.blob

package.json

@@ -128,14 +128,14 @@
   "workspaces": [
     "configs/eslint-config-mongosh",
     "configs/tsconfig-mongosh",
+    "scripts/docker",
     "packages/async-rewriter2",
     "packages/build",
     "packages/errors",
     "packages/history",
     "packages/java-shell",
     "packages/js-multiline-to-singleline",
     "packages/types",
-    "scripts/docker",
     "packages/i18n",
     "packages/logging",
     "packages/service-provider-core",

packages/build/package.json

@@ -65,7 +65,7 @@
     "@mongodb-js/dl-center": "^1.1.1",
     "@octokit/rest": "^17.9.0",
     "aws-sdk": "^2.674.0",
-    "boxednode": "^2.3.0",
+    "boxednode": "^2.4.0",
     "command-exists": "^1.2.9",
     "download": "^8.0.0",
     "es-aggregate-error": "^1.0.9",

packages/build/src/compile/signable-compiler.ts

@@ -180,11 +180,11 @@ export class SignableCompiler {
       preCompileHook,
       executableMetadata: this.executableMetadata,
       // Node.js startup snapshots are an experimental feature of Node.js.
-      // TODO(MONGOSH-1605): Re-enable startup snapshots after figuring out
-      // issues with running the binary when CPU features differ
-      // significantly.
-      useCodeCache: true,
-      // useNodeSnapshot: true,
+      // useCodeCache: true,
+      useNodeSnapshot: true,
+      // To account for the fact that we are manually patching Node.js to include
+      // https://github.com/nodejs/node/pull/50453 until we have caught up with upstream
+      nodeSnapshotConfigFlags: ['WithoutCodeCache'],
     });
   }
 }

packages/cli-repl/package.json

@@ -33,6 +33,7 @@
     "prepublish": "npm run compile",
     "webpack-build": "npm run compile && webpack --mode production",
     "webpack-build-dev": "npm run compile && webpack --mode development",
+    "start-snapshot": "rm -f snapshot.blob && node --snapshot-blob snapshot.blob --build-snapshot dist/mongosh.js && node --snapshot-blob snapshot.blob dist/mongosh.js",
     "prettier": "prettier",
     "reformat": "npm run prettier -- --write . && npm run eslint --fix"
   },

packages/cli-repl/src/run.ts

@@ -44,6 +44,9 @@ if ((v8 as any)?.startupSnapshot?.isBuildingSnapshot?.()) {
     const ConsoleCtor = console.Console;
     (v8 as any).startupSnapshot.addDeserializeCallback(() => {
       console.Console = ConsoleCtor;
+      // Work around Node.js caching the cwd when snapshotting
+      // https://github.com/nodejs/node/pull/51901
+      process.chdir('.');
     });
   }
 

packages/cli-repl/src/startup-timing.ts

@@ -3,6 +3,7 @@ import {
   type TimingCategory,
   type TimingInterface,
 } from '@mongosh/types';
+import v8 from 'v8';
 
 const jsTimingEntries: [string, string, bigint][] = [];
 
@@ -58,24 +59,36 @@ const timing: TimingInterface = linkTimingInterface();
 export const markTime = timing.markTime;
 export const getTimingData = timing.getTimingData;
 
-if (process.env.MONGOSH_SHOW_TIMING_DATA) {
-  process.on('exit', function () {
-    const rawTimingData = getTimingData();
-    if (process.env.MONGOSH_SHOW_TIMING_DATA === 'json') {
-      console.log(JSON.stringify(rawTimingData));
-    } else {
-      console.table(
-        rawTimingData.map(([category, label, time], i) => [
-          category,
-          label,
-          `${(time / 1_000_000).toFixed(2)}ms`,
-          i > 0
-            ? `+${((time - rawTimingData[i - 1][2]) / 1_000_000).toFixed(2)}ms`
-            : '',
-        ])
-      );
-    }
+function installExitHandler() {
+  if (process.env.MONGOSH_SHOW_TIMING_DATA) {
+    process.on('exit', function () {
+      const rawTimingData = getTimingData();
+      if (process.env.MONGOSH_SHOW_TIMING_DATA === 'json') {
+        console.log(JSON.stringify(rawTimingData));
+      } else {
+        console.table(
+          rawTimingData.map(([category, label, time], i) => [
+            category,
+            label,
+            `${(time / 1_000_000).toFixed(2)}ms`,
+            i > 0
+              ? `+${((time - rawTimingData[i - 1][2]) / 1_000_000).toFixed(
+                  2
+                )}ms`
+              : '',
+          ])
+        );
+      }
+    });
+  }
+}
+
+if ((v8 as any)?.startupSnapshot?.isBuildingSnapshot?.()) {
+  (v8 as any).startupSnapshot.addDeserializeCallback(() => {
+    installExitHandler();
   });
+} else {
+  installExitHandler();
 }
 
 markTime(TimingCategories.REPLInstantiation, 'cli-repl timing initialized');

packages/cli-repl/webpack.config.js

@@ -1,4 +1,6 @@
 'use strict';
+const fs = require('fs');
+const crypto = require('crypto');
 const { merge } = require('webpack-merge');
 const path = require('path');
 const { WebpackDependenciesPlugin } = require('@mongodb-js/sbom-tools');
@@ -30,6 +32,67 @@ const config = {
   },
   plugins: [webpackDependenciesPlugin],
   entry: './lib/run.js',
+  resolve: {
+    alias: {
+      // This is similar to https://github.com/babel/babel/issues/12442,
+      // @babel/code-frame loads chalk loads supports-color which checks
+      // for TTY color support during startup rather than at runtime
+      '@babel/code-frame': makeLazyForwardModule('@babel/code-frame'),
+    },
+  },
 };
 
 module.exports = merge(baseWebpackConfig, config);
+
+// Helper to create a module that lazily loads the actual target package
+// when it is being encountered. This is useful for snapshotting, where some
+// packages either cannot be snapshotted or perform initialization during
+// startup that should depend on runtime state.
+function makeLazyForwardModule(pkg) {
+  const S = JSON.stringify;
+  const tmpdir = path.resolve(__dirname, '..', 'tmp', 'lazy-webpack-modules');
+  fs.mkdirSync(tmpdir, { recursive: true });
+  const filename = path.join(
+    tmpdir,
+    crypto.createHash('sha256').update(pkg).digest('hex').slice(0, 16) + '.js'
+  );
+
+  const moduleContents = require(pkg);
+  let source = `'use strict';\nlet _cache;\n`;
+  source += `function orig() {\n_cache = require(${S(
+    require.resolve(pkg)
+  )}); orig = () => _cache; return _cache;\n}\n`;
+  if (typeof moduleContents === 'function') {
+    source += `module.exports = function(...args) { return orig().apply(this, args); };\n`;
+  } else {
+    source += `module.exports = {};\n`;
+  }
+  let i = 0;
+  for (const key of Object.keys(moduleContents)) {
+    if (typeof moduleContents[key] === 'function') {
+      source += `module.exports[${S(
+        key
+      )}] = function(...args) { return orig()[${S(
+        key
+      )}].apply(this, args); };\n`;
+    } else {
+      source += `let value_${i}, value_${i}_set = false;\n`;
+      source += `Object.defineProperty(module.exports, ${S(key)}, {
+        enumerable: true, configurable: true,
+        get() {
+          if (value_${i}_set) return value_${i};
+          value_${i} = orig()[${S(key)}];
+          value_${i}_set = true;
+          return value_${i};
+        }, set(v) {
+          value_${i}_set = true;
+          value_${i} = v;
+        }
+      })\n`;
+      i++;
+    }
+  }
+
+  fs.writeFileSync(filename, source);
+  return filename;
+}