Don't end resource notifications until service is disposed (#7108)

Contributes to #4878

Author: DamianEdwards

src/Aspire.Hosting/ApplicationModel/ResourceNotificationService.cs

@@ -14,13 +14,13 @@ namespace Aspire.Hosting.ApplicationModel;
 /// <summary>
 /// A service that allows publishing and subscribing to changes in the state of a resource.
 /// </summary>
-public class ResourceNotificationService
+public class ResourceNotificationService : IDisposable
 {
     // Resource state is keyed by the resource and the unique name of the resource. This could be the name of the resource, or a replica ID.
     private readonly ConcurrentDictionary<(IResource, string), ResourceNotificationState> _resourceNotificationStates = new();
     private readonly ILogger<ResourceNotificationService> _logger;
     private readonly IServiceProvider _serviceProvider;
-    private readonly CancellationToken _applicationStopping;
+    private readonly CancellationTokenSource _disposing = new();
     private readonly ResourceLoggerService _resourceLoggerService;
 
     private Action<ResourceEvent>? OnResourceUpdated { get; set; }
@@ -43,7 +43,6 @@ public ResourceNotificationService(ILogger<ResourceNotificationService> logger,
     {
         _logger = logger ?? throw new ArgumentNullException(nameof(logger));
         _serviceProvider = new NullServiceProvider();
-        _applicationStopping = hostApplicationLifetime?.ApplicationStopping ?? throw new ArgumentNullException(nameof(hostApplicationLifetime));
         _resourceLoggerService = new ResourceLoggerService();
     }
 
@@ -58,8 +57,10 @@ public ResourceNotificationService(ILogger<ResourceNotificationService> logger,
     {
         _logger = logger ?? throw new ArgumentNullException(nameof(logger));
         _serviceProvider = serviceProvider;
-        _applicationStopping = hostApplicationLifetime?.ApplicationStopping ?? throw new ArgumentNullException(nameof(hostApplicationLifetime));
         _resourceLoggerService = resourceLoggerService ?? throw new ArgumentNullException(nameof(resourceLoggerService));
+
+        // The IHostApplicationLifetime parameter is not used anymore, but we keep it for backwards compatibility.
+        // Notfication updates will be cancelled when the service is disposed.
     }
 
     private class NullServiceProvider : IServiceProvider
@@ -105,7 +106,7 @@ public Task WaitForResourceAsync(string resourceName, string? targetState = null
                                                      Justification = "targetState(s) parameters are mutually exclusive.")]
     public async Task<string> WaitForResourceAsync(string resourceName, IEnumerable<string> targetStates, CancellationToken cancellationToken = default)
     {
-        using var watchCts = CancellationTokenSource.CreateLinkedTokenSource(_applicationStopping, cancellationToken);
+        using var watchCts = CancellationTokenSource.CreateLinkedTokenSource(_disposing.Token, cancellationToken);
         var watchToken = watchCts.Token;
         await foreach (var resourceEvent in WatchAsync(watchToken).ConfigureAwait(false))
         {
@@ -273,7 +274,7 @@ public async Task WaitForDependenciesAsync(IResource resource, CancellationToken
                                                      Justification = "predicate and targetState(s) parameters are mutually exclusive.")]
     public async Task<ResourceEvent> WaitForResourceAsync(string resourceName, Func<ResourceEvent, bool> predicate, CancellationToken cancellationToken = default)
     {
-        using var watchCts = CancellationTokenSource.CreateLinkedTokenSource(_applicationStopping, cancellationToken);
+        using var watchCts = CancellationTokenSource.CreateLinkedTokenSource(_disposing.Token, cancellationToken);
         var watchToken = watchCts.Token;
         await foreach (var resourceEvent in WatchAsync(watchToken).ConfigureAwait(false))
         {
@@ -502,6 +503,12 @@ private static CustomResourceSnapshot GetCurrentSnapshot(IResource resource, Res
     private ResourceNotificationState GetResourceNotificationState(IResource resource, string resourceId) =>
         _resourceNotificationStates.GetOrAdd((resource, resourceId), _ => new ResourceNotificationState());
 
+    /// <inheritdoc/>
+    public void Dispose()
+    {
+        _disposing.Cancel();
+    }
+
     /// <summary>
     /// The annotation that allows publishing and subscribing to changes in the state of a resource.
     /// </summary>

src/Aspire.Hosting/PublicAPI.Unshipped.txt

@@ -133,6 +133,7 @@ Aspire.Hosting.ApplicationModel.ResourceCommandState
 Aspire.Hosting.ApplicationModel.ResourceCommandState.Disabled = 1 -> Aspire.Hosting.ApplicationModel.ResourceCommandState
 Aspire.Hosting.ApplicationModel.ResourceCommandState.Enabled = 0 -> Aspire.Hosting.ApplicationModel.ResourceCommandState
 Aspire.Hosting.ApplicationModel.ResourceCommandState.Hidden = 2 -> Aspire.Hosting.ApplicationModel.ResourceCommandState
+Aspire.Hosting.ApplicationModel.ResourceNotificationService.Dispose() -> void
 Aspire.Hosting.ApplicationModel.ResourceNotificationService.ResourceNotificationService(Microsoft.Extensions.Logging.ILogger<Aspire.Hosting.ApplicationModel.ResourceNotificationService!>! logger, Microsoft.Extensions.Hosting.IHostApplicationLifetime! hostApplicationLifetime) -> void
 Aspire.Hosting.ApplicationModel.ResourceNotificationService.ResourceNotificationService(Microsoft.Extensions.Logging.ILogger<Aspire.Hosting.ApplicationModel.ResourceNotificationService!>! logger, Microsoft.Extensions.Hosting.IHostApplicationLifetime! hostApplicationLifetime, System.IServiceProvider! serviceProvider, Aspire.Hosting.ApplicationModel.ResourceLoggerService! resourceLoggerService) -> void
 Aspire.Hosting.ApplicationModel.ResourceNotificationService.WaitForDependenciesAsync(Aspire.Hosting.ApplicationModel.IResource! resource, System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task!

tests/Aspire.Hosting.Tests/ResourceNotificationTests.cs

@@ -244,6 +244,24 @@ public async Task WaitingOnResourceReturnsCorrectStateWhenResourceReachesOneOfTa
         Assert.Equal("SomeOtherState", reachedState);
     }
 
+    [Fact]
+    public async Task WaitingOnResourceReturnsItReachesStateAfterApplicationStoppingCancellationTokenSignaled()
+    {
+        var resource1 = new CustomResource("myResource1");
+
+        using var hostApplicationLifetime = new TestHostApplicationLifetime();
+        var notificationService = ResourceNotificationServiceTestHelpers.Create(hostApplicationLifetime: hostApplicationLifetime);
+
+        var waitTask = notificationService.WaitForResourceAsync("myResource1", "SomeState");
+        hostApplicationLifetime.StopApplication();
+
+        await notificationService.PublishUpdateAsync(resource1, snapshot => snapshot with { State = "SomeState" }).DefaultTimeout();
+
+        await waitTask.DefaultTimeout();
+
+        Assert.True(waitTask.IsCompletedSuccessfully);
+    }
+
     [Fact]
     public async Task WaitingOnResourceThrowsOperationCanceledExceptionIfResourceDoesntReachStateBeforeCancellationTokenSignaled()
     {
@@ -261,13 +279,13 @@ await Assert.ThrowsAsync<OperationCanceledException>(async () =>
     }
 
     [Fact]
-    public async Task WaitingOnResourceThrowsOperationCanceledExceptionIfResourceDoesntReachStateBeforeApplicationStoppingCancellationTokenSignaled()
+    public async Task WaitingOnResourceThrowsOperationCanceledExceptionIfResourceDoesntReachStateBeforeServiceIsDisposed()
     {
-        using var hostApplicationLifetime = new TestHostApplicationLifetime();
-        var notificationService = ResourceNotificationServiceTestHelpers.Create(hostApplicationLifetime: hostApplicationLifetime);
+        var notificationService = ResourceNotificationServiceTestHelpers.Create();
 
         var waitTask = notificationService.WaitForResourceAsync("myResource1", "SomeState");
-        hostApplicationLifetime.StopApplication();
+
+        notificationService.Dispose();
 
         await Assert.ThrowsAsync<OperationCanceledException>(async () =>
         {