Overriding Start and Stop to support graceful shutdown and improve intent and readability

Author: ramonsmits

src/ServiceControl.Audit/Auditing/AuditIngestion.cs

@@ -194,11 +194,14 @@ async Task OnMessage(MessageContext messageContext, CancellationToken cancellati
             await taskCompletionSource.Task;
         }
 
-
-        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
+        public override async Task StartAsync(CancellationToken cancellationToken)
         {
             await watchdog.Start(() => applicationLifetime.StopApplication());
+            await base.StartAsync(cancellationToken);
+        }
 
+        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
+        {
             try
             {
                 var contexts = new List<MessageContext>(transportSettings.MaxConcurrency.Value);
@@ -243,18 +246,30 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
             }
             catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
             {
-                // Cancellation
+                // ExecuteAsync cancelled
             }
-            finally
+        }
+
+        public override async Task StopAsync(CancellationToken cancellationToken)
+        {
+            try
             {
                 await watchdog.Stop();
                 channel.Writer.Complete();
-
+                await base.StopAsync(cancellationToken);
+            }
+            finally
+            {
                 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);
+                    try
+                    {
+                        await transportInfrastructure.Shutdown(cancellationToken);
+                    }
+                    catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+                    {
+                        // StopAsync cancelled
+                    }
                 }
             }
         }

src/ServiceControl/Operations/ErrorIngestion.cs

@@ -70,10 +70,14 @@ public ErrorIngestion(
             );
         }
 
-        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
+        public override async Task StartAsync(CancellationToken cancellationToken)
         {
             await watchdog.Start(() => applicationLifetime.StopApplication());
+            await base.StartAsync(cancellationToken);
+        }
 
+        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
+        {
             try
             {
                 var contexts = new List<MessageContext>(transportSettings.MaxConcurrency.Value);
@@ -118,17 +122,30 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
             }
             catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
             {
-                // Cancellation
+                // ExecuteAsync cancelled
             }
-            finally
+        }
+
+        public override async Task StopAsync(CancellationToken cancellationToken)
+        {
+            try
             {
                 await watchdog.Stop();
                 channel.Writer.Complete();
+                await base.StopAsync(cancellationToken);
+            }
+            finally
+            {
                 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);
+                    try
+                    {
+                        await transportInfrastructure.Shutdown(cancellationToken);
+                    }
+                    catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+                    {
+                        // StopAsync cancelled
+                    }
                 }
             }
         }