From c02a7e7e937504d1c32e2b20330d5feabad2f010 Mon Sep 17 00:00:00 2001
From: Stephen Belanger <stephen.belanger@datadoghq.com>
Date: Thu, 13 Apr 2023 06:40:50 -0700
Subject: [PATCH] diagnostics_channel: fix ref counting bug when reaching zero
 subscribers
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PR-URL: https://github.com/nodejs/node/pull/47520
Reviewed-By: Robert Nagy <ronagy@icloud.com>
Reviewed-By: Rafael Gonzaga <rafael.nunu@hotmail.com>
Reviewed-By: Gerhard Stöbich <deb2001-github@yahoo.de>
---
 lib/diagnostics_channel.js                    | 52 +++++++++++--------
 .../test-diagnostics-channel-pub-sub.js       |  7 +++
 2 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/lib/diagnostics_channel.js b/lib/diagnostics_channel.js
index be50c3b8f6e48d..b399c1e2f87ad3 100644
--- a/lib/diagnostics_channel.js
+++ b/lib/diagnostics_channel.js
@@ -5,6 +5,7 @@ const {
   ArrayPrototypeIndexOf,
   ArrayPrototypePush,
   ArrayPrototypeSplice,
+  SafeFinalizationRegistry,
   ObjectGetPrototypeOf,
   ObjectSetPrototypeOf,
   Promise,
@@ -29,14 +30,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) {
@@ -81,7 +97,7 @@ class ActiveChannel {
   subscribe(subscription) {
     validateFunction(subscription, 'subscription');
     ArrayPrototypePush(this._subscribers, subscription);
-    incRef(this);
+    channels.incRef(this.name);
   }
 
   unsubscribe(subscription) {
@@ -90,7 +106,7 @@ class ActiveChannel {
 
     ArrayPrototypeSplice(this._subscribers, index, 1);
 
-    decRef(this);
+    channels.decRef(this.name);
     maybeMarkInactive(this);
 
     return true;
@@ -98,7 +114,7 @@ class ActiveChannel {
 
   bindStore(store, transform) {
     const replacing = this._stores.has(store);
-    if (!replacing) incRef(this);
+    if (!replacing) channels.incRef(this.name);
     this._stores.set(store, transform);
   }
 
@@ -109,7 +125,7 @@ class ActiveChannel {
 
     this._stores.delete(store);
 
-    decRef(this);
+    channels.decRef(this.name);
     maybeMarkInactive(this);
 
     return true;
@@ -154,7 +170,7 @@ class Channel {
     this._stores = undefined;
     this.name = name;
 
-    channels.set(name, new WeakReference(this));
+    channels.set(name, this);
   }
 
   static [SymbolHasInstance](instance) {
@@ -192,12 +208,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') {
@@ -216,12 +230,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
index 2317d90dbbc554..a7232ab58ce8a5 100644
--- 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);