Watchdog stop passed CancellationToken.None which could block stop till infinity

Author: ramonsmits

src/ServiceControl.Audit/Auditing/AuditIngestion.cs

@@ -260,7 +260,7 @@ public override async Task StopAsync(CancellationToken cancellationToken)
         {
             try
             {
-                await watchdog.Stop();
+                await watchdog.Stop(cancellationToken);
                 channel.Writer.Complete();
                 await base.StopAsync(cancellationToken);
             }

src/ServiceControl.Infrastructure.Tests/WatchdogTests.cs

@@ -31,7 +31,7 @@ public async Task It_shuts_down_gracefully()
 
             await started.Task;
 
-            await dog.Stop();
+            await dog.Stop(TestContext.CurrentContext.CancellationToken);
 
             await stopped.Task;
         }
@@ -52,7 +52,7 @@ public async Task When_stop_fails_it_reports_the_failure()
 
             await started.Task;
 
-            await dog.Stop();
+            await dog.Stop(TestContext.CurrentContext.CancellationToken);
 
             Assert.That(lastFailure, Is.EqualTo("Simulated"));
         }
@@ -91,7 +91,7 @@ public async Task On_failure_triggers_stopping()
 
             await restarted.Task;
 
-            await dog.Stop();
+            await dog.Stop(TestContext.CurrentContext.CancellationToken);
         }
 
         [Test]
@@ -129,7 +129,7 @@ public async Task When_first_start_attempt_works_it_recovers_from_further_errors
 
             await recoveredFromError.Task;
 
-            await dog.Stop();
+            await dog.Stop(TestContext.CurrentContext.CancellationToken);
 
             //Make sure failure is cleared
             Assert.That(lastFailure, Is.Null);

src/ServiceControl.Infrastructure/Watchdog.cs

@@ -17,9 +17,14 @@ public class Watchdog
         ILog log;
         string taskName;
 
-        public Watchdog(string taskName, Func<CancellationToken, Task> ensureStarted,
-            Func<CancellationToken, Task> ensureStopped, Action<string> reportFailure, Action clearFailure,
-            TimeSpan timeToWaitBetweenStartupAttempts, ILog log)
+        public Watchdog(
+            string taskName,
+            Func<CancellationToken, Task> ensureStarted,
+            Func<CancellationToken, Task> ensureStopped, Action<string> reportFailure,
+            Action clearFailure,
+            TimeSpan timeToWaitBetweenStartupAttempts,
+            ILog log
+        )
         {
             this.taskName = taskName;
             this.ensureStopped = ensureStopped;
@@ -87,25 +92,27 @@ public Task Start(Action onFailedOnStartup)
                         //Ignore, no need to log cancellation of delay
                     }
                 }
-                try
-                {
-                    log.Debug($"Stopping watching process {taskName}");
-                    //We don't pass the shutdown token here because it has already been cancelled and we want to ensure we stop the ingestion.
-                    await ensureStopped(CancellationToken.None).ConfigureAwait(false);
-                }
-                catch (Exception e)
-                {
-                    log.Error($"Error while trying to stop {taskName}.", e);
-                    reportFailure(e.Message);
-                }
             });
             return Task.CompletedTask;
         }
 
-        public Task Stop()
+        public async Task Stop(CancellationToken cancellationToken)
         {
-            shutdownTokenSource.Cancel();
-            return watchdog;
+            try
+            {
+                log.Debug($"Stopping watching process {taskName}");
+                await shutdownTokenSource.CancelAsync().ConfigureAwait(false);
+                await watchdog.ConfigureAwait(false);
+            }
+            catch (Exception e)
+            {
+                log.Error($"Error while trying to stop {taskName}.", e);
+                throw;
+            }
+            finally
+            {
+                await ensureStopped(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 }

src/ServiceControl/Operations/ErrorIngestion.cs

@@ -132,7 +132,7 @@ public override async Task StopAsync(CancellationToken cancellationToken)
         {
             try
             {
-                await watchdog.Stop();
+                await watchdog.Stop(cancellationToken);
                 channel.Writer.Complete();
                 await base.StopAsync(cancellationToken);
             }