esm: add initialize hook

Refs: nodejs/loaders#147
izaakschroeder · Jul 19, 2023 · ddfa9d4 · ddfa9d4
1 parent a2fc4a3
commit ddfa9d4
Show file tree

Hide file tree

Showing 5 changed files with 114 additions and 18 deletions.
diff --git a/lib/internal/modules/esm/hooks.js b/lib/internal/modules/esm/hooks.js
@@ -127,28 +127,31 @@ class Hooks {
    * Import and register custom/user-defined module loader hook(s).
    * @param {string} urlOrSpecifier
    * @param {string} parentURL
+   * @param {any} data
    */
-  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 await this.addCustomLoader(urlOrSpecifier, keyedExports, data);
   }
 
   /**
    * Collect custom/user-defined module loader hook(s).
    * 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 pass to loader's `initialize`
    */
-  addCustomLoader(url, exports) {
+  async addCustomLoader(url, exports, data) {
     const {
       globalPreload,
       resolve,
       load,
+      initialize,
     } = pluckHooks(exports);
 
     if (globalPreload) {
@@ -162,6 +165,9 @@ class Hooks {
       const next = this.#chains.load[this.#chains.load.length - 1];
       ArrayPrototypePush(this.#chains.load, { __proto__: null, fn: load, url, next });
     }
+    if (initialize) {
+      return await initialize(data);
+    }
   }
 
   /**
@@ -553,15 +559,18 @@ class HooksProxy {
     }
   }
 
-  async makeAsyncRequest(method, ...args) {
+  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 = transferList ?
+      [asyncCommChannel.port2, ...transferList] :
+      [asyncCommChannel.port2]
+    this.#worker.postMessage({ 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 +602,12 @@ class HooksProxy {
     return body;
   }
 
-  makeSyncRequest(method, ...args) {
+  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({ method, args }, transferList);
 
     let response;
     do {
@@ -710,6 +719,7 @@ function pluckHooks({
   globalPreload,
   resolve,
   load,
+  initialize,
 }) {
   const acceptedHooks = { __proto__: null };
 
@@ -723,6 +733,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
@@ -297,15 +297,15 @@ 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
       // order to preserve the git diff instead of moving the class.
       // 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);
   }
 
   /**
@@ -411,8 +411,8 @@ class CustomizedModuleLoader {
    *                           registered if using it package name as specifier
    * @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);
   }
 
   /**
@@ -425,12 +425,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);
   }
 
   /**
@@ -440,7 +440,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) {
@@ -497,6 +497,7 @@ function getHooksProxy() {
  * Register a single loader programmatically.
  * @param {string} specifier
  * @param {string} [parentURL]
+ * @param {any} [data]
  * @returns {void}
  * @example
  * ```js
@@ -506,9 +507,9 @@ function getHooksProxy() {
  * register(new URL('./myLoader.js', import.meta.url));
  * ```
  */
-function register(specifier, parentURL = 'data:') {
+function register(specifier, parentURL = 'data:', data = undefined, 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/test/es-module/test-esm-loader-hooks.mjs b/test/es-module/test-esm-loader-hooks.mjs
@@ -553,4 +553,65 @@ 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);
+  });
 });
diff --git a/test/fixtures/es-module-loaders/hooks-initialize-port.mjs 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
@@ -0,0 +1,3 @@
+export async function initialize() {
+  console.log('hooks initialize');
+}