Fix leaking transportInfrastructure instance

Author: ramonsmits

src/ServiceControl.Audit/Auditing/AuditIngestion.cs

@@ -78,65 +78,97 @@ async Task EnsureStarted(CancellationToken cancellationToken = default)
 
                 if (!unitOfWorkFactory.CanIngestMore())
                 {
-                    if (queueIngestor != null)
-                    {
-                        var stoppable = queueIngestor;
-                        queueIngestor = null;
-                        logger.Info("Shutting down due to failed persistence health check. Infrastructure shut down commencing");
-                        await stoppable.StopReceive(cancellationToken);
-                        logger.Info("Shutting down due to failed persistence health check. Infrastructure shut down completed");
-                    }
-
+                    await StopAndTeardownInfrastructure(cancellationToken);
                     return;
                 }
 
-                if (queueIngestor != null)
+                await SetUpAndStartInfrastructure(cancellationToken);
+            }
+            catch (Exception e)
+            {
+                try
                 {
-                    logger.Debug("Ensure started. Already started, skipping start up");
-                    return; //Already started
+                    await StopAndTeardownInfrastructure(cancellationToken);
                 }
+                catch (Exception teardownException)
+                {
+                    throw new AggregateException(e, teardownException);
+                }
+
+                throw;
+            }
+            finally
+            {
+                logger.Debug("Ensure started. Start/stop semaphore releasing");
+                startStopSemaphore.Release();
+                logger.Debug("Ensure started. Start/stop semaphore released");
+            }
+        }
 
-                logger.Info("Ensure started. Infrastructure starting");
+        async Task SetUpAndStartInfrastructure(CancellationToken cancellationToken)
+        {
+            if (queueIngestor != null)
+            {
+                logger.Debug("Infrastructure already Started");
+                return;
+            }
 
+            try
+            {
+                logger.Info("Starting infrastructure");
                 transportInfrastructure = await transportCustomization.CreateTransportInfrastructure(
                     inputEndpoint,
                     transportSettings,
                     OnMessage,
                     errorHandlingPolicy.OnError,
                     OnCriticalError,
-                    TransportTransactionMode.ReceiveOnly);
+                    TransportTransactionMode.ReceiveOnly
+                );
 
                 queueIngestor = transportInfrastructure.Receivers[inputEndpoint];
 
                 await auditIngestor.VerifyCanReachForwardingAddress();
-
                 await queueIngestor.StartReceive(cancellationToken);
 
-                logger.Info("Ensure started. Infrastructure started");
+                logger.Info("Started infrastructure");
             }
-            catch
+            catch (Exception e)
             {
-                if (queueIngestor != null)
+                logger.Error("Failed to start infrastructure", e);
+                throw;
+            }
+        }
+
+        async Task StopAndTeardownInfrastructure(CancellationToken cancellationToken)
+        {
+            if (transportInfrastructure == null)
+            {
+                logger.Debug("Infrastructure already Stopped");
+                return;
+            }
+
+            try
+            {
+                logger.Info("Stopping infrastructure");
+                try
                 {
-                    try
+                    if (queueIngestor != null)
                     {
                         await queueIngestor.StopReceive(cancellationToken);
                     }
-                    catch (OperationCanceledException e) when (e.CancellationToken == cancellationToken)
-                    {
-                        logger.Info("StopReceive cancelled");
-                    }
+                }
+                finally
+                {
+                    await transportInfrastructure.Shutdown(cancellationToken);
                 }
 
-                queueIngestor = null; // Setting to null so that it doesn't exit when it retries in line 185
-
-                throw;
+                queueIngestor = null;
+                logger.Info("Stopped infrastructure");
             }
-            finally
+            catch (Exception e)
             {
-                logger.Debug("Ensure started. Start/stop semaphore releasing");
-                startStopSemaphore.Release();
-                logger.Debug("Ensure started. Start/stop semaphore released");
+                logger.Error("Failed to stop infrastructure", e);
+                throw;
             }
         }
 

src/ServiceControl.Infrastructure/Watchdog.cs

@@ -12,7 +12,7 @@ public class Watchdog
         Action<string> reportFailure;
         Action clearFailure;
         Task watchdog;
-        CancellationTokenSource shutdownTokenSource = new CancellationTokenSource();
+        CancellationTokenSource shutdownTokenSource = new();
         TimeSpan timeToWaitBetweenStartupAttempts;
         ILog log;
         string taskName;
@@ -47,41 +47,37 @@ public Task Start(Action onFailedOnStartup)
             {
                 log.Debug($"Starting watching {taskName}");
 
-                bool? failedOnStartup = null;
+                bool startup = true;
 
                 while (!shutdownTokenSource.IsCancellationRequested)
                 {
                     try
                     {
+                        using var cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(shutdownTokenSource.Token);
+                        cancellationTokenSource.CancelAfter(15000);
+
                         log.Debug($"Ensuring {taskName} is running");
-                        await ensureStarted(shutdownTokenSource.Token).ConfigureAwait(false);
+                        await ensureStarted(cancellationTokenSource.Token).ConfigureAwait(false);
                         clearFailure();
-
-                        failedOnStartup ??= false;
+                        startup = false;
                     }
-                    catch (OperationCanceledException e) when (!shutdownTokenSource.IsCancellationRequested)
+                    catch (OperationCanceledException e) when (shutdownTokenSource.IsCancellationRequested)
                     {
-                        // Continue, as OCE is not from caller
-                        log.Info("Start cancelled, retrying...", e);
-                        continue;
+                        log.Debug("Cancelled", e);
+                        return;
                     }
                     catch (Exception e)
                     {
                         reportFailure(e.Message);
 
-                        if (failedOnStartup == null)
+                        if (startup)
                         {
-                            failedOnStartup = true;
-
                             log.Error($"Error during initial startup attempt for {taskName}.", e);
-
-                            //there was an error during startup hence we want to shut down the instance
                             onFailedOnStartup();
+                            return;
                         }
-                        else
-                        {
-                            log.Error($"Error while trying to start {taskName}. Starting will be retried in {timeToWaitBetweenStartupAttempts}.", e);
-                        }
+
+                        log.Error($"Error while trying to start {taskName}. Starting will be retried in {timeToWaitBetweenStartupAttempts}.", e);
                     }
                     try
                     {