Fix flaky HealthCheck test (#20741)

Author: BrennanConroy

src/HealthChecks/HealthChecks/src/HealthCheckPublisherHostedService.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
@@ -22,6 +22,7 @@ internal sealed class HealthCheckPublisherHostedService : IHostedService
 
         private CancellationTokenSource _stopping;
         private Timer _timer;
+        private CancellationTokenSource _runTokenSource;
 
         public HealthCheckPublisherHostedService(
             HealthCheckService healthCheckService,
@@ -69,7 +70,7 @@ public Task StartAsync(CancellationToken cancellationToken = default)
             }
 
             // IMPORTANT - make sure this is the last thing that happens in this method. The timer can
-            // fire before other code runs. 
+            // fire before other code runs.
             _timer = NonCapturingTimer.Create(Timer_Tick, null, dueTime: _options.Value.Delay, period: _options.Value.Period);
 
             return Task.CompletedTask;
@@ -104,6 +105,12 @@ private async void Timer_Tick(object state)
             await RunAsync();
         }
 
+        // Internal for testing
+        internal void CancelToken()
+        {
+            _runTokenSource.Cancel();
+        }
+
         // Internal for testing
         internal async Task RunAsync()
         {
@@ -116,6 +123,7 @@ internal async Task RunAsync()
                 var timeout = _options.Value.Timeout;
 
                 cancellation = CancellationTokenSource.CreateLinkedTokenSource(_stopping.Token);
+                _runTokenSource = cancellation;
                 cancellation.CancelAfter(timeout);
 
                 await RunAsyncCore(cancellation.Token);

src/HealthChecks/HealthChecks/test/HealthCheckPublisherHostedServiceTest.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
@@ -145,7 +145,7 @@ public async Task StopAsync_CancelsExecution()
                 await service.StopAsync(); // Trigger cancellation
 
                 // Assert
-                await AssertCancelledAsync(publishers[0].Entries[0].cancellationToken);
+                await AssertCanceledAsync(publishers[0].Entries[0].cancellationToken);
                 Assert.False(service.IsTimerRunning);
                 Assert.True(service.IsStopping);
 
@@ -286,10 +286,7 @@ public async Task RunAsync_PublishersCanTimeout()
                 new TestPublisher() { Wait = unblock.Task, },
             };
 
-            var service = CreateService(publishers, sink: sink, configure: (options) =>
-            {
-                options.Timeout = TimeSpan.FromMilliseconds(50);
-            });
+            var service = CreateService(publishers, sink: sink);
 
             try
             {
@@ -300,7 +297,9 @@ public async Task RunAsync_PublishersCanTimeout()
 
                 await publishers[0].Started.TimeoutAfter(TimeSpan.FromSeconds(10));
 
-                await AssertCancelledAsync(publishers[0].Entries[0].cancellationToken);
+                service.CancelToken();
+
+                await AssertCanceledAsync(publishers[0].Entries[0].cancellationToken);
 
                 unblock.SetResult(null);
 
@@ -483,7 +482,7 @@ private HealthCheckPublisherHostedService CreateService(
             return services.GetServices<IHostedService>().OfType< HealthCheckPublisherHostedService>().Single();
         }
 
-        private static async Task AssertCancelledAsync(CancellationToken cancellationToken)
+        private static async Task AssertCanceledAsync(CancellationToken cancellationToken)
         {
             await Assert.ThrowsAsync<TaskCanceledException>(() => Task.Delay(TimeSpan.FromSeconds(10), cancellationToken));
         }