feat: add code caching support MONGOSH-1334 (#39)

This speeds up mongosh startup time by 26.8 % when active
(but will likely be mutually exclusive with snapshot support,
once that is finished).

Author: addaleax

.github/workflows/nodejs.yml

@@ -36,7 +36,7 @@ jobs:
       fail-fast: false
       matrix:
         node-version: [14.x, 16.x, 18.x, 19.x]
-        shard: [1, 2, 3]
+        shard: [1, 2, 3, 4]
     runs-on: windows-latest
     steps:
       - uses: actions/checkout@v2
@@ -61,4 +61,4 @@ jobs:
       - name: Sharded tests
         run: npm test -- -g "shard ${{ matrix.shard }}"
       - name: Unsharded tests
-        run: npm test -- -g "shard [1-3]" -i
+        run: npm test -- -g "shard [1-4]" -i

bin/boxednode.js

@@ -13,7 +13,7 @@ const argv = require('yargs')
     alias: 't', type: 'string', demandOption: true, desc: 'Target executable file'
   })
   .option('node-version', {
-    alias: 'n', type: 'string', desc: 'Node.js version or semver version range or tarball file url', default: '*'
+    alias: 'n', type: 'string', desc: 'Node.js version or semver version range or .tar.gz file url', default: '*'
   })
   .option('configure-args', {
     alias: 'C', type: 'string', desc: 'Extra ./configure or vcbuild arguments, comma-separated'
@@ -27,9 +27,12 @@ const argv = require('yargs')
   .option('namespace', {
     alias: 'N', type: 'string', desc: 'Module identifier for the generated binary'
   })
-  .options('use-legacy-default-uv-loop', {
+  .option('use-legacy-default-uv-loop', {
     type: 'boolean', desc: 'Use the global singleton libuv event loop rather than a separate local one'
   })
+  .option('use-code-cache', {
+    alias: 'H', type: 'boolean', desc: 'Use Node.js code cache support to speed up startup'
+  })
   .example('$0 -s myProject.js -t myProject.exe -n ^14.0.0',
     'Create myProject.exe from myProject.js using Node.js v14')
   .help()
@@ -46,7 +49,8 @@ const argv = require('yargs')
       configureArgs: (argv.C || '').split(',').filter(Boolean),
       makeArgs: (argv.M || '').split(',').filter(Boolean),
       namespace: argv.N,
-      useLegacyDefaultUvLoop: argv.useLegacyDefaultUvLoop
+      useLegacyDefaultUvLoop: argv.useLegacyDefaultUvLoop,
+      useCodeCache: argv.H
     });
   } catch (err) {
     console.error(err);

resources/entry-point-trampoline.js

@@ -2,6 +2,7 @@
 const Module = require('module');
 const vm = require('vm');
 const path = require('path');
+const assert = require('assert');
 const {
   requireMappings,
   enableBindingsPatch
@@ -49,7 +50,7 @@ if (enableBindingsPatch) {
   });
 }
 
-module.exports = (src) => {
+module.exports = (src, codeCacheMode, codeCache) => {
   const __filename = process.execPath;
   const __dirname = path.dirname(process.execPath);
   const innerRequire = Module.createRequire(__filename);
@@ -68,6 +69,7 @@ module.exports = (src) => {
   Object.setPrototypeOf(require, Object.getPrototypeOf(innerRequire));
 
   process.argv.unshift(__filename);
+  process.boxednode = {};
 
   const module = {
     exports,
@@ -77,10 +79,23 @@ module.exports = (src) => {
     path: __dirname,
     require
   };
-  vm.compileFunction(src, [
+  const mainFunction = vm.compileFunction(src, [
     '__filename', '__dirname', 'require', 'exports', 'module'
   ], {
-    filename: __filename
-  })(__filename, __dirname, require, exports, module);
+    filename: __filename,
+    cachedData: codeCache.length > 0 ? codeCache : undefined,
+    produceCachedData: codeCacheMode === 'generate'
+  });
+  if (codeCacheMode === 'generate') {
+    assert.strictEqual(mainFunction.cachedDataProduced, true);
+    process.stdout.write(mainFunction.cachedData);
+    return;
+  }
+
+  process.boxednode.hasCodeCache = codeCache.length > 0;
+  // https://github.com/nodejs/node/pull/46320
+  process.boxednode.rejectedCodeCache = mainFunction.cachedDataRejected;
+
+  mainFunction(__filename, __dirname, require, exports, module);
   return module.exports;
 };

resources/main-template.cc

@@ -6,6 +6,7 @@
 #include "node.h"
 #include "node_api.h"
 #include "uv.h"
+#include "brotli/decode.h"
 #if HAVE_OPENSSL
 #include <openssl/err.h>
 #include <openssl/ssl.h>
@@ -43,6 +44,7 @@ void TearDownOncePerProcess();
 #endif
 namespace boxednode {
 Local<String> GetBoxednodeMainScriptSource(Isolate* isolate);
+Local<Uint8Array> GetBoxednodeCodeCacheBuffer(Isolate* isolate);
 }
 
 extern "C" {
@@ -151,9 +153,18 @@ static int RunNodeInstance(MultiIsolatePlatform* platform,
       return 1; // There has been a JS exception.
     }
     assert(loadenv_ret->IsFunction());
-    Local<Value> source = boxednode::GetBoxednodeMainScriptSource(isolate);
-    if (loadenv_ret.As<Function>()->Call(context, Null(isolate), 1, &source).IsEmpty())
+    Local<Value> trampoline_args[] = {
+      boxednode::GetBoxednodeMainScriptSource(isolate),
+      String::NewFromUtf8Literal(isolate, BOXEDNODE_CODE_CACHE_MODE),
+      boxednode::GetBoxednodeCodeCacheBuffer(isolate),
+    };
+    if (loadenv_ret.As<Function>()->Call(
+        context,
+        Null(isolate),
+        sizeof(trampoline_args) / sizeof(trampoline_args[0]),
+        trampoline_args).IsEmpty()) {
       return 1; // JS exception.
+    }
 
     {
       // SealHandleScope protects against handle leaks from callbacks.

src/helpers.ts

@@ -5,6 +5,7 @@ import childProcess from 'child_process';
 import { promisify } from 'util';
 import tar from 'tar';
 import stream from 'stream';
+import zlib from 'zlib';
 import { once } from 'events';
 
 export const pipeline = promisify(stream.pipeline);
@@ -101,3 +102,51 @@ export function createCppJsStringDefinition (fnName: string, source: string): st
   }
   `;
 }
+
+export async function createCompressedBlobDefinition (fnName: string, source: Uint8Array): Promise<string> {
+  const compressed = await promisify(zlib.brotliCompress)(source, {
+    params: {
+      [zlib.constants.BROTLI_PARAM_QUALITY]: zlib.constants.BROTLI_MAX_QUALITY,
+      [zlib.constants.BROTLI_PARAM_SIZE_HINT]: source.length
+    }
+  });
+  return `
+  static const uint8_t ${fnName}_source_[] = {
+    ${Uint8Array.prototype.toString.call(compressed)}
+  };
+
+  std::string ${fnName}() {
+    ${source.length === 0 ? 'return {};' : `
+    size_t decoded_size = ${source.length};
+    std::string dst(decoded_size, 0);
+    const auto result = BrotliDecoderDecompress(
+      ${compressed.length},
+      ${fnName}_source_,
+      &decoded_size,
+      reinterpret_cast<uint8_t*>(&dst[0]));
+    assert(result == BROTLI_DECODER_RESULT_SUCCESS);
+    assert(decoded_size == ${source.length});
+    return dst;`}
+  }
+
+  std::shared_ptr<v8::BackingStore> ${fnName}BackingStore() {
+    std::string* str = new std::string(std::move(${fnName}()));
+    return v8::SharedArrayBuffer::NewBackingStore(
+      &str->front(),
+      str->size(),
+      [](void*, size_t, void* deleter_data) {
+        delete static_cast<std::string*>(deleter_data);
+      },
+      static_cast<void*>(str));
+  }
+
+  v8::Local<v8::Uint8Array> ${fnName}Buffer(v8::Isolate* isolate) {
+    ${source.length === 0 ? `
+    auto array_buffer = v8::SharedArrayBuffer::New(isolate, 0);
+    ` : `
+    auto array_buffer = v8::SharedArrayBuffer::New(isolate, ${fnName}BackingStore());
+    `}
+    return v8::Uint8Array::New(array_buffer, 0, array_buffer->ByteLength());
+  }
+  `;
+}

src/index.ts

@@ -11,10 +11,11 @@ import { promisify } from 'util';
 import { promises as fs, createReadStream, createWriteStream } from 'fs';
 import { AddonConfig, loadGYPConfig, storeGYPConfig, modifyAddonGyp } from './native-addons';
 import { ExecutableMetadata, generateRCFile } from './executable-metadata';
-import { spawnBuildCommand, ProcessEnv, pipeline, createCppJsStringDefinition } from './helpers';
+import { spawnBuildCommand, ProcessEnv, pipeline, createCppJsStringDefinition, createCompressedBlobDefinition } from './helpers';
 import { Readable } from 'stream';
 import nv from '@pkgjs/nv';
 import { fileURLToPath, URL } from 'url';
+import { execFile } from 'child_process';
 
 // Download and unpack a tarball containing the code for a specific Node.js version.
 async function getNodeSourceForVersion (range: string, dir: string, logger: Logger, retries = 2): Promise<string> {
@@ -26,7 +27,8 @@ async function getNodeSourceForVersion (range: string, dir: string, logger: Logg
   } catch { /* not a valid URL */ }
 
   if (inputIsFileUrl) {
-    logger.stepStarting(`Extracting tarball from ${range}`);
+    logger.stepStarting(`Extracting tarball from ${range} to ${dir}`);
+    await fs.mkdir(dir, { recursive: true });
     await pipeline(
       createReadStream(fileURLToPath(range)),
       zlib.createGunzip(),
@@ -179,9 +181,9 @@ async function compileNode (
   const nodeVersion = await getNodeVersionFromSourceDirectory(sourcePath);
   if (nodeVersion[0] > 19 || (nodeVersion[0] === 19 && nodeVersion[1] >= 4)) {
     if (process.platform !== 'win32') {
-      buildArgs.unshift('--disable-shared-readonly-heap');
+      buildArgs = ['--disable-shared-readonly-heap', ...buildArgs];
     } else {
-      buildArgs.unshift('no-shared-roheap');
+      buildArgs = ['no-shared-roheap', ...buildArgs];
     }
   }
 
@@ -201,6 +203,13 @@ async function compileNode (
 
     return path.join(sourcePath, 'out', 'Release', 'node');
   } else {
+    // On Windows, running vcbuild multiple times may result in errors
+    // when the source data changes in between runs.
+    await fs.rm(path.join(sourcePath, 'out', 'Release'), {
+      recursive: true,
+      force: true
+    });
+
     // These defaults got things to work locally. We only include them if no
     // conflicting arguments have been passed manually.
     const vcbuildArgs: string[] = [...buildArgs, ...makeArgs, 'projgen'];
@@ -230,6 +239,7 @@ type CompilationOptions = {
   addons?: AddonConfig[],
   enableBindingsPatch?: boolean,
   useLegacyDefaultUvLoop?: boolean;
+  useCodeCache?: boolean,
   executableMetadata?: ExecutableMetadata,
   preCompileHook?: (nodeSourceTree: string, options: CompilationOptions) => void | Promise<void>
 }
@@ -258,12 +268,12 @@ async function compileJSFileAsBinaryImpl (options: CompilationOptions, logger: L
   const enableBindingsPatch = options.enableBindingsPatch ?? options.addons?.length > 0;
 
   const jsMainSource = await fs.readFile(options.sourceFile, 'utf8');
+  const registerFunctions: string[] = [];
 
   // We use the official embedder API for stability, which is available in all
   // supported versions of Node.js.
   {
     const extraGypDependencies: string[] = [];
-    const registerFunctions: string[] = [];
     for (const addon of (options.addons || [])) {
       const addonResult = await modifyAddonGyp(
         addon, nodeSourcePath, options.env || process.env, logger);
@@ -290,23 +300,6 @@ async function compileJSFileAsBinaryImpl (options: CompilationOptions, logger: L
       await fs.writeFile(path.join(nodeSourcePath, 'src', header), source);
     }
     logger.stepCompleted();
-
-    logger.stepStarting('Handling main file source');
-    let mainSource = await fs.readFile(
-      path.join(__dirname, '..', 'resources', 'main-template.cc'), 'utf8');
-    mainSource = mainSource.replace(/\bREPLACE_WITH_ENTRY_POINT\b/g,
-      JSON.stringify(JSON.stringify(`${namespace}/${namespace}`)));
-    mainSource = mainSource.replace(/\bREPLACE_DECLARE_LINKED_MODULES\b/g,
-      registerFunctions.map((fn) => `void ${fn}(const void**,const void**);\n`).join(''));
-    mainSource = mainSource.replace(/\bREPLACE_DEFINE_LINKED_MODULES\b/g,
-      registerFunctions.map((fn) => `${fn},`).join(''));
-    mainSource = mainSource.replace(/\bREPLACE_WITH_MAIN_SCRIPT_SOURCE_GETTER\b/g,
-      createCppJsStringDefinition('GetBoxednodeMainScriptSource', jsMainSource));
-    if (options.useLegacyDefaultUvLoop) {
-      mainSource = `#define BOXEDNODE_USE_DEFAULT_UV_LOOP 1\n${mainSource}`;
-    }
-    await fs.writeFile(path.join(nodeSourcePath, 'src', 'node_main.cc'), mainSource);
-    logger.stepCompleted();
   }
 
   logger.stepStarting('Inserting custom code into Node.js source');
@@ -338,13 +331,61 @@ async function compileJSFileAsBinaryImpl (options: CompilationOptions, logger: L
     logger.stepCompleted();
   }
 
-  const binaryPath = await compileNode(
-    nodeSourcePath,
-    extraJSSourceFiles,
-    options.configureArgs,
-    options.makeArgs,
-    options.env || process.env,
-    logger);
+  async function writeMainFileAndCompile ({ codeCacheBlob, codeCacheMode }: {
+    codeCacheBlob: Uint8Array,
+    codeCacheMode: 'ignore' | 'generate' | 'consume'
+  }): Promise<string> {
+    logger.stepStarting('Handling main file source');
+    let mainSource = await fs.readFile(
+      path.join(__dirname, '..', 'resources', 'main-template.cc'), 'utf8');
+    mainSource = mainSource.replace(/\bREPLACE_WITH_ENTRY_POINT\b/g,
+      JSON.stringify(JSON.stringify(`${namespace}/${namespace}`)));
+    mainSource = mainSource.replace(/\bREPLACE_DECLARE_LINKED_MODULES\b/g,
+      registerFunctions.map((fn) => `void ${fn}(const void**,const void**);\n`).join(''));
+    mainSource = mainSource.replace(/\bREPLACE_DEFINE_LINKED_MODULES\b/g,
+      registerFunctions.map((fn) => `${fn},`).join(''));
+    mainSource = mainSource.replace(/\bREPLACE_WITH_MAIN_SCRIPT_SOURCE_GETTER\b/g,
+      createCppJsStringDefinition('GetBoxednodeMainScriptSource', jsMainSource) + '\n' +
+      await createCompressedBlobDefinition('GetBoxednodeCodeCache', codeCacheBlob));
+    mainSource = mainSource.replace(/\bBOXEDNODE_CODE_CACHE_MODE\b/g,
+      JSON.stringify(codeCacheMode));
+    if (options.useLegacyDefaultUvLoop) {
+      mainSource = `#define BOXEDNODE_USE_DEFAULT_UV_LOOP 1\n${mainSource}`;
+    }
+    await fs.writeFile(path.join(nodeSourcePath, 'src', 'node_main.cc'), mainSource);
+    logger.stepCompleted();
+
+    return await compileNode(
+      nodeSourcePath,
+      extraJSSourceFiles,
+      options.configureArgs,
+      options.makeArgs,
+      options.env || process.env,
+      logger);
+  }
+
+  let binaryPath: string;
+  if (!options.useCodeCache) {
+    binaryPath = await writeMainFileAndCompile({
+      codeCacheBlob: new Uint8Array(0),
+      codeCacheMode: 'ignore'
+    });
+  } else {
+    binaryPath = await writeMainFileAndCompile({
+      codeCacheBlob: new Uint8Array(0),
+      codeCacheMode: 'generate'
+    });
+    logger.stepStarting('Running code cache generation');
+    const codeCacheResult = await promisify(execFile)(binaryPath, { encoding: 'buffer' });
+    if (codeCacheResult.stdout.length === 0) {
+      throw new Error('Empty code cache result');
+    }
+    logger.stepCompleted();
+    binaryPath = await writeMainFileAndCompile({
+      codeCacheBlob: codeCacheResult.stdout,
+      codeCacheMode: 'consume'
+    });
+  }
 
   logger.stepStarting(`Moving resulting binary to ${options.targetFile}`);
   await fs.mkdir(path.dirname(options.targetFile), { recursive: true });