From e1d4107ce20161a1182cfdb7d75333bdae87ca59 Mon Sep 17 00:00:00 2001
From: Izaak Schroeder <izaak.schroeder@gmail.com>
Date: Mon, 10 Jul 2023 16:37:14 -0700
Subject: [PATCH] esm: add `initialize` hook

Refs: https://github.com/nodejs/loaders/issues/147
---
 doc/api/esm.md                                |  72 +++++++++-
 doc/api/module.md                             |  21 +++
 lib/internal/modules/esm/hooks.js             |  48 +++++--
 lib/internal/modules/esm/loader.js            |  27 ++--
 lib/internal/modules/esm/utils.js             |  10 ++
 lib/internal/process/pre_execution.js         |  13 +-
 test/es-module/test-esm-loader-hooks.mjs      | 136 ++++++++++++++++++
 .../hooks-initialize-port.mjs                 |  17 +++
 .../es-module-loaders/hooks-initialize.mjs    |   3 +
 9 files changed, 323 insertions(+), 24 deletions(-)
 create mode 100644 test/fixtures/es-module-loaders/hooks-initialize-port.mjs
 create mode 100644 test/fixtures/es-module-loaders/hooks-initialize.mjs

diff --git a/doc/api/esm.md b/doc/api/esm.md
index 68ecad93b8ad3b8..9e3718241c686c4 100644
--- a/doc/api/esm.md
+++ b/doc/api/esm.md
@@ -685,6 +685,12 @@ of Node.js applications.
 <!-- YAML
 added: v8.8.0
 changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/48842
+    description: Added `initialize` hook and deprecated `getGlobalPreload`.
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/48559
+    description: Allow loaders to use `import()` and unflag `Module.register`.
   - version:
     - v18.6.0
     - v16.17.0
@@ -938,10 +944,73 @@ export async function load(url, context, nextLoad) {
 In a more advanced scenario, this can also be used to transform an unsupported
 source to a supported one (see [Examples](#examples) below).
 
+#### `initialize()`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+> The loaders API is being redesigned. This hook may disappear or its
+> signature may change. Do not rely on the API described below.
+
+> In a previous version of this API, this was split across 3 separate, now
+> deprecated, hooks (`getFormat`, `getSource`, and `transformSource`).
+
+* `data` {any} The data provided via `Module.register(loader, parentUrl, data)`
+* Returns: {any} The data returned to the caller of `Module.register`
+
+The `initialize` hook provides a way to define a custom method of that runs
+in the loader's thread when the loader is initialized. This hook can send
+and receive data from a `Module.register` invocation, including ports and
+other transferrable objects.
+
+Loader code:
+
+```js
+// In the below example this file is referenced as
+// '/path-to-my-loader.js'
+
+export function initialize({ number, port }) {
+  port.postMessage(`increment: ${number + 1}`);
+  return 'ok';
+}
+```
+
+Caller code:
+
+```js
+import assert from 'node:assert';
+import { register } from 'node:module';
+import { MessageChannel } from 'node:worker_threads';
+
+// In this example '/path-to-my-loader.js' is replaced with the
+// path to the file containing the loader contents above.
+
+// This example showcases how a message channel can be used to
+// communicate to the loader, by sending `port2` to the loader.
+const { port1, port2 } = new MessageChannel();
+
+port1.on('message', (msg) => {
+  assert(msg === 'increment: 2');
+});
+
+const result = register('/path-to-my-loader.js', import.meta.url, {
+  data: { number: 1, port: port2 },
+  transferList: [port2],
+});
+
+assert(result === 'ok');
+```
+
 #### `globalPreload()`
 
 <!-- YAML
+deprecated: REPLACEME
 changes:
+  - version:
+    - REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/48842
+    description: Deprecated in favor of `initialize` hook.
   - version:
     - v18.6.0
     - v16.17.0
@@ -949,8 +1018,7 @@ changes:
     description: Add support for chaining globalPreload hooks.
 -->
 
-> The loaders API is being redesigned. This hook may disappear or its
-> signature may change. Do not rely on the API described below.
+> Deprecated: Use `initialize` instead.
 
 > In a previous version of this API, this hook was named
 > `getGlobalPreloadCode`.
diff --git a/doc/api/module.md b/doc/api/module.md
index 8dd5fd4fa59f6c1..3193f1da6c1ab1e 100644
--- a/doc/api/module.md
+++ b/doc/api/module.md
@@ -175,6 +175,27 @@ globalPreload: http-to-https
 globalPreload: unpkg
 ```
 
+This function can also be used to pass data to the loader's `initialize`
+hook include transferrable objects like ports.
+
+```mjs
+import { register } from 'node:module';
+import { MessageChannel } from 'node:worker_threads';
+
+// This example showcases how a message channel can be used to
+// communicate to the loader, by sending `port2` to the loader.
+const { port1, port2 } = new MessageChannel();
+
+port1.on('message', (msg) => {
+  console.log(msg);
+});
+
+register('./my-programmatic-loader.mjs', import.meta.url, {
+  data: { number: 1, port: port2 },
+  transferList: [port2],
+});
+```
+
 ### `module.syncBuiltinESMExports()`
 
 <!-- YAML
diff --git a/lib/internal/modules/esm/hooks.js b/lib/internal/modules/esm/hooks.js
index 85811ad6c0f11c5..f51f9537701c279 100644
--- a/lib/internal/modules/esm/hooks.js
+++ b/lib/internal/modules/esm/hooks.js
@@ -2,6 +2,7 @@
 
 const {
   ArrayPrototypePush,
+  ArrayPrototypePushApply,
   FunctionPrototypeCall,
   Int32Array,
   ObjectAssign,
@@ -127,15 +128,17 @@ class Hooks {
    * Import and register custom/user-defined module loader hook(s).
    * @param {string} urlOrSpecifier
    * @param {string} parentURL
+   * @param {any} [data] Arbitrary data to be passed from the custom
+   * loader (user-land) to the worker.
    */
-  async register(urlOrSpecifier, parentURL) {
+  async register(urlOrSpecifier, parentURL, data) {
     const moduleLoader = require('internal/process/esm_loader').esmLoader;
     const keyedExports = await moduleLoader.import(
       urlOrSpecifier,
       parentURL,
       kEmptyObject,
     );
-    this.addCustomLoader(urlOrSpecifier, keyedExports);
+    return this.addCustomLoader(urlOrSpecifier, keyedExports, data);
   }
 
   /**
@@ -143,12 +146,15 @@ class Hooks {
    * After all hooks have been collected, the global preload hook(s) must be initialized.
    * @param {string} url Custom loader specifier
    * @param {Record<string, unknown>} exports
+   * @param {any} [data] Arbitrary data to be passed from the custom loader (user-land)
+   * to the worker.
    */
-  addCustomLoader(url, exports) {
+  addCustomLoader(url, exports, data) {
     const {
       globalPreload,
       resolve,
       load,
+      initialize,
     } = pluckHooks(exports);
 
     if (globalPreload) {
@@ -162,6 +168,7 @@ class Hooks {
       const next = this.#chains.load[this.#chains.load.length - 1];
       ArrayPrototypePush(this.#chains.load, { __proto__: null, fn: load, url, next });
     }
+    return initialize?.(data);
   }
 
   /**
@@ -553,15 +560,26 @@ class HooksProxy {
     }
   }
 
-  async makeAsyncRequest(method, ...args) {
+  /**
+   * Invoke a remote method asynchronously.
+   * @param {any[]} transferList Objects in `args` to be transferred
+   * @param {string} method Method to invoke
+   * @param  {any[]} args Arguments to pass to `method`
+   * @returns {Promise<any>}
+   */
+  async makeAsyncRequest(transferList, method, ...args) {
     this.waitForWorker();
 
     MessageChannel ??= require('internal/worker/io').MessageChannel;
     const asyncCommChannel = new MessageChannel();
 
     // Pass work to the worker.
-    debug('post async message to worker', { method, args });
-    this.#worker.postMessage({ method, args, port: asyncCommChannel.port2 }, [asyncCommChannel.port2]);
+    debug('post async message to worker', { method, args, transferList });
+    const finalTransferList = [asyncCommChannel.port2];
+    if (transferList) {
+      ArrayPrototypePushApply(finalTransferList, transferList);
+    }
+    this.#worker.postMessage({ __proto__: null, method, args, port: asyncCommChannel.port2 }, finalTransferList);
 
     if (this.#numberOfPendingAsyncResponses++ === 0) {
       // On the next lines, the main thread will await a response from the worker thread that might
@@ -593,12 +611,19 @@ class HooksProxy {
     return body;
   }
 
-  makeSyncRequest(method, ...args) {
+  /**
+   * Invoke a remote method synchronously.
+   * @param {any[]} transferList Objects in `args` to be transferred
+   * @param {string} method Method to invoke
+   * @param  {any[]} args Arguments to pass to `method`
+   * @returns {any}
+   */
+  makeSyncRequest(transferList, method, ...args) {
     this.waitForWorker();
 
     // Pass work to the worker.
-    debug('post sync message to worker', { method, args });
-    this.#worker.postMessage({ method, args });
+    debug('post sync message to worker', { method, args, transferList });
+    this.#worker.postMessage({ __proto__: null, method, args }, transferList);
 
     let response;
     do {
@@ -710,6 +735,7 @@ function pluckHooks({
   globalPreload,
   resolve,
   load,
+  initialize,
 }) {
   const acceptedHooks = { __proto__: null };
 
@@ -723,6 +749,10 @@ function pluckHooks({
     acceptedHooks.load = load;
   }
 
+  if (initialize) {
+    acceptedHooks.initialize = initialize;
+  }
+
   return acceptedHooks;
 }
 
diff --git a/lib/internal/modules/esm/loader.js b/lib/internal/modules/esm/loader.js
index 93ab847bf979d0f..00ebb0cda1764bd 100644
--- a/lib/internal/modules/esm/loader.js
+++ b/lib/internal/modules/esm/loader.js
@@ -307,7 +307,7 @@ class ModuleLoader {
     return module.getNamespace();
   }
 
-  register(specifier, parentUrl) {
+  register(specifier, parentUrl, data, transferList) {
     if (!this.#customizations) {
       // `CustomizedModuleLoader` is defined at the bottom of this file and
       // available well before this line is ever invoked. This is here in
@@ -315,7 +315,7 @@ class ModuleLoader {
       // eslint-disable-next-line no-use-before-define
       this.setCustomizations(new CustomizedModuleLoader());
     }
-    return this.#customizations.register(specifier, parentUrl);
+    return this.#customizations.register(specifier, parentUrl, data, transferList);
   }
 
   /**
@@ -426,10 +426,13 @@ class CustomizedModuleLoader {
    *                                   be registered.
    * @param {string} parentURL The parent URL from where the loader will be
    *                           registered if using it package name as specifier
+   * @param {any} data Arbitrary data to be passed from the custom loader
+   * (user-land) to the worker.
+   * @param {any[]} transferList Objects in `data` that are changing ownership
    * @returns {{ format: string, url: URL['href'] }}
    */
-  register(originalSpecifier, parentURL) {
-    return hooksProxy.makeSyncRequest('register', originalSpecifier, parentURL);
+  register(originalSpecifier, parentURL, data, transferList) {
+    return hooksProxy.makeSyncRequest(transferList, 'register', originalSpecifier, parentURL, data);
   }
 
   /**
@@ -442,12 +445,12 @@ class CustomizedModuleLoader {
    * @returns {{ format: string, url: URL['href'] }}
    */
   resolve(originalSpecifier, parentURL, importAssertions) {
-    return hooksProxy.makeAsyncRequest('resolve', originalSpecifier, parentURL, importAssertions);
+    return hooksProxy.makeAsyncRequest(undefined, 'resolve', originalSpecifier, parentURL, importAssertions);
   }
 
   resolveSync(originalSpecifier, parentURL, importAssertions) {
     // This happens only as a result of `import.meta.resolve` calls, which must be sync per spec.
-    return hooksProxy.makeSyncRequest('resolve', originalSpecifier, parentURL, importAssertions);
+    return hooksProxy.makeSyncRequest(undefined, 'resolve', originalSpecifier, parentURL, importAssertions);
   }
 
   /**
@@ -457,7 +460,7 @@ class CustomizedModuleLoader {
    * @returns {Promise<{ format: ModuleFormat, source: ModuleSource }>}
    */
   load(url, context) {
-    return hooksProxy.makeAsyncRequest('load', url, context);
+    return hooksProxy.makeAsyncRequest(undefined, 'load', url, context);
   }
 
   importMetaInitialize(meta, context, loader) {
@@ -514,18 +517,22 @@ function getHooksProxy() {
  * Register a single loader programmatically.
  * @param {string} specifier
  * @param {string} [parentURL]
- * @returns {void}
+ * @param {any} [data] Arbitrary data passed to loader's `initialize` hook
+ * @param {any[]} transferList Objects in `data` that are changing ownership
+ * @returns {any}
  * @example
  * ```js
  * register('./myLoader.js');
  * register('ts-node/esm', import.meta.url);
  * register('./myLoader.js', import.meta.url);
  * register(new URL('./myLoader.js', import.meta.url));
+ * register('./myLoader.js', import.meta.url, {banana: 'tasty'});
+ * register('./myLoader.js', import.meta.url, someArrayBuffer, [someArrayBuffer]);
  * ```
  */
-function register(specifier, parentURL = 'data:') {
+function register(specifier, parentURL = 'data:', data, transferList) {
   const moduleLoader = require('internal/process/esm_loader').esmLoader;
-  moduleLoader.register(`${specifier}`, parentURL);
+  return moduleLoader.register(`${specifier}`, parentURL, data, transferList);
 }
 
 module.exports = {
diff --git a/lib/internal/modules/esm/utils.js b/lib/internal/modules/esm/utils.js
index 3da53dd25574e88..82cf4a14fa52002 100644
--- a/lib/internal/modules/esm/utils.js
+++ b/lib/internal/modules/esm/utils.js
@@ -123,6 +123,16 @@ async function initializeHooks() {
   const hooks = new Hooks();
   esmLoader.setCustomizations(hooks);
 
+  // We need the loader customizations to be set _before_ we start invoking
+  // `--require`, otherwise loops can happen because a `--require` script
+  // might call `register(...)` before we've installed ourselves. These
+  // global values are magically set in `setupUserModules` just for us and
+  // we call them in the correct order.
+  // N.B.  This block appears here specifically in order to ensure that
+  // `--require` calls occur before `--loader` ones do.
+  globalThis.loadPreloadModules();
+  globalThis.initializeFrozenIntrinsics();
+
   const parentURL = pathToFileURL(cwd).href;
   for (let i = 0; i < customLoaderURLs.length; i++) {
     await hooks.register(
diff --git a/lib/internal/process/pre_execution.js b/lib/internal/process/pre_execution.js
index f785fa8f80a3c15..1a8c27e42184ef1 100644
--- a/lib/internal/process/pre_execution.js
+++ b/lib/internal/process/pre_execution.js
@@ -154,9 +154,16 @@ function setupUserModules(isLoaderWorker = false) {
   initializeESMLoader(isLoaderWorker);
   const CJSLoader = require('internal/modules/cjs/loader');
   assert(!CJSLoader.hasLoadedAnyUserCJSModule);
-  loadPreloadModules();
-  // Need to be done after --require setup.
-  initializeFrozenIntrinsics();
+  if (isLoaderWorker) {
+    // Loader workers are responsible for doing this themselves.
+    globalThis.loadPreloadModules = loadPreloadModules;
+    globalThis.initializeFrozenIntrinsics = initializeFrozenIntrinsics;
+  } else {
+    loadPreloadModules();
+    // Need to be done after --require setup.
+    initializeFrozenIntrinsics();
+  }
+
 }
 
 function refreshRuntimeOptions() {
diff --git a/test/es-module/test-esm-loader-hooks.mjs b/test/es-module/test-esm-loader-hooks.mjs
index a5e43eb51c5345a..52a339fff2522ca 100644
--- a/test/es-module/test-esm-loader-hooks.mjs
+++ b/test/es-module/test-esm-loader-hooks.mjs
@@ -553,4 +553,140 @@ describe('Loader hooks', { concurrency: true }, () => {
     assert.strictEqual(code, 0);
     assert.strictEqual(signal, null);
   });
+
+  it('should invoke `initialize` correctly', async () => {
+    const { code, signal, stdout, stderr } = await spawnPromisified(execPath, [
+      '--no-warnings',
+      '--experimental-loader',
+      fixtures.fileURL('/es-module-loaders/hooks-initialize.mjs'),
+      '--input-type=module',
+      '--eval',
+      `
+        import os from 'node:os';
+        import fs from 'node:fs';
+      `,
+    ]);
+
+    const lines = stdout.trim().split('\n');
+
+    assert.strictEqual(lines.length, 1);
+    assert.strictEqual(lines[0], 'hooks initialize');
+
+    assert.strictEqual(stderr, '');
+
+    assert.strictEqual(code, 0);
+    assert.strictEqual(signal, null);
+  });
+
+  it('should allow communicating with loader via `register` ports', async () => {
+    const { code, signal, stdout, stderr } = await spawnPromisified(execPath, [
+      '--no-warnings',
+      '--input-type=module',
+      '--eval',
+      `
+        import {MessageChannel} from 'node:worker_threads';
+        import {register} from 'node:module';
+        const {port1, port2} = new MessageChannel();
+        port1.on('message', (msg) => {
+          console.log('message', msg);
+        });
+        const result = register(
+          ${JSON.stringify(fixtures.fileURL('/es-module-loaders/hooks-initialize-port.mjs'))},
+          'data:',
+          port2,
+          [port2],
+        );
+        console.log('register', result);
+
+        await import('node:os');
+        port1.close();
+      `,
+    ]);
+
+    const lines = stdout.split('\n');
+
+    assert.strictEqual(lines[0], 'register ok');
+    assert.strictEqual(lines[1], 'message initialize');
+    assert.strictEqual(lines[2], 'message resolve node:os');
+
+    assert.strictEqual(stderr, '');
+
+    assert.strictEqual(code, 0);
+    assert.strictEqual(signal, null);
+  });
+
+  it('should have `register` work with cjs', async () => {
+    const { code, signal, stdout, stderr } = await spawnPromisified(execPath, [
+      '--no-warnings',
+      '--input-type=commonjs',
+      '--eval',
+      `
+        const {register} = require('node:module');
+        register(
+          ${JSON.stringify(fixtures.fileURL('/es-module-loaders/hooks-initialize.mjs'))},
+        );
+        register(
+          ${JSON.stringify(fixtures.fileURL('/es-module-loaders/loader-load-foo-or-42.mjs'))},
+        );
+
+        import('node:os').then((result) => {
+          console.log(result.default);
+        });
+      `,
+    ]);
+
+    const lines = stdout.split('\n');
+
+    assert.strictEqual(lines[0], 'hooks initialize');
+    assert.strictEqual(lines[1], 'foo');
+
+    assert.strictEqual(stderr, '');
+
+    assert.strictEqual(code, 0);
+    assert.strictEqual(signal, null);
+  });
+
+  it('`register` should work with `require`', async () => {
+    const { code, signal, stdout, stderr } = await spawnPromisified(execPath, [
+      '--no-warnings',
+      '--require',
+      fixtures.path('/es-module-loaders/register-loader.cjs'),
+      '--input-type=module',
+      '--eval',
+      `
+        import 'node:os';
+      `,
+    ]);
+
+    const lines = stdout.split('\n');
+
+    assert.strictEqual(lines[0], 'resolve passthru');
+
+    assert.strictEqual(stderr, '');
+
+    assert.strictEqual(code, 0);
+    assert.strictEqual(signal, null);
+  });
+
+  it('`register` should work with `import`', async () => {
+    const { code, signal, stdout, stderr } = await spawnPromisified(execPath, [
+      '--no-warnings',
+      '--import',
+      fixtures.fileURL('/es-module-loaders/register-loader.mjs'),
+      '--input-type=module',
+      '--eval',
+      `
+        import 'node:os';
+      `,
+    ]);
+
+    const lines = stdout.split('\n');
+
+    assert.strictEqual(lines[0], 'resolve passthru');
+
+    assert.strictEqual(stderr, '');
+
+    assert.strictEqual(code, 0);
+    assert.strictEqual(signal, null);
+  });
 });
diff --git a/test/fixtures/es-module-loaders/hooks-initialize-port.mjs b/test/fixtures/es-module-loaders/hooks-initialize-port.mjs
new file mode 100644
index 000000000000000..c522e3fa8bfd988
--- /dev/null
+++ b/test/fixtures/es-module-loaders/hooks-initialize-port.mjs
@@ -0,0 +1,17 @@
+let thePort = null;
+
+export async function initialize(port) {
+  port.postMessage('initialize');
+  thePort = port;
+  return 'ok';
+}
+
+export async function resolve(specifier, context, next) {
+  if (specifier === 'node:fs' || specifier.includes('loader')) {
+    return next(specifier);
+  }
+
+  thePort.postMessage(`resolve ${specifier}`);
+
+  return next(specifier);
+}
diff --git a/test/fixtures/es-module-loaders/hooks-initialize.mjs b/test/fixtures/es-module-loaders/hooks-initialize.mjs
new file mode 100644
index 000000000000000..dbea050dd3a92f5
--- /dev/null
+++ b/test/fixtures/es-module-loaders/hooks-initialize.mjs
@@ -0,0 +1,3 @@
+export async function initialize() {
+  console.log('hooks initialize');
+}