diagnostics_channel: fix ref counting bug when reaching zero subscribers

PR-URL: #47520 Reviewed-By: Robert Nagy <[email protected]> Reviewed-By: Rafael Gonzaga <[email protected]> Reviewed-By: Gerhard Stöbich <[email protected]>
nodejs · Nov 10, 2023 · 88f9ebb · 88f9ebb
1 parent 4f78233
commit 88f9ebb
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 21 deletions.
diff --git a/lib/diagnostics_channel.js b/lib/diagnostics_channel.js
@@ -4,6 +4,7 @@ const {
   ArrayPrototypeIndexOf,
   ArrayPrototypePush,
   ArrayPrototypeSplice,
+  SafeFinalizationRegistry,
   ObjectGetPrototypeOf,
   ObjectSetPrototypeOf,
   Promise,
@@ -28,14 +29,29 @@ const { triggerUncaughtException } = internalBinding('errors');
 
 const { WeakReference } = internalBinding('util');
 
-function decRef(channel) {
-  if (channels.get(channel.name).decRef() === 0) {
-    channels.delete(channel.name);
+// Can't delete when weakref count reaches 0 as it could increment again.
+// Only GC can be used as a valid time to clean up the channels map.
+class WeakRefMap extends SafeMap {
+  #finalizers = new SafeFinalizationRegistry((key) => {
+    this.delete(key);
+  });
+
+  set(key, value) {
+    this.#finalizers.register(value, key);
+    return super.set(key, new WeakReference(value));
   }
-}
 
-function incRef(channel) {
-  channels.get(channel.name).incRef();
+  get(key) {
+    return super.get(key)?.get();
+  }
+
+  incRef(key) {
+    return super.get(key)?.incRef();
+  }
+
+  decRef(key) {
+    return super.get(key)?.decRef();
+  }
 }
 
 function markActive(channel) {
@@ -80,7 +96,7 @@ class ActiveChannel {
   subscribe(subscription) {
     validateFunction(subscription, 'subscription');
     ArrayPrototypePush(this._subscribers, subscription);
-    incRef(this);
+    channels.incRef(this.name);
   }
 
   unsubscribe(subscription) {
@@ -89,15 +105,15 @@ class ActiveChannel {
 
     ArrayPrototypeSplice(this._subscribers, index, 1);
 
-    decRef(this);
+    channels.decRef(this.name);
     maybeMarkInactive(this);
 
     return true;
   }
 
   bindStore(store, transform) {
     const replacing = this._stores.has(store);
-    if (!replacing) incRef(this);
+    if (!replacing) channels.incRef(this.name);
     this._stores.set(store, transform);
   }
 
@@ -108,7 +124,7 @@ class ActiveChannel {
 
     this._stores.delete(store);
 
-    decRef(this);
+    channels.decRef(this.name);
     maybeMarkInactive(this);
 
     return true;
@@ -153,7 +169,7 @@ class Channel {
     this._stores = undefined;
     this.name = name;
 
-    channels.set(name, new WeakReference(this));
+    channels.set(name, this);
   }
 
   static [SymbolHasInstance](instance) {
@@ -191,12 +207,10 @@ class Channel {
   }
 }
 
-const channels = new SafeMap();
+const channels = new WeakRefMap();
 
 function channel(name) {
-  let channel;
-  const ref = channels.get(name);
-  if (ref) channel = ref.get();
+  const channel = channels.get(name);
   if (channel) return channel;
 
   if (typeof name !== 'string' && typeof name !== 'symbol') {
@@ -215,12 +229,8 @@ function unsubscribe(name, subscription) {
 }
 
 function hasSubscribers(name) {
-  let channel;
-  const ref = channels.get(name);
-  if (ref) channel = ref.get();
-  if (!channel) {
-    return false;
-  }
+  const channel = channels.get(name);
+  if (!channel) return false;
 
   return channel.hasSubscribers;
 }

diff --git a/test/parallel/test-diagnostics-channel-pub-sub.js b/test/parallel/test-diagnostics-channel-pub-sub.js
@@ -42,3 +42,10 @@ assert.ok(!dc.unsubscribe(name, subscriber));
 assert.throws(() => {
   dc.subscribe(name, null);
 }, { code: 'ERR_INVALID_ARG_TYPE' });
+
+// Reaching zero subscribers should not delete from the channels map as there
+// will be no more weakref to incRef if another subscribe happens while the
+// channel object itself exists.
+channel.subscribe(subscriber);
+channel.unsubscribe(subscriber);
+channel.subscribe(subscriber);