Promptly remove invalid resolvers (dotnet#3800)

Author: ReubenBond

Directory.Packages.props

@@ -87,6 +87,7 @@
     <PackageVersion Include="Microsoft.Extensions.Options" Version="$(MicrosoftExtensionsOptionsPackageVersion)" />
     <PackageVersion Include="Microsoft.Extensions.Primitives" Version="$(MicrosoftExtensionsPrimitivesPackageVersion)" />
     <PackageVersion Include="Microsoft.Extensions.Http.Resilience" Version="$(MicrosoftExtensionsHttpResiliencePackageVersion)" />
+    <PackageVersion Include="Microsoft.Extensions.TimeProvider.Testing" Version="$(MicrosoftExtensionsTimeProviderTestingVersion)" />
     <!-- external dependencies -->
     <PackageVersion Include="Confluent.Kafka" Version="2.3.0" />
     <PackageVersion Include="Dapper" Version="2.1.37" />

eng/Versions.props

@@ -46,6 +46,7 @@
     <MicrosoftExtensionsDiagnosticsHealthChecksEntityFrameworkCorePackageVersion>8.0.4</MicrosoftExtensionsDiagnosticsHealthChecksEntityFrameworkCorePackageVersion>
     <MicrosoftExtensionsDiagnosticsHealthChecksPackageVersion>8.0.4</MicrosoftExtensionsDiagnosticsHealthChecksPackageVersion>
     <MicrosoftExtensionsFeaturesPackageVersion>8.0.4</MicrosoftExtensionsFeaturesPackageVersion>
+    <MicrosoftExtensionsTimeProviderTestingVersion>8.4.0</MicrosoftExtensionsTimeProviderTestingVersion>
     <!-- EF -->
     <MicrosoftEntityFrameworkCoreCosmosPackageVersion>8.0.4</MicrosoftEntityFrameworkCoreCosmosPackageVersion>
     <MicrosoftEntityFrameworkCoreDesignPackageVersion>8.0.4</MicrosoftEntityFrameworkCoreDesignPackageVersion>

src/Microsoft.Extensions.ServiceDiscovery/Http/HttpServiceEndpointResolver.cs

@@ -57,6 +57,10 @@ public async ValueTask<ServiceEndpoint> GetEndpointAsync(HttpRequestMessage requ
 
                 return endpoint;
             }
+            else
+            {
+                _resolvers.TryRemove(KeyValuePair.Create(resolver.ServiceName, resolver));
+            }
         }
     }
 
@@ -140,6 +144,7 @@ private async Task CleanupResolversAsyncCore()
                 cleanupTasks.Add(resolver.DisposeAsync().AsTask());
             }
         }
+
         if (cleanupTasks is not null)
         {
             await Task.WhenAll(cleanupTasks).ConfigureAwait(false);

src/Microsoft.Extensions.ServiceDiscovery/ServiceEndpointResolver.cs

@@ -49,6 +49,7 @@ public async ValueTask<ServiceEndpointSource> GetEndpointsAsync(string serviceNa
         while (true)
         {
             ObjectDisposedException.ThrowIf(_disposed, this);
+            cancellationToken.ThrowIfCancellationRequested();
             var resolver = _resolvers.GetOrAdd(
                 serviceName,
                 static (name, self) => self.CreateResolver(name),
@@ -64,6 +65,10 @@ public async ValueTask<ServiceEndpointSource> GetEndpointsAsync(string serviceNa
 
                 return result;
             }
+            else
+            {
+                _resolvers.TryRemove(KeyValuePair.Create(resolver.ServiceName, resolver));
+            }
         }
     }
 
@@ -148,6 +153,7 @@ private async Task CleanupResolversAsyncCore()
                 cleanupTasks.Add(resolver.DisposeAsync().AsTask());
             }
         }
+
         if (cleanupTasks is not null)
         {
             await Task.WhenAll(cleanupTasks).ConfigureAwait(false);

src/Microsoft.Extensions.ServiceDiscovery/ServiceEndpointWatcher.cs

@@ -32,6 +32,7 @@ internal sealed partial class ServiceEndpointWatcher(
     private ServiceEndpointSource? _cachedEndpoints;
     private Task _refreshTask = Task.CompletedTask;
     private volatile CacheStatus _cacheState;
+    private IDisposable? _changeTokenRegistration;
 
     /// <summary>
     /// Gets the service name.
@@ -60,6 +61,8 @@ public void Start()
     public ValueTask<ServiceEndpointSource> GetEndpointsAsync(CancellationToken cancellationToken = default)
     {
         ThrowIfNoProviders();
+        ObjectDisposedException.ThrowIf(_disposalCancellation.IsCancellationRequested, this);
+        cancellationToken.ThrowIfCancellationRequested();
 
         // If the cache is valid, return the cached value.
         if (_cachedEndpoints is { ChangeToken.HasChanged: false } cached)
@@ -76,9 +79,11 @@ async ValueTask<ServiceEndpointSource> GetEndpointsInternal(CancellationToken ca
             ServiceEndpointSource? result;
             do
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 await RefreshAsync(force: false).WaitAsync(cancellationToken).ConfigureAwait(false);
                 result = _cachedEndpoints;
             } while (result is null);
+
             return result;
         }
     }
@@ -89,7 +94,7 @@ private Task RefreshAsync(bool force)
         lock (_lock)
         {
             // If the cache is invalid or needs invalidation, refresh the cache.
-            if (_refreshTask.IsCompleted && (_cacheState == CacheStatus.Invalid || _cachedEndpoints is null or { ChangeToken.HasChanged: true } || force))
+            if (!_disposalCancellation.IsCancellationRequested && _refreshTask.IsCompleted && (_cacheState == CacheStatus.Invalid || _cachedEndpoints is null or { ChangeToken.HasChanged: true } || force))
             {
                 // Indicate that the cache is being updated and start a new refresh task.
                 _cacheState = CacheStatus.Refreshing;
@@ -128,10 +133,18 @@ private async Task RefreshAsyncInternal()
         CacheStatus newCacheState;
         try
         {
+            lock (_lock)
+            {
+                // Dispose the existing change token registration, if any.
+                _changeTokenRegistration?.Dispose();
+                _changeTokenRegistration = null;
+            }
+
             Log.ResolvingEndpoints(_logger, ServiceName);
             var builder = new ServiceEndpointBuilder();
             foreach (var provider in _providers)
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 await provider.PopulateAsync(builder, cancellationToken).ConfigureAwait(false);
             }
 
@@ -143,13 +156,12 @@ private async Task RefreshAsyncInternal()
                 // Check if we need to poll for updates or if we can register for change notification callbacks.
                 if (endpoints.ChangeToken.ActiveChangeCallbacks)
                 {
-                    // Initiate a background refresh, if necessary.
-                    endpoints.ChangeToken.RegisterChangeCallback(static state => _ = ((ServiceEndpointWatcher)state!).RefreshAsync(force: false), this);
-                    if (_pollingTimer is { } timer)
-                    {
-                        _pollingTimer = null;
-                        timer.Dispose();
-                    }
+                    // Initiate a background refresh when the change token fires.
+                    _changeTokenRegistration = endpoints.ChangeToken.RegisterChangeCallback(static state => _ = ((ServiceEndpointWatcher)state!).RefreshAsync(force: false), this);
+
+                    // Dispose the existing timer, if any, since we are reliant on change tokens for updates.
+                    _pollingTimer?.Dispose();
+                    _pollingTimer = null;
                 }
                 else
                 {
@@ -211,6 +223,13 @@ private void SchedulePollingTimer()
     {
         lock (_lock)
         {
+            if (_disposalCancellation.IsCancellationRequested)
+            {
+                _pollingTimer?.Dispose();
+                _pollingTimer = null;
+                return;
+            }
+
             if (_pollingTimer is null)
             {
                 _pollingTimer = _timeProvider.CreateTimer(s_pollingAction, this, _options.RefreshPeriod, TimeSpan.Zero);
@@ -227,14 +246,15 @@ public async ValueTask DisposeAsync()
     {
         lock (_lock)
         {
-            if (_pollingTimer is { } timer)
-            {
-                _pollingTimer = null;
-                timer.Dispose();
-            }
+            _disposalCancellation.Cancel();
+
+            _changeTokenRegistration?.Dispose();
+            _changeTokenRegistration = null;
+
+            _pollingTimer?.Dispose();
+            _pollingTimer = null;
         }
 
-        _disposalCancellation.Cancel();
         if (_refreshTask is { } task)
         {
             await task.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);

tests/Microsoft.Extensions.ServiceDiscovery.Dns.Tests/DnsServiceEndpointResolverTests.cs

@@ -0,0 +1,35 @@
+// 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.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Time.Testing;
+using Xunit;
+
+namespace Microsoft.Extensions.ServiceDiscovery.Dns.Tests;
+
+public class DnsServiceEndpointResolverTests
+{
+    [Fact]
+    public async Task ResolveServiceEndpoint_Dns_MultiShot()
+    {
+        var timeProvider = new FakeTimeProvider();
+        var services = new ServiceCollection()
+            .AddSingleton<TimeProvider>(timeProvider)
+            .AddServiceDiscoveryCore()
+            .AddDnsServiceEndpointProvider(o => o.DefaultRefreshPeriod = TimeSpan.FromSeconds(30))
+            .BuildServiceProvider();
+        var resolver = services.GetRequiredService<ServiceEndpointResolver>();
+        var initialResult = await resolver.GetEndpointsAsync("https://localhost", CancellationToken.None);
+        Assert.NotNull(initialResult);
+        Assert.True(initialResult.Endpoints.Count > 0);
+        timeProvider.Advance(TimeSpan.FromSeconds(7));
+        var secondResult = await resolver.GetEndpointsAsync("https://localhost", CancellationToken.None);
+        Assert.NotNull(secondResult);
+        Assert.True(initialResult.Endpoints.Count > 0);
+        timeProvider.Advance(TimeSpan.FromSeconds(80));
+        var thirdResult = await resolver.GetEndpointsAsync("https://localhost", CancellationToken.None);
+        Assert.NotNull(thirdResult);
+        Assert.True(initialResult.Endpoints.Count > 0);
+    }
+}

tests/Microsoft.Extensions.ServiceDiscovery.Dns.Tests/DnsSrvServiceEndpointResolverTests.cs

@@ -73,7 +73,7 @@ private sealed class FakeDnsQueryResponse : IDnsQueryResponse
     }
 
     [Fact]
-    public async Task ResolveServiceEndpoint_Dns()
+    public async Task ResolveServiceEndpoint_DnsSrv()
     {
         var dnsClientMock = new FakeDnsClient
         {
@@ -134,7 +134,7 @@ public async Task ResolveServiceEndpoint_Dns()
     [InlineData(true)]
     [InlineData(false)]
     [Theory]
-    public async Task ResolveServiceEndpoint_Dns_MultipleProviders_PreventMixing(bool dnsFirst)
+    public async Task ResolveServiceEndpoint_DnsSrv_MultipleProviders_PreventMixing(bool dnsFirst)
     {
         var dnsClientMock = new FakeDnsClient
         {
@@ -233,19 +233,4 @@ public async Task ResolveServiceEndpoint_Dns_MultipleProviders_PreventMixing(boo
             }
         }
     }
-
-    public class MyConfigurationProvider : ConfigurationProvider, IConfigurationSource
-    {
-        public IConfigurationProvider Build(IConfigurationBuilder builder) => this;
-        public void SetValues(IEnumerable<KeyValuePair<string, string?>> values)
-        {
-            Data.Clear();
-            foreach (var (key, value) in values)
-            {
-                Data[key] = value;
-            }
-
-            OnReload();
-        }
-    }
 }

tests/Microsoft.Extensions.ServiceDiscovery.Dns.Tests/Microsoft.Extensions.ServiceDiscovery.Dns.Tests.csproj

@@ -9,6 +9,7 @@
   <ItemGroup>
     <PackageReference Include="Microsoft.DotNet.XUnitExtensions" />
     <PackageReference Include="Microsoft.Extensions.Hosting" />
+    <PackageReference Include="Microsoft.Extensions.TimeProvider.Testing" />
   </ItemGroup>
 
   <ItemGroup>