From 56815474627b0facdf0de775de2236dddfbc93ee Mon Sep 17 00:00:00 2001 From: Ismail Bennani Date: Mon, 3 Aug 2026 01:40:26 +0200 Subject: [PATCH 1/3] Stop leaking change token registrations ServiceEndpointBuilder.Build wrapped the change tokens contributed by the endpoint providers in a CompositeChangeToken. That composite registers on its inner tokens on behalf of its consumers and releases those registrations only when it signals, so linking a token which never signals, such as the reload token of a configuration which is never reloaded, keeps the composite and everything it references alive for the lifetime of that source. ServiceEndpointResolver evicts a watcher whose service name has not been resolved in the last ten seconds, and the next resolution creates a new one. Each of those lifecycles added a registration on the application's configuration reload token which was never released, even though the watcher disposes its own registration correctly both on refresh and in DisposeAsync. Replace the composite with LinkedChangeToken, which registers the consumer's callback directly on each source and hands those source registrations to the consumer's own registration. Releasing them then needs nothing from the owner of the token: the consumer disposing its registration, which consumers already do, is enough. A lone change token, which is what the Configuration and PassThrough providers produce together, is now returned as-is with no wrapper at all. Fixes #7673 Co-Authored-By: Claude Opus 5 (1M context) --- .../Internal/LinkedChangeToken.cs | 187 +++++++++ .../ServiceEndpointBuilder.cs | 9 +- .../LinkedChangeTokenTests.cs | 390 ++++++++++++++++++ 3 files changed, 585 insertions(+), 1 deletion(-) create mode 100644 src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs create mode 100644 test/Libraries/Microsoft.Extensions.ServiceDiscovery.Tests/LinkedChangeTokenTests.cs diff --git a/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs b/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs new file mode 100644 index 00000000000..c4ef8cacb10 --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs @@ -0,0 +1,187 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using Microsoft.Extensions.Primitives; + +namespace Microsoft.Extensions.ServiceDiscovery.Internal; + +/// +/// An which signals when any of the change tokens it is linked to signals. +/// +/// +/// +/// This serves the same purpose as , but it holds callbacks on its sources +/// only for as long as a consumer is listening. registers on its sources on +/// behalf of its consumers and releases those registrations only when it signals, so linking a token which +/// never signals roots the composite and everything it references for the lifetime of that source. +/// +/// +/// Here a consumer's registration is made directly on each source and owns those source registrations, so a +/// consumer releasing its own registration, which is what consumers already do, is all that is needed. Nothing +/// has to remember to release this token. +/// +/// +/// Registering directly on the sources also means the callback behaviour of this token is whatever its sources +/// provide, rather than being reshaped by an intermediate . As +/// allows, callbacks are best effort; , which polls the +/// sources, is the reliable way to observe a change. +/// +/// +internal sealed class LinkedChangeToken : IChangeToken +{ + private readonly IReadOnlyList _sources; + private volatile bool _hasChanged; + + /// + /// Initializes a new instance. + /// + /// The change tokens to link to. + public LinkedChangeToken(IReadOnlyList sources) + { + ArgumentNullException.ThrowIfNull(sources); + + _sources = sources; + + for (var i = 0; i < sources.Count; i++) + { + if (sources[i].ActiveChangeCallbacks) + { + ActiveChangeCallbacks = true; + break; + } + } + } + + /// + /// + /// Callbacks are raised only by sources which raise them. Changes to the other sources are observed only by + /// polling , which matches . + /// + public bool ActiveChangeCallbacks { get; } + + /// + public bool HasChanged + { + get + { + if (_hasChanged) + { + return true; + } + + for (var i = 0; i < _sources.Count; i++) + { + if (_sources[i].HasChanged) + { + _hasChanged = true; + return true; + } + } + + return false; + } + } + + /// + public IDisposable RegisterChangeCallback(Action callback, object? state) + { + var registration = new Registration(this, callback, state); + + // Linked after construction rather than in the constructor, because a source which has already signaled + // raises the callback during linking and must not observe a partially constructed registration. + registration.LinkToSources(); + return registration; + } + + /// + /// A consumer's registration, which holds that consumer's registration on each of the sources and releases + /// them when it is disposed or when one of the sources signals. + /// + private sealed class Registration : IDisposable + { + // Cached so that registering on a source does not allocate a delegate. The callback shape is dictated by + // IChangeToken.RegisterChangeCallback; passing the registration as its state keeps it closure-free. + private static readonly Action s_onSourceSignaled = static state => ((Registration)state!).OnSourceSignaled(); + + private readonly LinkedChangeToken _token; + private readonly IDisposable?[] _sourceRegistrations; + private Action? _callback; + private object? _state; + + public Registration(LinkedChangeToken token, Action callback, object? state) + { + _token = token; + _callback = callback; + _state = state; + _sourceRegistrations = new IDisposable?[token._sources.Count]; + } + + /// + /// Registers this consumer's callback on each source which raises callbacks. + /// + public void LinkToSources() + { + var sources = _token._sources; + + for (var i = 0; i < sources.Count; i++) + { + if (sources[i].ActiveChangeCallbacks) + { + // A source which has already signaled may raise the callback here, synchronously. Sources + // backed by a CancellationToken do, but IChangeToken does not require it, so a change is + // only reliably observed by polling HasChanged. + _sourceRegistrations[i] = sources[i].RegisterChangeCallback(s_onSourceSignaled, this); + } + } + + // A null callback means a source signaled, or the consumer disposed, while this loop was still + // running, so Release could not see every registration it was meant to release. Release the rest. + if (Volatile.Read(ref _callback) is null) + { + Release(); + } + } + + public void Dispose() + { + // Cleared before releasing, so that a concurrent LinkToSources sees that it has to release the + // registrations it makes after this point. + Interlocked.Exchange(ref _callback, null); + _state = null; + Release(); + } + + private void OnSourceSignaled() + { + // Only the first source to signal raises the consumer's callback, and a consumer which has disposed + // its registration is not called at all. + if (Interlocked.Exchange(ref _callback, null) is not { } callback) + { + return; + } + + var state = _state; + _state = null; + _token._hasChanged = true; + + try + { + callback(state); + } + finally + { + Release(); + } + } + + private void Release() + { + for (var i = 0; i < _sourceRegistrations.Length; i++) + { + // Exchanged so that releasing more than once, which linking and a signaling source can both + // cause, disposes each source registration exactly once. + Interlocked.Exchange(ref _sourceRegistrations[i], null)?.Dispose(); + } + } + } +} diff --git a/src/Libraries/Microsoft.Extensions.ServiceDiscovery/ServiceEndpointBuilder.cs b/src/Libraries/Microsoft.Extensions.ServiceDiscovery/ServiceEndpointBuilder.cs index 947f24b2f81..cbbe86c4c84 100644 --- a/src/Libraries/Microsoft.Extensions.ServiceDiscovery/ServiceEndpointBuilder.cs +++ b/src/Libraries/Microsoft.Extensions.ServiceDiscovery/ServiceEndpointBuilder.cs @@ -3,6 +3,7 @@ using Microsoft.AspNetCore.Http.Features; using Microsoft.Extensions.Primitives; +using Microsoft.Extensions.ServiceDiscovery.Internal; namespace Microsoft.Extensions.ServiceDiscovery; @@ -40,7 +41,13 @@ public void AddChangeToken(IChangeToken changeToken) /// The service endpoint source. public ServiceEndpointSource Build() { - return new ServiceEndpointSource(_endpoints, new CompositeChangeToken(_changeTokens), _features); + // A single change token, which is the common case, is returned as-is: there is nothing to link, and the + // consumer's registration on the token is its own to release. + var changeToken = _changeTokens.Count == 1 + ? _changeTokens[0] + : new LinkedChangeToken(_changeTokens); + + return new ServiceEndpointSource(_endpoints, changeToken, _features); } } diff --git a/test/Libraries/Microsoft.Extensions.ServiceDiscovery.Tests/LinkedChangeTokenTests.cs b/test/Libraries/Microsoft.Extensions.ServiceDiscovery.Tests/LinkedChangeTokenTests.cs new file mode 100644 index 00000000000..35e7fab0296 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.ServiceDiscovery.Tests/LinkedChangeTokenTests.cs @@ -0,0 +1,390 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System.Diagnostics.CodeAnalysis; +using System.Net; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Primitives; +using Microsoft.Extensions.ServiceDiscovery.Internal; +using Xunit; + +namespace Microsoft.Extensions.ServiceDiscovery.Tests; + +/// +/// Tests for and for the change token registrations which endpoint resolution +/// leaves behind on the tokens its providers contribute. +/// +public class LinkedChangeTokenTests +{ + [Fact] + public void ActiveChangeCallbacks_IsTrue_WhenAnySourceRaisesCallbacks() + { + Assert.False(new LinkedChangeToken([]).ActiveChangeCallbacks); + Assert.False(new LinkedChangeToken([new PassiveChangeToken()]).ActiveChangeCallbacks); + Assert.True(new LinkedChangeToken([new PassiveChangeToken(), new TrackingChangeToken()]).ActiveChangeCallbacks); + } + + [Fact] + public void HasChanged_PollsSources_WhichDoNotRaiseCallbacks() + { + var passive = new PassiveChangeToken(); + var token = new LinkedChangeToken([passive]); + + Assert.False(token.HasChanged); + + passive.HasChanged = true; + Assert.True(token.HasChanged); + } + + [Fact] + public void RegisterChangeCallback_IsInvoked_WhenAnySourceSignals() + { + var first = new TrackingChangeToken(); + var second = new TrackingChangeToken(); + var token = new LinkedChangeToken([first, second]); + + var signaled = 0; + using var registration = token.RegisterChangeCallback(_ => signaled++, null); + + Assert.Equal(0, signaled); + Assert.False(token.HasChanged); + + second.Signal(); + + Assert.Equal(1, signaled); + Assert.True(token.HasChanged); + + // Signalling the other source must not raise the callback a second time. + first.Signal(); + Assert.Equal(1, signaled); + } + + [Fact] + public void SourceRegistrations_AreReleased_WhenLastConsumerDisposes() + { + var first = new TrackingChangeToken(); + var second = new TrackingChangeToken(); + var token = new LinkedChangeToken([first, second]); + + var registration = token.RegisterChangeCallback(static _ => { }, null); + + Assert.Equal(1, first.OutstandingRegistrations); + Assert.Equal(1, second.OutstandingRegistrations); + + registration.Dispose(); + + // Nothing is listening any more, so the token must stop listening to its sources. This is the leak the + // type exists to avoid: CompositeChangeToken would hold these until it signaled, which for a source + // which never signals is never. + Assert.Equal(0, first.OutstandingRegistrations); + Assert.Equal(0, second.OutstandingRegistrations); + } + + [Fact] + public void SourceRegistrations_AreRetained_WhileAnyConsumerRemains() + { + var source = new TrackingChangeToken(); + var token = new LinkedChangeToken([source]); + + var first = token.RegisterChangeCallback(static _ => { }, null); + var second = token.RegisterChangeCallback(static _ => { }, null); + + // Each consumer holds its own registration on the source. + Assert.Equal(2, source.OutstandingRegistrations); + + first.Dispose(); + Assert.Equal(1, source.OutstandingRegistrations); + + second.Dispose(); + Assert.Equal(0, source.OutstandingRegistrations); + } + + [Fact] + public void Constructor_Throws_WhenSourcesIsNull() + => Assert.Throws(() => new LinkedChangeToken(null!)); + + [Fact] + public void SourceRegistrations_AreReleased_WhenSignaled() + { + var source = new TrackingChangeToken(); + var token = new LinkedChangeToken([source]); + + using var registration = token.RegisterChangeCallback(static _ => { }, null); + Assert.Equal(1, source.OutstandingRegistrations); + + source.Signal(); + + Assert.True(token.HasChanged); + Assert.Equal(0, source.OutstandingRegistrations); + } + + [Fact] + public void DisposingARegistrationTwice_ReleasesItsSourcesOnce() + { + var source = new TrackingChangeToken(); + var token = new LinkedChangeToken([source]); + + var first = token.RegisterChangeCallback(static _ => { }, null); + var second = token.RegisterChangeCallback(static _ => { }, null); + + first.Dispose(); + first.Dispose(); + + // The second consumer is still listening, so a double dispose of the first must not have released it. + Assert.Equal(1, source.OutstandingRegistrations); + + second.Dispose(); + Assert.Equal(0, source.OutstandingRegistrations); + } + + [Fact] + public void RegisteringAgain_ListensToSourcesAgain() + { + var source = new TrackingChangeToken(); + var token = new LinkedChangeToken([source]); + + token.RegisterChangeCallback(static _ => { }, null).Dispose(); + Assert.Equal(0, source.OutstandingRegistrations); + + var signaled = false; + using var registration = token.RegisterChangeCallback(_ => signaled = true, null); + Assert.Equal(1, source.OutstandingRegistrations); + + source.Signal(); + Assert.True(signaled); + } + + [Fact] + public void RegisteringAfterASourceHasSignaled_PropagatesWhateverTheSourceDoes() + { + var source = new TrackingChangeToken(); + var token = new LinkedChangeToken([source]); + + source.Signal(); + + var signaled = false; + using var registration = token.RegisterChangeCallback(_ => signaled = true, null); + + // Raising the callback when registering on a source which has already signaled is not required by + // IChangeToken. A source backed by a CancellationToken does it, and because this token registers on its + // sources directly rather than proxying them, that behaviour reaches the consumer unchanged. + Assert.True(signaled); + + // HasChanged, unlike the callback, is reliable whatever the source does, because it polls. + Assert.True(token.HasChanged); + + // The callback ran during registration, so the source registration it made was released there too. + Assert.Equal(0, source.OutstandingRegistrations); + } + + [Fact] + public void HasChanged_IsTrue_WhenAPollOnlySourceChangedBeforeRegistering() + { + // A source which raises no callbacks cannot notify a consumer at all, so polling is the only way its + // change is ever seen. This is the case the contract has in mind when it says callbacks are best effort. + var passive = new PassiveChangeToken(); + var token = new LinkedChangeToken([passive, new TrackingChangeToken()]); + + passive.HasChanged = true; + + var signaled = false; + using var registration = token.RegisterChangeCallback(_ => signaled = true, null); + + Assert.False(signaled); + Assert.True(token.HasChanged); + } + + [Fact] + public void ASourceSignalingWhileLinking_DoesNotOrphanTheLaterSourceRegistrations() + { + var first = new TrackingChangeToken(); + var second = new TrackingChangeToken(); + var third = new TrackingChangeToken(); + + // Signaled up front, so that registering on it raises the callback from inside the linking loop. The + // release that triggers runs before the later sources have been registered on, so it cannot see them and + // the check made once linking finishes is the only thing which can release them. + first.Signal(); + + var token = new LinkedChangeToken([first, second, third]); + using var registration = token.RegisterChangeCallback(static _ => { }, null); + + // The later sources were registered on... + Assert.Equal(1, second.TotalRegistrations); + Assert.Equal(1, third.TotalRegistrations); + + // ...and none of those registrations was left behind. + Assert.Equal(0, first.OutstandingRegistrations); + Assert.Equal(0, second.OutstandingRegistrations); + Assert.Equal(0, third.OutstandingRegistrations); + } + + [Fact] + public async Task ASourceSignalingWhileACallbackRuns_NeitherInvokesItAgainNorBlocks() + { + // How long a step which should complete immediately is given before it is treated as blocked. A passing + // run never waits for it, since every wait below ends as soon as the step it waits for happens; it only + // bounds how long a regression takes to fail. + var blockedTimeout = TimeSpan.FromSeconds(5); + + using var first = new CancellationTokenSource(); + using var second = new CancellationTokenSource(); + + // Real cancellation-backed tokens rather than the tracking fake, which is not built for concurrent use. + var token = new LinkedChangeToken([new CancellationChangeToken(first.Token), new CancellationChangeToken(second.Token)]); + + using var callbackRunning = new ManualResetEventSlim(false); + using var releaseCallback = new ManualResetEventSlim(false); + var invocations = 0; + + using var registration = token.RegisterChangeCallback( + _ => + { + Interlocked.Increment(ref invocations); + callbackRunning.Set(); + releaseCallback.Wait(blockedTimeout); + }, + null); + + // Signalling the first source raises the callback, which parks. Waiting for it makes the overlap below a + // fact rather than something the thread pool may or may not produce. + var firstChange = Task.Run(first.Cancel); + Assert.True(callbackRunning.Wait(blockedTimeout), "The callback was never raised."); + + // The second source now signals while that callback is definitely still running. It must not wait on it, + // which it would if this token serialised callbacks behind a lock, and nothing here would release it. + var secondChange = Task.Run(second.Cancel); + var secondCompleted = await Task.WhenAny(secondChange, Task.Delay(blockedTimeout)); + Assert.True(secondCompleted == secondChange, "Signalling a source blocked behind a callback which was still running."); + + releaseCallback.Set(); + var firstCompleted = await Task.WhenAny(firstChange, Task.Delay(blockedTimeout)); + Assert.True(firstCompleted == firstChange, "Signalling a source did not complete once its callback returned."); + + await Task.WhenAll(firstChange, secondChange); + + Assert.Equal(1, Volatile.Read(ref invocations)); + Assert.True(token.HasChanged); + } + + [Theory] + [InlineData(1)] // One token is returned by the builder as-is, with no linking involved. + [InlineData(2)] // Two or more are linked. + public async Task WatcherLifecycles_DoNotAccumulateRegistrationsOnProviderChangeTokens(int changeTokenCount) + { + // Regression test for https://github.com/dotnet/extensions/issues/7673: a watcher is created and disposed + // every time the resolver evicts an idle service name, and each lifecycle used to add a registration on + // the provider's change token which was never released. + var sources = Enumerable.Range(0, changeTokenCount).Select(_ => new TrackingChangeToken()).ToArray(); + var provider = new FakeEndpointProvider(builder => + { + foreach (var source in sources) + { + builder.AddChangeToken(source); + } + + builder.Endpoints.Add(ServiceEndpoint.Create(new IPEndPoint(IPAddress.Loopback, 8080))); + }); + + var services = new ServiceCollection() + .AddSingleton(new FakeEndpointProviderFactory(provider)) + .AddServiceDiscoveryCore() + .BuildServiceProvider(); + var watcherFactory = services.GetRequiredService(); + + const int Lifecycles = 5; + for (var i = 0; i < Lifecycles; i++) + { + ServiceEndpointWatcher watcher; + await using ((watcher = watcherFactory.CreateWatcher("http://basket")).ConfigureAwait(false)) + { + var endpoints = await watcher.GetEndpointsAsync(CancellationToken.None); + Assert.Single(endpoints.Endpoints); + } + } + + foreach (var source in sources) + { + // Sanity check that the watcher did register, so that the assertion below is meaningful. + Assert.Equal(Lifecycles, source.TotalRegistrations); + Assert.Equal(0, source.OutstandingRegistrations); + } + } + + private sealed class FakeEndpointProviderFactory(IServiceEndpointProvider provider) : IServiceEndpointProviderFactory + { + public bool TryCreateProvider(ServiceEndpointQuery query, [NotNullWhen(true)] out IServiceEndpointProvider? resolver) + { + resolver = provider; + return true; + } + } + + private sealed class FakeEndpointProvider(Action populate) : IServiceEndpointProvider + { + public ValueTask PopulateAsync(IServiceEndpointBuilder endpoints, CancellationToken cancellationToken) + { + populate(endpoints); + return default; + } + + public ValueTask DisposeAsync() => default; + } + + /// + /// A change token which raises callbacks and keeps count of the registrations which have not been released. + /// + private sealed class TrackingChangeToken : IChangeToken + { + private readonly CancellationTokenSource _cts = new(); + + public bool ActiveChangeCallbacks => true; + + public bool HasChanged => _cts.IsCancellationRequested; + + /// Gets the number of registrations which have been made and not yet released. + public int OutstandingRegistrations { get; private set; } + + /// Gets the number of registrations which have been made. + public int TotalRegistrations { get; private set; } + + public void Signal() => _cts.Cancel(); + + public IDisposable RegisterChangeCallback(Action callback, object? state) + { + OutstandingRegistrations++; + TotalRegistrations++; + return new Registration(this, _cts.Token.Register(callback, state)); + } + + private sealed class Registration(TrackingChangeToken owner, CancellationTokenRegistration registration) : IDisposable + { + private bool _disposed; + + public void Dispose() + { + if (_disposed) + { + return; + } + + _disposed = true; + registration.Dispose(); + owner.OutstandingRegistrations--; + } + } + } + + /// + /// A change token whose changes are observable only by polling . + /// + private sealed class PassiveChangeToken : IChangeToken + { + public bool ActiveChangeCallbacks => false; + + public bool HasChanged { get; set; } + + public IDisposable RegisterChangeCallback(Action callback, object? state) + => throw new InvalidOperationException("This token does not raise callbacks and must not be registered on."); + } +} From e2fb4fc125ae000eedc4c4080bd93fadad866682 Mon Sep 17 00:00:00 2001 From: Ismail Bennani Date: Tue, 11 Aug 2026 22:13:03 +0200 Subject: [PATCH 2/3] Add a test which forces disposal into a signal's handover A consumer's callback and the state it registered with belong together: whichever of a signalling source and a disposing consumer claims the callback owns that state too. Disposal currently drops the state whether or not it claimed the callback, so a source which has just claimed it finds the state gone by the time it reads it and calls the consumer with null. ServiceEndpointWatcher's callback casts its state to a watcher, so that throws. Claiming the callback and reading the state are adjacent instructions, which leaves a test nothing to interleave with from the outside. Give the token an OnCallbackClaimed callout at that point, which only tests set, and the interleaving becomes reachable: a test disposes from it, and the two paths then run in exactly that order on one thread, with no dependence on how threads happen to be scheduled. A counter in the test asserts that the moment was reached, so the test cannot quietly assert nothing if that ordering ever stops happening. Both orders are covered: the signal claiming the callback first, which is the failing one, and the consumer disposing first, where the signal has nothing left to raise. The first test fails against the current implementation on every target framework, on every run. Co-Authored-By: Claude Opus 5 (1M context) --- .../Internal/LinkedChangeToken.cs | 14 +++++ .../LinkedChangeTokenTests.cs | 59 +++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs b/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs index c4ef8cacb10..44943401918 100644 --- a/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs +++ b/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs @@ -82,6 +82,18 @@ public bool HasChanged } } + /// + /// Gets or sets a callout made by a signalling source once it has claimed a consumer's callback and before it + /// reads the state that callback was registered with. + /// + /// + /// Only tests set this; it is null everywhere else. Those two steps are adjacent instructions, so setting this + /// is the only way a test can decide how the race between Registration.Dispose and + /// Registration.OnSourceSignaled comes out, and therefore the only way a test can show that this race is not a + /// problem. + /// + internal Action? OnCallbackClaimed { get; set; } + /// public IDisposable RegisterChangeCallback(Action callback, object? state) { @@ -160,6 +172,8 @@ private void OnSourceSignaled() return; } + _token.OnCallbackClaimed?.Invoke(); + var state = _state; _state = null; _token._hasChanged = true; diff --git a/test/Libraries/Microsoft.Extensions.ServiceDiscovery.Tests/LinkedChangeTokenTests.cs b/test/Libraries/Microsoft.Extensions.ServiceDiscovery.Tests/LinkedChangeTokenTests.cs index 35e7fab0296..b7461ce964f 100644 --- a/test/Libraries/Microsoft.Extensions.ServiceDiscovery.Tests/LinkedChangeTokenTests.cs +++ b/test/Libraries/Microsoft.Extensions.ServiceDiscovery.Tests/LinkedChangeTokenTests.cs @@ -267,6 +267,65 @@ public async Task ASourceSignalingWhileACallbackRuns_NeitherInvokesItAgainNorBlo Assert.True(token.HasChanged); } + [Fact] + public void AConsumerDisposingAfterASignalHasClaimedItsCallback_StillGetsTheStateItRegisteredWith() + { + // A signal has claimed the consumer's callback and not yet read the state it was registered with, and the + // consumer disposes right then. That state belongs to the signal now, so disposal must leave it alone. + // Disposing from OnCallbackClaimed is what puts the two paths in that order. + var source = new TrackingChangeToken(); + var token = new LinkedChangeToken([source]); + var expectedState = new object(); + + IDisposable registration = null!; + object? observedState = null; + var invocations = 0; + var forced = 0; + + token.OnCallbackClaimed = () => + { + forced++; + registration.Dispose(); + }; + + registration = token.RegisterChangeCallback( + state => + { + invocations++; + observedState = state; + }, + expectedState); + + source.Signal(); + + // The interleaving really was forced, rather than the test having quietly asserted nothing. + Assert.Equal(1, forced); + + Assert.Equal(1, invocations); + Assert.Same(expectedState, observedState); + + // Disposal losing the callback does not stop it from releasing what it holds on the sources. + Assert.Equal(0, source.OutstandingRegistrations); + } + + [Fact] + public void ASignalAfterTheConsumerHasDisposed_DoesNotReachIt() + { + // The other order: disposal claimed the callback first, so the signal has nothing left to raise. + var source = new TrackingChangeToken(); + var token = new LinkedChangeToken([source]); + + var invocations = 0; + var registration = token.RegisterChangeCallback(_ => invocations++, new object()); + + registration.Dispose(); + source.Signal(); + + Assert.Equal(0, invocations); + Assert.Equal(0, source.OutstandingRegistrations); + Assert.True(token.HasChanged); + } + [Theory] [InlineData(1)] // One token is returned by the builder as-is, with no linking involved. [InlineData(2)] // Two or more are linked. From 5ba03530bb06efb28c5b94b0f8bf3ed034d47b74 Mon Sep 17 00:00:00 2001 From: Ismail Bennani Date: Tue, 11 Aug 2026 22:33:25 +0200 Subject: [PATCH 3/3] Let claiming the callback claim the state with it Disposal cleared the state whether or not it was the one which had suppressed the callback, so a source which had just claimed that callback could find the state gone by the time it read it and hand the consumer null instead. ServiceEndpointWatcher's callback casts its state to a watcher, so that throws. Drop the state only on the disposal which won the exchange for the callback. The other path keeps it and consumes it, which is what it already does, and that exchange is what orders its read after the write registration made. Co-Authored-By: Claude Opus 5 (1M context) --- .../Internal/LinkedChangeToken.cs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs b/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs index 44943401918..e2b9fafd07d 100644 --- a/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs +++ b/src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs @@ -158,8 +158,14 @@ public void Dispose() { // Cleared before releasing, so that a concurrent LinkToSources sees that it has to release the // registrations it makes after this point. - Interlocked.Exchange(ref _callback, null); - _state = null; + if (Interlocked.Exchange(ref _callback, null) is not null) + { + // Taking the callback is what claims the state, so the state is dropped here only when it was this + // disposal which suppressed the callback. A source signalling concurrently may have taken the + // callback instead, and it has to be able to hand the consumer the state it registered with. + _state = null; + } + Release(); } @@ -174,6 +180,8 @@ private void OnSourceSignaled() _token.OnCallbackClaimed?.Invoke(); + // Taking the callback above claimed the state, so a concurrent disposal cannot clear it from underneath + // this callback, and that exchange orders this read after the write which registration made. var state = _state; _state = null; _token._hasChanged = true;