From 498cbb8441baf64f93abbb1a2c2f1254d2b4da23 Mon Sep 17 00:00:00 2001
From: Antoine du Hamel <duhamelantoine1995@gmail.com>
Date: Thu, 11 May 2023 19:55:12 +0200
Subject: [PATCH] esm: do not use `'beforeExit'` on the main thread

---
 lib/internal/modules/esm/hooks.js        | 21 ++++-----------------
 test/es-module/test-esm-loader-hooks.mjs | 16 ++++++++++++++++
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/lib/internal/modules/esm/hooks.js b/lib/internal/modules/esm/hooks.js
index df6016914f8a7e..b37b66d1d17ef2 100644
--- a/lib/internal/modules/esm/hooks.js
+++ b/lib/internal/modules/esm/hooks.js
@@ -523,14 +523,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();
 
@@ -546,11 +538,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;
@@ -559,16 +549,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 59632fec26a47b..7e60932c6f6145 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);
+  });
 });