From e8217d56fbfd39fb28351dbc48c2c89f5d2d4c3a Mon Sep 17 00:00:00 2001
From: Matthew Little <zone117x@gmail.com>
Date: Thu, 31 Oct 2024 15:18:12 +0100
Subject: [PATCH] fix: memleak in timeout w/ abort signal

---
 src/helpers/__tests__/helpers.test.ts | 52 +++++++++++++++++++++++++++
 src/helpers/time.ts                   | 22 ++++++------
 2 files changed, 64 insertions(+), 10 deletions(-)
 create mode 100644 src/helpers/__tests__/helpers.test.ts

diff --git a/src/helpers/__tests__/helpers.test.ts b/src/helpers/__tests__/helpers.test.ts
new file mode 100644
index 0000000..3ae3784
--- /dev/null
+++ b/src/helpers/__tests__/helpers.test.ts
@@ -0,0 +1,52 @@
+import * as events from 'node:events';
+import { timeout } from '../time';
+
+describe('Helper tests', () => {
+  test('timeout function should not cause memory leak by accumulating abort listeners on abort', async () => {
+    const controller = new AbortController();
+    const { signal } = controller;
+
+    const countListeners = () => events.getEventListeners(signal, 'abort').length;
+
+    // Ensure the initial listener count is zero
+    expect(countListeners()).toBe(0);
+
+    // Run enough iterations to detect a pattern
+    for (let i = 0; i < 100; i++) {
+      try {
+        const sleepPromise = timeout(1000, signal);
+        controller.abort(); // Abort immediately
+        await sleepPromise;
+      } catch (err: any) {
+        expect(err.toString()).toMatch(/aborted/i);
+      }
+
+      // Assert that listener count does not increase
+      expect(countListeners()).toBeLessThanOrEqual(1); // 1 listener may temporarily be added and removed
+    }
+
+    // Final check to confirm listeners are cleaned up
+    expect(countListeners()).toBe(0);
+  });
+
+  test('timeout function should not cause memory leak by accumulating abort listeners on successful completion', async () => {
+    const controller = new AbortController();
+    const { signal } = controller;
+
+    const countListeners = () => events.getEventListeners(signal, 'abort').length;
+
+    // Ensure the initial listener count is zero
+    expect(countListeners()).toBe(0);
+
+    // Run enough iterations to detect a pattern
+    for (let i = 0; i < 100; i++) {
+      await timeout(2, signal); // Complete sleep without abort
+
+      // Assert that listener count does not increase
+      expect(countListeners()).toBe(0); // No listeners should remain after successful sleep completion
+    }
+
+    // Final check to confirm listeners are cleaned up
+    expect(countListeners()).toBe(0);
+  });
+});
diff --git a/src/helpers/time.ts b/src/helpers/time.ts
index 3f4042d..37b7faa 100644
--- a/src/helpers/time.ts
+++ b/src/helpers/time.ts
@@ -1,22 +1,24 @@
+import { addAbortListener } from 'node:events';
+
 /**
  * Wait a set amount of milliseconds or until the timer is aborted.
  * @param ms - Number of milliseconds to wait
- * @param abortController - Abort controller
+ * @param abort - Abort controller
  * @returns Promise
  */
-export function timeout(ms: number, abortController?: AbortController): Promise<void> {
+export function timeout(ms: number, abort?: AbortController | AbortSignal): Promise<void> {
   return new Promise((resolve, reject) => {
+    const signal = abort && 'signal' in abort ? abort.signal : abort;
+    if (signal?.aborted) return reject(signal.reason);
+    const disposable = signal ? addAbortListener(signal, onAbort) : undefined;
     const timeout = setTimeout(() => {
+      disposable?.[Symbol.dispose ?? (Symbol.for('nodejs.dispose') as typeof Symbol.dispose)]();
       resolve();
     }, ms);
-    abortController?.signal.addEventListener(
-      'abort',
-      () => {
-        clearTimeout(timeout);
-        reject(new Error(`Timeout aborted`));
-      },
-      { once: true }
-    );
+    function onAbort() {
+      clearTimeout(timeout);
+      reject(signal?.reason);
+    }
   });
 }