From 8fb6c2a50acbf0118a39aeec6af0df898acc142f Mon Sep 17 00:00:00 2001
From: Ryan Nowak <nowakra@gmail.com>
Date: Wed, 19 Sep 2018 21:22:48 -0700
Subject: [PATCH] Allow cancellation to propagate

---
 .../DefaultHealthCheckService.cs              |  4 ++-
 .../DefaultHealthCheckServiceTest.cs          | 33 +++++++++++++++++--
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/Microsoft.Extensions.Diagnostics.HealthChecks/DefaultHealthCheckService.cs b/src/Microsoft.Extensions.Diagnostics.HealthChecks/DefaultHealthCheckService.cs
index b7a34ff35a..eae633dca0 100644
--- a/src/Microsoft.Extensions.Diagnostics.HealthChecks/DefaultHealthCheckService.cs
+++ b/src/Microsoft.Extensions.Diagnostics.HealthChecks/DefaultHealthCheckService.cs
@@ -77,7 +77,9 @@ namespace Microsoft.Extensions.Diagnostics.HealthChecks
 
                             Log.HealthCheckEnd(_logger, registration, entry, stopwatch.GetElapsedTime());
                         }
-                        catch (Exception ex)
+
+                        // Allow cancellation to propagate.
+                        catch (Exception ex) when (ex as OperationCanceledException == null)
                         {
                             entry = new HealthReportEntry(HealthStatus.Failed, ex.Message, ex, data: null);
                             Log.HealthCheckError(_logger, registration, ex, stopwatch.GetElapsedTime());
diff --git a/test/Microsoft.Extensions.Diagnostics.HealthChecks.Tests/DefaultHealthCheckServiceTest.cs b/test/Microsoft.Extensions.Diagnostics.HealthChecks.Tests/DefaultHealthCheckServiceTest.cs
index 024309a024..6c74ddcdf0 100644
--- a/test/Microsoft.Extensions.Diagnostics.HealthChecks.Tests/DefaultHealthCheckServiceTest.cs
+++ b/test/Microsoft.Extensions.Diagnostics.HealthChecks.Tests/DefaultHealthCheckServiceTest.cs
@@ -33,7 +33,7 @@ namespace Microsoft.Extensions.Diagnostics.HealthChecks
                 .AddCheck("Baz", new DelegateHealthCheck(_ => Task.FromResult(HealthCheckResult.Passed())));
 
             var services = serviceCollection.BuildServiceProvider();
-            
+
             var scopeFactory = services.GetRequiredService<IServiceScopeFactory>();
             var options = services.GetRequiredService<IOptions<HealthCheckServiceOptions>>();
             var logger = services.GetRequiredService<ILogger<DefaultHealthCheckService>>();
@@ -188,6 +188,35 @@ namespace Microsoft.Extensions.Diagnostics.HealthChecks
                 });
         }
 
+        [Fact]
+        public async Task CheckHealthAsync_Cancellation_CanPropagate()
+        {
+            // Arrange
+            var insideCheck = new TaskCompletionSource<object>();
+
+            var service = CreateHealthChecksService(b =>
+            {
+                b.AddAsyncCheck("cancels", async ct =>
+                {
+                    insideCheck.SetResult(null);
+
+                    await Task.Delay(10000, ct);
+                    return HealthCheckResult.Failed();
+                });
+            });
+
+            var cancel = new CancellationTokenSource();
+            var task = service.CheckHealthAsync(cancel.Token);
+
+            // After this returns we know the check has started
+            await insideCheck.Task;
+
+            cancel.Cancel();
+
+            // Act & Assert
+            await Assert.ThrowsAsync<TaskCanceledException>(async () => await task);
+        }
+
         [Fact]
         public async Task CheckHealthAsync_ConvertsExceptionInHealthCheckToFailedResultAsync()
         {
@@ -366,7 +395,7 @@ namespace Microsoft.Extensions.Diagnostics.HealthChecks
             public CheckWithServiceDependency(AnotherService _)
             {
             }
-            
+
             public Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
             {
                 return Task.FromResult(HealthCheckResult.Passed());