lib: remove settled dependant signals when they are GCed

geeksilva97 · aduh95 · commit a9473bb8e375 · 2024-10-19T10:54:36.000+02:00
PR-URL: #55354 Reviewed-By: James M Snell <jasnell@gmail.com> Reviewed-By: Chemi Atlow <chemi@atlow.co.il>
diff --git a/lib/internal/abort_controller.js b/lib/internal/abort_controller.js
@@ -84,6 +84,17 @@ function lazyMessageChannel() {
 }
 
 const clearTimeoutRegistry = new SafeFinalizationRegistry(clearTimeout);
+const dependantSignalsCleanupRegistry = new SafeFinalizationRegistry((signalWeakRef) => {
+  const signal = signalWeakRef.deref();
+  if (signal === undefined) {
+    return;
+  }
+  signal[kDependantSignals].forEach((ref) => {
+    if (ref.deref() === undefined) {
+      signal[kDependantSignals].delete(ref);
+    }
+  });
+});
 const gcPersistentSignals = new SafeSet();
 
 const kAborted = Symbol('kAborted');
@@ -243,24 +254,27 @@ class AbortSignal extends EventTarget {
       }
       signal[kDependantSignals] ??= new SafeSet();
       if (!signal[kComposite]) {
-        resultSignal[kSourceSignals].add(new SafeWeakRef(signal));
+        const signalWeakRef = new SafeWeakRef(signal);
+        resultSignal[kSourceSignals].add(signalWeakRef);
         signal[kDependantSignals].add(resultSignalWeakRef);
+        dependantSignalsCleanupRegistry.register(resultSignal, signalWeakRef);
       } else if (!signal[kSourceSignals]) {
         continue;
       } else {
-        for (const sourceSignal of signal[kSourceSignals]) {
-          const sourceSignalRef = sourceSignal.deref();
-          if (!sourceSignalRef) {
+        for (const sourceSignalWeakRef of signal[kSourceSignals]) {
+          const sourceSignal = sourceSignalWeakRef.deref();
+          if (!sourceSignal) {
             continue;
           }
-          assert(!sourceSignalRef.aborted);
-          assert(!sourceSignalRef[kComposite]);
+          assert(!sourceSignal.aborted);
+          assert(!sourceSignal[kComposite]);
 
-          if (resultSignal[kSourceSignals].has(sourceSignal)) {
+          if (resultSignal[kSourceSignals].has(sourceSignalWeakRef)) {
             continue;
           }
-          resultSignal[kSourceSignals].add(sourceSignal);
-          sourceSignalRef[kDependantSignals].add(resultSignalWeakRef);
+          resultSignal[kSourceSignals].add(sourceSignalWeakRef);
+          sourceSignal[kDependantSignals].add(resultSignalWeakRef);
+          dependantSignalsCleanupRegistry.register(resultSignal, sourceSignalWeakRef);
         }
       }
     }
diff --git a/test/parallel/test-abortsignal-drop-settled-signals.mjs b/test/parallel/test-abortsignal-drop-settled-signals.mjs
@@ -0,0 +1,122 @@
+// Flags: --expose_gc
+//
+import '../common/index.mjs';
+import { describe, it } from 'node:test';
+
+function makeSubsequentCalls(limit, done, holdReferences = false) {
+  let dependantSymbol;
+  let signalRef;
+  const ac = new AbortController();
+  const retainedSignals = [];
+  const handler = () => { };
+
+  function run(iteration) {
+    if (iteration > limit) {
+      // This setImmediate is necessary to ensure that in the last iteration the remaining signal is GCed (if not
+      // retained)
+      setImmediate(() => {
+        global.gc();
+        done(ac.signal, dependantSymbol);
+      });
+      return;
+    }
+
+    if (holdReferences) {
+      retainedSignals.push(AbortSignal.any([ac.signal]));
+    } else {
+      // Using a WeakRef to avoid retaining information that will interfere with the test
+      signalRef = new WeakRef(AbortSignal.any([ac.signal]));
+      signalRef.deref().addEventListener('abort', handler);
+    }
+
+    dependantSymbol ??= Object.getOwnPropertySymbols(ac.signal).find(
+      (s) => s.toString() === 'Symbol(kDependantSignals)'
+    );
+
+    setImmediate(() => {
+      // Removing the event listener at some moment in the future
+      // Which will then allow the signal to be GCed
+      signalRef?.deref()?.removeEventListener('abort', handler);
+      run(iteration + 1);
+    });
+  }
+
+  run(1);
+};
+
+function runShortLivedSourceSignal(limit, done) {
+  const signalRefs = new Set();
+
+  function run(iteration) {
+    if (iteration > limit) {
+      global.gc();
+      done(signalRefs);
+      return;
+    }
+
+    const ac = new AbortController();
+    signalRefs.add(new WeakRef(ac.signal));
+    AbortSignal.any([ac.signal]);
+
+    setImmediate(() => run(iteration + 1));
+  }
+
+  run(1);
+};
+
+const limit = 10_000;
+
+describe('when there is a long-lived signal', () => {
+  it('drops settled dependant signals', (t, done) => {
+    makeSubsequentCalls(limit, (signal, depandantSignalsKey) => {
+      setImmediate(() => {
+        t.assert.strictEqual(signal[depandantSignalsKey].size, 0);
+        done();
+      });
+    });
+  });
+
+  it('keeps all active dependant signals', (t, done) => {
+    makeSubsequentCalls(limit, (signal, depandantSignalsKey) => {
+      t.assert.strictEqual(signal[depandantSignalsKey].size, limit);
+
+      done();
+    }, true);
+  });
+});
+
+it('does not prevent source signal from being GCed if it is short-lived', (t, done) => {
+  runShortLivedSourceSignal(limit, (signalRefs) => {
+    setImmediate(() => {
+      const unGCedSignals = [...signalRefs].filter((ref) => ref.deref());
+
+      t.assert.strictEqual(unGCedSignals.length, 0);
+      done();
+    });
+  });
+});
+
+it('drops settled dependant signals when signal is composite', (t, done) => {
+  const controllers = Array.from({ length: 2 }, () => new AbortController());
+  const composedSignal1 = AbortSignal.any([controllers[0].signal]);
+  const composedSignalRef = new WeakRef(AbortSignal.any([composedSignal1, controllers[1].signal]));
+
+  const kDependantSignals = Object.getOwnPropertySymbols(controllers[0].signal).find(
+    (s) => s.toString() === 'Symbol(kDependantSignals)'
+  );
+
+  setImmediate(() => {
+    global.gc();
+
+    t.assert.strictEqual(composedSignalRef.deref(), undefined);
+    t.assert.strictEqual(controllers[0].signal[kDependantSignals].size, 2);
+    t.assert.strictEqual(controllers[1].signal[kDependantSignals].size, 1);
+
+    setImmediate(() => {
+      t.assert.strictEqual(controllers[0].signal[kDependantSignals].size, 0);
+      t.assert.strictEqual(controllers[1].signal[kDependantSignals].size, 0);
+
+      done();
+    });
+  });
+});
