Refactored to use BackgroundService

Author: ramonsmits

src/ServiceControl.Audit/Auditing/AuditIngestion.cs

@@ -17,7 +17,7 @@
     using ServiceControl.Infrastructure.Metrics;
     using Transports;
 
-    class AuditIngestion : IHostedService
+    class AuditIngestion : BackgroundService
     {
         static readonly long FrequencyInMilliseconds = Stopwatch.Frequency / 1000;
 
@@ -70,26 +70,6 @@ public AuditIngestion(
             );
         }
 
-        public async Task StartAsync(CancellationToken cancellationToken)
-        {
-            stopSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-            ingestionWorker = Task.Run(() => Loop(stopSource.Token), stopSource.Token);
-            await watchdog.Start(() => applicationLifetime.StopApplication());
-        }
-
-        public async Task StopAsync(CancellationToken cancellationToken)
-        {
-            await stopSource.CancelAsync();
-            await watchdog.Stop();
-            channel.Writer.Complete();
-            await ingestionWorker;
-
-            if (transportInfrastructure != null)
-            {
-                await transportInfrastructure.Shutdown(cancellationToken);
-            }
-        }
-
         Task OnCriticalError(string failure, Exception exception)
         {
             logger.Fatal($"OnCriticalError. '{failure}'", exception);
@@ -214,13 +194,16 @@ async Task OnMessage(MessageContext messageContext, CancellationToken cancellati
             await taskCompletionSource.Task;
         }
 
-        async Task Loop(CancellationToken cancellationToken)
+
+        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
         {
+            await watchdog.Start(() => applicationLifetime.StopApplication());
+
             try
             {
                 var contexts = new List<MessageContext>(transportSettings.MaxConcurrency.Value);
 
-                while (await channel.Reader.WaitToReadAsync(cancellationToken))
+                while (await channel.Reader.WaitToReadAsync(stoppingToken))
                 {
                     // will only enter here if there is something to read.
                     try
@@ -237,7 +220,7 @@ async Task Loop(CancellationToken cancellationToken)
                             await auditIngestor.Ingest(contexts);
                         }
                     }
-                    catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+                    catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
                     {
                         throw; // Catch again in outer catch
                     }
@@ -258,16 +241,26 @@ async Task Loop(CancellationToken cancellationToken)
                 }
                 // will fall out here when writer is completed
             }
-            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+            catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
             {
                 // Cancellation
             }
+            finally
+            {
+                await watchdog.Stop();
+                channel.Writer.Complete();
+
+                if (transportInfrastructure != null)
+                {
+                    // stoppingToken is cancelled, invoke so transport infrastructure will run teardown
+                    // No need to await, as this will throw OperationCancelledException
+                    _ = transportInfrastructure.Shutdown(stoppingToken);
+                }
+            }
         }
 
         TransportInfrastructure transportInfrastructure;
         IMessageReceiver queueIngestor;
-        Task ingestionWorker;
-        CancellationTokenSource stopSource;
 
         readonly SemaphoreSlim startStopSemaphore = new SemaphoreSlim(1);
         readonly string inputEndpoint;

src/ServiceControl/Operations/ErrorIngestion.cs

@@ -17,7 +17,7 @@
     using ServiceBus.Management.Infrastructure.Settings;
     using Transports;
 
-    class ErrorIngestion : IHostedService
+    class ErrorIngestion : BackgroundService
     {
         static readonly long FrequencyInMilliseconds = Stopwatch.Frequency / 1000;
 
@@ -70,31 +70,15 @@ public ErrorIngestion(
             );
         }
 
-        public async Task StartAsync(CancellationToken cancellationToken)
+        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
         {
-            stopSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-            ingestionWorker = Task.Run(() => Loop(stopSource.Token), stopSource.Token);
             await watchdog.Start(() => applicationLifetime.StopApplication());
-        }
-
-        public async Task StopAsync(CancellationToken cancellationToken)
-        {
-            await watchdog.Stop();
-            channel.Writer.Complete();
-            await ingestionWorker;
-            if (transportInfrastructure != null)
-            {
-                await transportInfrastructure.Shutdown(cancellationToken);
-            }
-        }
 
-        async Task Loop(CancellationToken cancellationToken)
-        {
             try
             {
                 var contexts = new List<MessageContext>(transportSettings.MaxConcurrency.Value);
 
-                while (await channel.Reader.WaitToReadAsync(cancellationToken))
+                while (await channel.Reader.WaitToReadAsync(stoppingToken))
                 {
                     // will only enter here if there is something to read.
                     try
@@ -111,7 +95,7 @@ async Task Loop(CancellationToken cancellationToken)
                             await ingestor.Ingest(contexts);
                         }
                     }
-                    catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+                    catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
                     {
                         throw; // Catch again in outer catch
                     }
@@ -132,10 +116,21 @@ async Task Loop(CancellationToken cancellationToken)
                 }
                 // will fall out here when writer is completed
             }
-            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+            catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
             {
                 // Cancellation
             }
+            finally
+            {
+                await watchdog.Stop();
+                channel.Writer.Complete();
+                if (transportInfrastructure != null)
+                {
+                    // stoppingToken is cancelled, invoke so transport infrastructure will run teardown
+                    // No need to await, as this will throw OperationCancelledException
+                    _ = transportInfrastructure.Shutdown(stoppingToken);
+                }
+            }
         }
 
         async Task EnsureStarted(CancellationToken cancellationToken = default)
@@ -244,8 +239,6 @@ async Task EnsureStopped(CancellationToken cancellationToken = default)
         ErrorIngestionFaultPolicy errorHandlingPolicy;
         TransportInfrastructure transportInfrastructure;
         IMessageReceiver messageReceiver;
-        Task ingestionWorker;
-        CancellationTokenSource stopSource;
 
         readonly Settings settings;
         readonly ITransportCustomization transportCustomization;