Stop leaking change token registrations

[ismailbennani]

Fixes #7673

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.

Microsoft Reviewers: Open in CodeFlow

[ismailbennani]

@dotnet-policy-service agree

[Copilot]

Pull request overview

This PR fixes a memory leak in ServiceDiscovery endpoint resolution caused by wrapping provider change tokens in CompositeChangeToken, which can keep proxy registrations alive until the composite signals (problematic when a source token never signals, e.g., typical configuration reload tokens). It replaces that behavior with a new LinkedChangeToken that makes consumer registrations directly against the underlying sources so disposing the consumer registration releases the source registrations immediately.

Changes:

• Replace CompositeChangeToken usage in ServiceEndpointBuilder.Build() with either the single original token (when only one exists) or a new LinkedChangeToken (when multiple exist).
• Introduce LinkedChangeToken (internal) to link multiple change tokens without retaining proxy registrations beyond consumer lifetime.
• Add comprehensive unit/regression tests validating callback/HasChanged behavior and that repeated watcher lifecycles do not accumulate registrations.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
test/Libraries/Microsoft.Extensions.ServiceDiscovery.Tests/LinkedChangeTokenTests.cs	Adds behavioral + regression tests ensuring linked tokens and watcher lifecycles don’t accumulate registrations.
src/Libraries/Microsoft.Extensions.ServiceDiscovery/ServiceEndpointBuilder.cs	Stops wrapping tokens in CompositeChangeToken; returns single token as-is or uses LinkedChangeToken for multiple tokens.
src/Libraries/Microsoft.Extensions.ServiceDiscovery/Internal/LinkedChangeToken.cs	Implements a linked token that ties source registrations to the consumer registration lifetime to avoid long-lived proxy registrations.

[jeffhandley]

Request changes

There is a race between Dispose and OnSourceSignaled in this implementation. If OnSourceSignaled wins the callback exchange at line 158, Dispose can still clear _state here before the callback reads it at line 163. The callback is then invoked with null state.

This can cause the watcher callback in ServiceEndpointWatcher.cs:175 to throw when it casts and uses the state as a ServiceEndpointWatcher.

Please make callback and state ownership atomic. Disposal should clear the state only if it successfully suppresses the callback; otherwise, the signaling path must retain and consume the original state.

I tried the following regression test locally, including the ManualChangeToken helper below, and it reproduced the suspected race on the current implementation:

[Fact]
public async Task DisposeRacingWithSource_DoesNotInvokeCallbackWithNullState()
{
    const int iterations = 2000;
    var failures = 0;

    for (int i = 0; i < iterations; i++)
    {
        var source = new ManualChangeToken();
        var token = new LinkedChangeToken([source]);
        var expectedState = new object();

        using var registration = token.RegisterChangeCallback(state =>
        {
            if (state is null)
            {
                Interlocked.Increment(ref failures);
            }
        }, expectedState);

        await Task.WhenAll(
            Task.Run(source.Fire),
            Task.Run(registration.Dispose));
    }

    Assert.Equal(0, failures);
}

private sealed class ManualChangeToken : IChangeToken
{
    private Action<object?>? _callback;
    private object? _state;

    public bool ActiveChangeCallbacks => true;
    public bool HasChanged => false;

    public IDisposable RegisterChangeCallback(Action<object?> callback, object? state)
    {
        Interlocked.Exchange(ref _callback, callback);
        _state = state;

        return new Registration(() =>
        {
            Interlocked.Exchange(ref _callback, null);
            _state = null;
        });
    }

    public void Fire()
    {
        var callback = Volatile.Read(ref _callback);
        var state = _state;
        callback?.Invoke(state);
    }

    private sealed class Registration(Action dispose) : IDisposable
    {
        public void Dispose() => dispose();
    }
}

Please add an equivalent regression test, preferably with a deterministic synchronization point around the callback/state handoff so the race does not depend solely on iteration count, and update the implementation so the callback always receives the state supplied during registration.

[ismailbennani]

The only way I found to force the disposal between the exchange and the read of the state was to introduce an event callback between them. I made it internal and explained in its documentation that it is there only for tests.