esm: do not use `'beforeExit'` on the main thread

PR-URL: https://github.com/nodejs/node/pull/47964
Fixes: https://github.com/nodejs/node/issues/47929
Reviewed-By: Geoffrey Booth <webadmin@geoffreybooth.com>
Reviewed-By: Yagiz Nizipli <yagiz@nizipli.com>
Reviewed-By: Moshe Atlow <moshe@atlow.co.il>

2 files changed, 20 insertions, 17 deletions

diff --git a/lib/internal/modules/esm/hooks.js b/lib/internal/modules/esm/hooks.js
index 667c4c0576..d927e3c9ba 100644
--- a/lib/internal/modules/esm/hooks.js
+++ b/lib/internal/modules/esm/hooks.js
@@ -521,14 +521,6 @@ class HooksProxy {
     }
   }
 
-  #beforeExitHandler = () => {
-    debug('beforeExit main thread', this.#lock, this.#numberOfPendingAsyncResponses);
-    if (this.#numberOfPendingAsyncResponses !== 0) {
-      // The worker still has some work to do, let's wait for it before terminating the process.
-      this.#worker.ref();
-    }
-  };
-
   async makeAsyncRequest(method, ...args) {
     this.#waitForWorker();
 
@@ -544,11 +536,9 @@ class HooksProxy {
       // come AFTER the last task in the event loop has run its course and there would be nothing
       // left keeping the thread alive (and once the main thread dies, the whole process stops).
       // However we want to keep the process alive until the worker thread responds (or until the
-      // event loop of the worker thread is also empty). So we add the beforeExit handler whose
-      // mission is to lock the main thread until we hear back from the worker thread. The `if`
-      // condition is there so we only add the event handler once (if there are already pending
-      // async responses, the previous calls have added the event listener).
-      process.on('beforeExit', this.#beforeExitHandler);
+      // event loop of the worker thread is also empty), so we ref the worker until we get all the
+      // responses back.
+      this.#worker.ref();
     }
 
     let response;
@@ -557,16 +547,13 @@ class HooksProxy {
       await AtomicsWaitAsync(this.#lock, WORKER_TO_MAIN_THREAD_NOTIFICATION, this.#workerNotificationLastId).value;
       this.#workerNotificationLastId = AtomicsLoad(this.#lock, WORKER_TO_MAIN_THREAD_NOTIFICATION);
 
-      // In case the beforeExit handler was called during the await, we revert its actions.
-      this.#worker.unref();
-
       response = receiveMessageOnPort(asyncCommChannel.port1);
     } while (response == null);
     debug('got async response from worker', { method, args }, this.#lock);
 
     if (--this.#numberOfPendingAsyncResponses === 0) {
       // We got all the responses from the worker, its job is done (until next time).
-      process.off('beforeExit', this.#beforeExitHandler);
+      this.#worker.unref();
     }
 
     const body = this.#unwrapMessage(response);
diff --git a/test/es-module/test-esm-loader-hooks.mjs b/test/es-module/test-esm-loader-hooks.mjs
index 59632fec26..7e60932c6f 100644
--- a/test/es-module/test-esm-loader-hooks.mjs
+++ b/test/es-module/test-esm-loader-hooks.mjs
@@ -243,4 +243,20 @@ describe('Loader hooks', { concurrency: true }, () => {
     assert.strictEqual(code, 42);
     assert.strictEqual(signal, null);
   });
+
+  it('should be fine to call `process.removeAllListeners("beforeExit")` from the main thread', async () => {
+    const { code, signal, stdout, stderr } = await spawnPromisified(execPath, [
+      '--no-warnings',
+      '--experimental-loader',
+      'data:text/javascript,export function load(a,b,c){return new Promise(d=>setTimeout(()=>d(c(a,b)),99))}',
+      '--input-type=module',
+      '--eval',
+      'setInterval(() => process.removeAllListeners("beforeExit"),1).unref();await import("data:text/javascript,")',
+    ]);
+
+    assert.strictEqual(stderr, '');
+    assert.strictEqual(stdout, '');
+    assert.strictEqual(code, 0);
+    assert.strictEqual(signal, null);
+  });
 });