diff --git a/lib/diagnostics_channel.js b/lib/diagnostics_channel.js index e8c5529167a078..d7558a88cc7e3f 100644 --- 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,7 +105,7 @@ class ActiveChannel { ArrayPrototypeSplice(this._subscribers, index, 1); - decRef(this); + channels.decRef(this.name); maybeMarkInactive(this); return true; @@ -97,7 +113,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); } @@ -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 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);