perf: load binaries in the same process (#97)

Author: paul-soporan

package.json

@@ -60,7 +60,7 @@
   },
   "scripts": {
     "build": "rm -rf dist shims && webpack && ts-node ./mkshims.ts",
-    "corepack": "ts-node ./sources/main.ts",
+    "corepack": "ts-node ./sources/_entryPoint.ts",
     "prepack": "node ./.yarn/releases/*.*js build",
     "postpack": "rm -rf dist shims",
     "typecheck": "tsc --noEmit",

sources/Engine.ts

@@ -5,11 +5,11 @@ import semver                                                 from 'semver';
 
 import defaultConfig                                          from '../config.json';
 
-import * as folderUtils                                       from './folderUtils';
 import * as corepackUtils                                     from './corepackUtils';
+import * as folderUtils                                       from './folderUtils';
 import * as semverUtils                                       from './semverUtils';
-import {SupportedPackageManagers, SupportedPackageManagerSet} from './types';
 import {Config, Descriptor, Locator}                          from './types';
+import {SupportedPackageManagers, SupportedPackageManagerSet} from './types';
 
 
 export class Engine {

sources/_entryPoint.ts

@@ -0,0 +1,8 @@
+import {runMain} from './main';
+
+// Used by the generated shims
+export {runMain};
+
+// Using `eval` to be sure that Webpack doesn't transform it
+if (process.mainModule === eval(`module`))
+  runMain(process.argv.slice(2));

sources/commands/Enable.ts

@@ -5,6 +5,7 @@ import path                                                              from 'p
 import which                                                             from 'which';
 
 import {Context}                                                         from '../main';
+import * as nodeUtils                                                    from '../nodeUtils';
 import {isSupportedPackageManager, SupportedPackageManagerSetWithoutNpm} from '../types';
 
 export class EnableCommand extends Command<Context> {
@@ -51,7 +52,7 @@ export class EnableCommand extends Command<Context> {
     installDirectory = fs.realpathSync(installDirectory);
 
     // We use `eval` so that Webpack doesn't statically transform it.
-    const manifestPath = eval(`require`).resolve(`corepack/package.json`);
+    const manifestPath = nodeUtils.dynamicRequire.resolve(`corepack/package.json`);
 
     const distFolder = path.join(path.dirname(manifestPath), `dist`);
     if (!fs.existsSync(distFolder))

sources/corepackUtils.ts

@@ -1,4 +1,3 @@
-import {StdioOptions, spawn, ChildProcess}                     from 'child_process';
 import fs                                                      from 'fs';
 import path                                                    from 'path';
 import semver                                                  from 'semver';
@@ -7,11 +6,9 @@ import * as debugUtils                                         from './debugUtil
 import * as folderUtils                                        from './folderUtils';
 import * as fsUtils                                            from './fsUtils';
 import * as httpUtils                                          from './httpUtils';
-import {Context}                                               from './main';
+import * as nodeUtils                                          from './nodeUtils';
 import {RegistrySpec, Descriptor, Locator, PackageManagerSpec} from './types';
 
-declare const __non_webpack_require__: unknown;
-
 export async function fetchAvailableTags(spec: RegistrySpec): Promise<Record<string, string>> {
   switch (spec.type) {
     case `npm`: {
@@ -133,7 +130,10 @@ export async function installVersion(installTarget: string, locator: Locator, {s
   return installFolder;
 }
 
-export async function runVersion(installSpec: { location: string, spec: PackageManagerSpec }, locator: Locator, binName: string, args: Array<string>, context: Context) {
+/**
+ * Loads the binary, taking control of the current process.
+ */
+export async function runVersion(installSpec: { location: string, spec: PackageManagerSpec }, binName: string, args: Array<string>): Promise<void> {
   let binPath: string | null = null;
   if (Array.isArray(installSpec.spec.bin)) {
     if (installSpec.spec.bin.some(bin => bin === binName)) {
@@ -155,82 +155,23 @@ export async function runVersion(installSpec: { location: string, spec: PackageM
   if (!binPath)
     throw new Error(`Assertion failed: Unable to locate path for bin '${binName}'`);
 
-  return new Promise<number>((resolve, reject) => {
-    process.on(`SIGINT`, () => {
-      // We don't want to exit the process before the child, so we just
-      // ignore SIGINT and wait for the regular exit to happen (the child
-      // will receive SIGINT too since it's part of the same process grp)
-    });
-
-    const stdio: StdioOptions = [`pipe`, `pipe`, `pipe`];
-
-    if (context.stdin === process.stdin)
-      stdio[0] = `inherit`;
-    if (context.stdout === process.stdout)
-      stdio[1] = `inherit`;
-    if (context.stderr === process.stderr)
-      stdio[2] = `inherit`;
-
-    const v8CompileCache = typeof __non_webpack_require__ !== `undefined`
-      ? eval(`require`).resolve(`./vcc.js`)
-      : eval(`require`).resolve(`corepack/dist/vcc.js`);
-
-    const child = spawn(process.execPath, [`--require`, v8CompileCache, binPath!, ...args], {
-      cwd: context.cwd,
-      stdio,
-      env: {
-        ...process.env,
-        COREPACK_ROOT: path.dirname(eval(`__dirname`)),
-      },
-    });
-
-    activeChildren.add(child);
-
-    if (activeChildren.size === 1) {
-      process.on(`SIGINT`, sigintHandler);
-      process.on(`SIGTERM`, sigtermHandler);
-    }
-
-    if (context.stdin !== process.stdin)
-      context.stdin.pipe(child.stdin!);
-    if (context.stdout !== process.stdout)
-      child.stdout!.pipe(context.stdout);
-    if (context.stderr !== process.stderr)
-      child.stderr!.pipe(context.stderr);
+  nodeUtils.registerV8CompileCache();
 
-    child.on(`error`, error => {
-      activeChildren.delete(child);
+  // We load the binary into the current process,
+  // while making it think it was spawned.
 
-      if (activeChildren.size === 0) {
-        process.off(`SIGINT`, sigintHandler);
-        process.off(`SIGTERM`, sigtermHandler);
-      }
+  // Non-exhaustive list of requirements:
+  // - Yarn uses process.argv[1] to determine its own path: https://github.com/yarnpkg/berry/blob/0da258120fc266b06f42aed67e4227e81a2a900f/packages/yarnpkg-cli/sources/main.ts#L80
+  // - pnpm uses `require.main == null` to determine its own version: https://github.com/pnpm/pnpm/blob/e2866dee92991e979b2b0e960ddf5a74f6845d90/packages/cli-meta/src/index.ts#L14
 
-      reject(error);
-    });
+  process.env.COREPACK_ROOT = path.dirname(eval(`__dirname`));
 
-    child.on(`exit`, exitCode => {
-      activeChildren.delete(child);
+  process.argv = [
+    process.execPath,
+    binPath,
+    ...args,
+  ];
+  process.execArgv = [];
 
-      if (activeChildren.size === 0) {
-        process.off(`SIGINT`, sigintHandler);
-        process.off(`SIGTERM`, sigtermHandler);
-      }
-
-      resolve(exitCode !== null ? exitCode : 1);
-    });
-  });
-}
-
-const activeChildren = new Set<ChildProcess>();
-
-function sigintHandler() {
-  // We don't want SIGINT to kill our process; we want it to kill the
-  // innermost process, whose end will cause our own to exit.
-}
-
-function sigtermHandler() {
-  for (const child of activeChildren) {
-    child.kill();
-  }
+  return nodeUtils.loadMainModule(binPath);
 }

sources/main.ts

@@ -5,8 +5,8 @@ import {DisableCommand}                                          from './command
 import {EnableCommand}                                           from './commands/Enable';
 import {HydrateCommand}                                          from './commands/Hydrate';
 import {PrepareCommand}                                          from './commands/Prepare';
-import * as miscUtils                                            from './miscUtils';
 import * as corepackUtils                                        from './corepackUtils';
+import * as miscUtils                                            from './miscUtils';
 import * as specUtils                                            from './specUtils';
 import {Locator, SupportedPackageManagers, Descriptor}           from './types';
 
@@ -19,7 +19,7 @@ type PackageManagerRequest = {
   binaryVersion: string | null;
 };
 
-function getPackageManagerRequestFromCli(parameter: string | undefined, context: CustomContext & Partial<Context>): PackageManagerRequest {
+function getPackageManagerRequestFromCli(parameter: string | undefined, context: CustomContext & Partial<Context>): PackageManagerRequest | null {
   if (!parameter)
     return null;
 
@@ -82,14 +82,20 @@ async function executePackageManagerRequest({packageManager, binaryName, binaryV
     throw new UsageError(`Failed to successfully resolve '${descriptor.range}' to a valid ${descriptor.name} release`);
 
   const installSpec = await context.engine.ensurePackageManager(resolved);
-  const exitCode = await corepackUtils.runVersion(installSpec, resolved, binaryName, args, context);
 
-  return exitCode;
+  return await corepackUtils.runVersion(installSpec, binaryName, args);
 }
 
-export async function main(argv: Array<string>, context: CustomContext & Partial<Context>) {
+async function main(argv: Array<string>) {
   const corepackVersion = require(`../package.json`).version;
 
+  // Because we load the binaries in the same process, we don't support custom contexts.
+  const context = {
+    ...Cli.defaultContext,
+    cwd: process.cwd(),
+    engine: new Engine(),
+  };
+
   const [firstArg, ...restArgs] = argv;
   const request = getPackageManagerRequestFromCli(firstArg, context);
 
@@ -110,10 +116,7 @@ export async function main(argv: Array<string>, context: CustomContext & Partial
     cli.register(HydrateCommand);
     cli.register(PrepareCommand);
 
-    return await cli.run(argv, {
-      ...Cli.defaultContext,
-      ...context,
-    });
+    return await cli.run(argv, context);
   } else {
     // Otherwise, we create a single-command CLI to run the specified package manager (we still use Clipanion in order to pretty-print usage errors).
     const cli = new Cli({
@@ -129,25 +132,16 @@ export async function main(argv: Array<string>, context: CustomContext & Partial
       }
     });
 
-    return await cli.run(restArgs, {
-      ...Cli.defaultContext,
-      ...context,
-    });
+    return await cli.run(restArgs, context);
   }
 }
 
+// Important: this is the only function that the corepack binary exports.
 export function runMain(argv: Array<string>) {
-  main(argv, {
-    cwd: process.cwd(),
-    engine: new Engine(),
-  }).then(exitCode => {
+  main(argv).then(exitCode => {
     process.exitCode = exitCode;
   }, err => {
     console.error(err.stack);
     process.exitCode = 1;
   });
 }
-
-// Using `eval` to be sure that Webpack doesn't transform it
-if (process.mainModule === eval(`module`))
-  runMain(process.argv.slice(2));

sources/module.d.ts

@@ -0,0 +1,16 @@
+import 'module';
+
+declare module 'module' {
+  const _cache: {[p: string]: NodeModule};
+
+  function _nodeModulePaths(from: string): Array<string>;
+  function _resolveFilename(request: string, parent: NodeModule | null | undefined, isMain: boolean): string;
+}
+
+declare global {
+  namespace NodeJS {
+    interface Module {
+      load(path: string): void;
+    }
+  }
+}

sources/nodeUtils.ts

@@ -0,0 +1,43 @@
+import Module from 'module';
+import path   from 'path';
+
+declare const __non_webpack_require__: NodeRequire | undefined;
+
+export const dynamicRequire: NodeRequire = typeof __non_webpack_require__ !== `undefined`
+  ? __non_webpack_require__
+  : require;
+
+function getV8CompileCachePath() {
+  return typeof __non_webpack_require__ !== `undefined`
+    ? `./vcc.js`
+    : `corepack/dist/vcc.js`;
+}
+
+export function registerV8CompileCache() {
+  const vccPath = getV8CompileCachePath();
+  dynamicRequire(vccPath);
+}
+
+/**
+ * Loads a module as a main module, enabling the `require.main === module` pattern.
+ */
+export function loadMainModule(id: string): void {
+  const modulePath = Module._resolveFilename(id, null, true);
+
+  const module = new Module(modulePath, undefined);
+
+  module.filename = modulePath;
+  module.paths = Module._nodeModulePaths(path.dirname(modulePath));
+
+  Module._cache[modulePath] = module;
+
+  process.mainModule = module;
+  module.id = `.`;
+
+  try {
+    return module.load(modulePath);
+  } catch (error) {
+    delete Module._cache[modulePath];
+    throw error;
+  }
+}

tests/_runCli.ts

@@ -1,36 +1,35 @@
 import {PortablePath, npath} from '@yarnpkg/fslib';
-import {PassThrough}         from 'stream';
-
-import {Engine}              from '../sources/Engine';
-import {main}                from '../sources/main';
+import {spawn}               from 'child_process';
 
 export async function runCli(cwd: PortablePath, argv: Array<string>) {
-  const stdin = new PassThrough();
-  const stdout = new PassThrough();
-  const stderr = new PassThrough();
-
   const out: Array<Buffer> = [];
   const err: Array<Buffer> = [];
 
-  stdout.on(`data`, chunk => {
-    out.push(chunk);
-  });
+  return new Promise((resolve, reject) => {
+    const child = spawn(process.execPath, [require.resolve(`corepack/dist/corepack.js`), ...argv], {
+      cwd: npath.fromPortablePath(cwd),
+      env: process.env,
+      stdio: `pipe`,
+    });
 
-  stderr.on(`data`, chunk => {
-    err.push(chunk);
-  });
+    child.stdout.on(`data`, chunk => {
+      out.push(chunk);
+    });
 
-  const exitCode = await main(argv, {
-    cwd: npath.fromPortablePath(cwd),
-    engine: new Engine(),
-    stdin,
-    stdout,
-    stderr,
-  });
+    child.stderr.on(`data`, chunk => {
+      err.push(chunk);
+    });
+
+    child.on(`error`, error => {
+      reject(error);
+    });
 
-  return {
-    exitCode,
-    stdout: Buffer.concat(out).toString(),
-    stderr: Buffer.concat(err).toString(),
-  };
+    child.on(`exit`, exitCode => {
+      resolve({
+        exitCode,
+        stdout: Buffer.concat(out).toString(),
+        stderr: Buffer.concat(err).toString(),
+      });
+    });
+  });
 }

tsconfig.json

@@ -11,6 +11,10 @@
     "module": "commonjs",
     "resolveJsonModule": true,
     "skipLibCheck": true,
+    "strict": true,
     "target": "es2017"
+  },
+  "ts-node": {
+    "transpileOnly": true
   }
 }