Ensure all messages to loader are acknowledged

Author: timfish

README.md

@@ -74,26 +74,27 @@ module.register('import-in-the-middle/hook.mjs', import.meta.url, {
 ### Only Intercepting Hooked modules
 
 If you are `Hook`'ing all modules before they are imported, for example in a
-module loaded via the Node.js `--import` argument, you can configure the loader
-hook to only intercept those specific modules: 
+module loaded via the Node.js `--import` CLI argument, you can configure the
+loader to intercept only modules that were specifically hooked.
 
 `instrument.mjs`
 ```js
 import { register } from 'module'
 import { Hook, createAddHookMessageChannel } from 'import-in-the-middle'
 
-const addHookMessagePort = createAddHookMessageChannel()
+const { addHookMessagePort, waitForAllMessagesAcknowledged } = createAddHookMessageChannel()
 
-const options = { 
-  data: { addHookMessagePort }, 
-  transferList: [addHookMessagePort] 
-}
+const options = { data: { addHookMessagePort }, transferList: [addHookMessagePort] }
 
 register('import-in-the-middle/hook.mjs', import.meta.url, options)
 
 Hook(['fs'], (exported, name, baseDir) => {
   // Instrument the fs module
 })
+
+// Ensure that the loader has acknowledged all the modules 
+// before we allow execution to continue
+await waitForAllMessagesAcknowledged()
 ```
 `my-app.mjs`
 ```js

hook.js

@@ -263,15 +263,16 @@ function createHook (meta) {
           }
 
           includeModules.push(...modules)
-        })
+          data.addHookMessagePort.postMessage('ack')
+        }).unref()
       }
     }
   }
 
   async function resolve (specifier, context, parentResolve) {
     cachedResolve = parentResolve
 
-    // See github.com/nodejs/import-in-the-middle/pull/76.
+    // See https://github.com/nodejs/import-in-the-middle/pull/76.
     if (specifier === iitmURL) {
       return {
         url: specifier,

index.js

@@ -32,16 +32,39 @@ function callHookFn (hookFn, namespace, name, baseDir) {
   }
 }
 
-let sendToMessageChannel
+let sendModulesToLoader
 
 function createAddHookMessageChannel () {
   const { port1, port2 } = new MessageChannel()
+  let pendingAckCount = 0
+  let resolveFn
 
-  sendToMessageChannel = (modules) => {
+  sendModulesToLoader = (modules) => {
+    pendingAckCount++
     port1.postMessage(modules)
   }
 
-  return port2
+  port1.on('message', () => {
+    pendingAckCount--
+
+    if (resolveFn && pendingAckCount <= 0) {
+      resolveFn()
+    }
+  }).unref()
+
+  function waitForAllMessagesAcknowledged () {
+    const promise = new Promise((resolve) => {
+      resolveFn = resolve
+    })
+
+    if (pendingAckCount === 0) {
+      resolveFn()
+    }
+
+    return promise
+  }
+
+  return { addHookMessagePort: port2, waitForAllMessagesAcknowledged }
 }
 
 function Hook (modules, options, hookFn) {
@@ -56,8 +79,8 @@ function Hook (modules, options, hookFn) {
   }
   const internals = options ? options.internals === true : false
 
-  if (sendToMessageChannel && Array.isArray(modules)) {
-    sendToMessageChannel(modules)
+  if (sendModulesToLoader && Array.isArray(modules)) {
+    sendModulesToLoader(modules)
   }
 
   this._iitmHook = (name, namespace) => {

test/fixtures/import.mjs

@@ -1,9 +1,10 @@
 import { register } from 'module'
 import { Hook, createAddHookMessageChannel } from '../../index.js'
-// We've imported path here to ensure that the hook is still applied later.
+// We've imported path here to ensure that the hook is still applied later even
+// if the library is used here.
 import * as path from 'path'
 
-const addHookMessagePort = createAddHookMessageChannel()
+const { addHookMessagePort, waitForAllMessagesAcknowledged } = createAddHookMessageChannel()
 
 register('../../hook.mjs', import.meta.url, { data: { addHookMessagePort }, transferList: [addHookMessagePort] })
 
@@ -12,3 +13,5 @@ Hook(['path'], (exports) => {
 })
 
 console.assert(path.sep !== '@')
+
+await waitForAllMessagesAcknowledged()