Add token to timer stop

Author: andreasohlund

src/ServiceControl.Infrastructure/AsyncTimer.cs

@@ -57,7 +57,7 @@ public TimerJob(Func<CancellationToken, Task<TimerJobExecutionResult>> callback,
         }, CancellationToken.None);
     }
 
-    public async Task Stop()
+    public async Task Stop(CancellationToken cancellationToken)
     {
         if (tokenSource == null)
         {
@@ -67,16 +67,18 @@ public async Task Stop()
         await tokenSource.CancelAsync().ConfigureAwait(false);
         tokenSource.Dispose();
 
-        if (task != null)
+        if (task == null)
         {
-            try
-            {
-                await task.ConfigureAwait(false);
-            }
-            catch (OperationCanceledException) when (tokenSource.IsCancellationRequested)
-            {
-                //NOOP
-            }
+            return;
+        }
+
+        try
+        {
+            Task.WaitAll([task], cancellationToken);
+        }
+        catch (OperationCanceledException) when (tokenSource.IsCancellationRequested)
+        {
+            //NOOP
         }
     }
 

src/ServiceControl.Monitoring/Licensing/LicenseCheckHostedService.cs

@@ -20,7 +20,7 @@ public Task StartAsync(CancellationToken cancellationToken)
             return Task.CompletedTask;
         }
 
-        public Task StopAsync(CancellationToken cancellationToken) => timer.Stop();
+        public Task StopAsync(CancellationToken cancellationToken) => timer.Stop(cancellationToken);
 
         TimerJob timer;
 

src/ServiceControl/CustomChecks/InternalCustomChecks/InternalCustomCheckManager.cs

@@ -67,7 +67,7 @@ async Task<TimerJobExecutionResult> Run(CancellationToken cancellationToken)
                 : TimerJobExecutionResult.DoNotContinueExecuting;
         }
 
-        public Task Stop() => timer?.Stop() ?? Task.CompletedTask;
+        public Task Stop() => timer?.Stop(CancellationToken.None) ?? Task.CompletedTask;
 
         TimerJob timer;
         readonly ICustomCheck check;

src/ServiceControl/Licensing/LicenseCheckHostedService.cs

@@ -21,7 +21,7 @@ public Task StartAsync(CancellationToken cancellationToken)
             return Task.CompletedTask;
         }
 
-        public Task StopAsync(CancellationToken cancellationToken) => timer.Stop();
+        public Task StopAsync(CancellationToken cancellationToken) => timer.Stop(cancellationToken);
 
         TimerJob timer;
 

src/ServiceControl/Monitoring/HeartbeatMonitoringHostedService.cs

@@ -24,17 +24,7 @@ public async Task StartAsync(CancellationToken cancellationToken)
             timer = scheduler.Schedule(_ => CheckEndpoints(), TimeSpan.Zero, TimeSpan.FromSeconds(5), e => { log.Error("Exception occurred when monitoring endpoint instances", e); });
         }
 
-        public async Task StopAsync(CancellationToken cancellationToken)
-        {
-            try
-            {
-                await timer.Stop();
-            }
-            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
-            {
-                //NOOP, invoked Stop does not
-            }
-        }
+        public Task StopAsync(CancellationToken cancellationToken) => timer.Stop(cancellationToken);
 
         async Task<TimerJobExecutionResult> CheckEndpoints()
         {

src/ServiceControl/Recoverability/RecoverabilityComponent.cs

@@ -174,10 +174,7 @@ public Task StartAsync(CancellationToken cancellationToken)
                 return Task.CompletedTask;
             }
 
-            public Task StopAsync(CancellationToken cancellationToken)
-            {
-                return timer?.Stop() ?? Task.CompletedTask;
-            }
+            public Task StopAsync(CancellationToken cancellationToken) => timer?.Stop(cancellationToken) ?? Task.CompletedTask;
 
             async Task<TimerJobExecutionResult> ProcessRequestedBulkRetryOperations()
             {
@@ -237,10 +234,7 @@ public Task StartAsync(CancellationToken cancellationToken)
                 return Task.CompletedTask;
             }
 
-            public Task StopAsync(CancellationToken cancellationToken)
-            {
-                return timer.Stop();
-            }
+            public Task StopAsync(CancellationToken cancellationToken) => timer.Stop(cancellationToken);
 
             TimerJob timer;
             readonly IAsyncTimer scheduler;
@@ -266,10 +260,7 @@ public Task StartAsync(CancellationToken cancellationToken)
                 return Task.CompletedTask;
             }
 
-            public async Task StopAsync(CancellationToken cancellationToken)
-            {
-                await timer.Stop();
-            }
+            public Task StopAsync(CancellationToken cancellationToken) => timer.Stop(cancellationToken);
 
             async Task<TimerJobExecutionResult> Process(CancellationToken cancellationToken)
             {