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);