From 73bbc6949c21c17283f8f37cddc482d8f2c958ea Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Mon, 21 Oct 2024 13:20:15 -0700
Subject: [PATCH] [dynamicIO] Avoid triggering memory leak false positive with
 makeHangingPromise (#71576)

makeHangingPromise returns a promise that rejects when an AbortSignal
aborts. If too many are created however Node warns about a possible
memory leak. We shouldn't increase the max count becuase this may
prevent detection of legitmate memory leaks in user code. As an
alternative we delegate aborts through a weakmap and attach only a
single listener. This is trivially slower and more memory but safer
semantics. This code only runs during revalidate and build so it does
not have a perf impact on production streaming use cases.
---
 .../src/server/dynamic-rendering-utils.ts     | 33 +++++++++++++++++--
 .../e2e/app-dir/dynamic-io/dynamic-io.test.ts |  6 ++++
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/packages/next/src/server/dynamic-rendering-utils.ts b/packages/next/src/server/dynamic-rendering-utils.ts
index 3a816fb01089a..b72865c19d985 100644
--- a/packages/next/src/server/dynamic-rendering-utils.ts
+++ b/packages/next/src/server/dynamic-rendering-utils.ts
@@ -10,17 +10,44 @@ export function makeHangingPromise<T>(
   expression: string
 ): Promise<T> {
   const hangingPromise = new Promise<T>((_, reject) => {
-    signal.addEventListener('abort', () => {
+    function abort() {
       reject(
         new Error(
           `During prerendering, ${expression} rejects when the prerender is complete. Typically these errors are handled by React but if you move ${expression} to a different context by using \`setTimeout\`, \`unstable_after\`, or similar functions you may observe this error and you should handle it in that context.`
         )
       )
-    })
+    }
+
+    if (signal.aborted) {
+      abort()
+      return
+    }
+
+    let listeners = listenersForSignal.get(signal)
+    if (listeners) {
+      listeners.push(abort)
+    } else {
+      listeners = [abort]
+      listenersForSignal.set(signal, listeners)
+      signal.addEventListener('abort', onAbort.bind(null, listeners), {
+        once: true,
+      })
+    }
   })
   // We are fine if no one actually awaits this promise. We shouldn't consider this an unhandled rejection so
   // we attach a noop catch handler here to suppress this warning. If you actually await somewhere or construct
   // your own promise out of it you'll need to ensure you handle the error when it rejects.
-  hangingPromise.catch(() => {})
+  hangingPromise.catch(ignoreReject)
   return hangingPromise
 }
+
+function ignoreReject() {}
+
+function onAbort(listeners: Array<AbortListener>) {
+  for (let i = 0; i < listeners.length; i++) {
+    listeners[i]()
+  }
+}
+
+type AbortListener = () => void
+const listenersForSignal = new WeakMap<AbortSignal, Array<AbortListener>>()
diff --git a/test/e2e/app-dir/dynamic-io/dynamic-io.test.ts b/test/e2e/app-dir/dynamic-io/dynamic-io.test.ts
index c4998caa6bc71..130bce6577621 100644
--- a/test/e2e/app-dir/dynamic-io/dynamic-io.test.ts
+++ b/test/e2e/app-dir/dynamic-io/dynamic-io.test.ts
@@ -20,6 +20,12 @@ describe('dynamic-io', () => {
     expect(next.cliOutput).not.toMatch('Error occurred prerendering page')
   })
 
+  if (!isNextDev) {
+    it('should not warn about potential memory leak for even listeners on AbortSignal', async () => {
+      expect(next.cliOutput).not.toMatch('MaxListenersExceededWarning')
+    })
+  }
+
   it('should prerender fully static pages', async () => {
     let $ = await next.render$('/cases/static', {})
     if (isNextDev) {