Introduce BackgroundService utility class in StartupHook

src/Microsoft.Diagnostics.Monitoring.HostingStartup/ParameterCapturing/ParameterCapturingService.cs

@@ -9,13 +9,13 @@
 using Microsoft.Diagnostics.Tools.Monitor.ParameterCapturing;
 using Microsoft.Diagnostics.Tools.Monitor.Profiler;
 using Microsoft.Diagnostics.Tools.Monitor.StartupHook;
-using Microsoft.Extensions.Hosting;
 using System;
 using System.Collections.Generic;
 using System.Globalization;
 using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
+using BackgroundService = Microsoft.Extensions.Hosting.BackgroundService;
 
 namespace Microsoft.Diagnostics.Monitoring.HostingStartup.ParameterCapturing
 {

src/Microsoft.Diagnostics.Monitoring.StartupHook/BackgroundService.cs

@@ -0,0 +1,55 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Diagnostics.Monitoring.StartupHook
+{
+    internal abstract class BackgroundService : IDisposable
+    {
+        private readonly CancellationTokenSource _cts = new();
+        private Task? _executeTask;
+        private long _disposedState;
+
+        public void Start()
+        {
+            _executeTask = Task.Run(async () =>
+            {
+                try
+                {
+                    await ExecuteAsync(_cts.Token);
+                }
+                catch (OperationCanceledException)
+                {
+                    // Ignore
+                }
+                catch (Exception ex)
+                {
+                    BackgroundTaskException = ex;
+                }
+            });
+        }
+
+        public void Stop()
+        {
+            _cts.Cancel();
+        }
+
+        public Exception? BackgroundTaskException { get; private set; }
+
+        public virtual void Dispose()
+        {
+            if (!DisposableHelper.CanDispose(ref _disposedState))
+                return;
+
+            _executeTask?.Wait();
+            _executeTask = null;
+
+            _cts.Dispose();
+        }
+
+        protected abstract Task ExecuteAsync(CancellationToken stoppingToken);
+    }
+}

src/Tests/Microsoft.Diagnostics.Monitoring.StartupHook.UnitTests/BackgroundService/BackgroundServiceTests.cs

@@ -0,0 +1,124 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.Diagnostics.Monitoring.TestCommon;
+using System;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Microsoft.Diagnostics.Monitoring.StartupHook
+{
+    [TargetFrameworkMonikerTrait(TargetFrameworkMonikerExtensions.CurrentTargetFrameworkMoniker)]
+    public sealed class BackgroundServiceTests
+    {
+        [Fact]
+        public void ConstructionWorks()
+        {
+            using BackgroundService _ = new MockBackgroundService();
+        }
+
+        [Fact]
+        public async Task StartBackgroundTask()
+        {
+            // Arrange
+            using MockBackgroundService service = new MockBackgroundService();
+
+            // Act
+            service.Start();
+
+            // Assert
+            await service.BackgroundTaskStarted.Task;
+        }
+
+        [Fact]
+        public async Task StopTriggersCancellation()
+        {
+            // Arrange
+            TaskCompletionSource backgroundTaskCompletion = new();
+            using MockBackgroundService service = new MockBackgroundService(backgroundTaskCompletion.Task);
+
+            // Act
+            service.Start();
+            await service.BackgroundTaskStarted.Task;
+
+            // Assert
+            Assert.False(service.BackgroundTaskWasCancelled);
+
+            service.Stop();
+            backgroundTaskCompletion.SetResult();
+            await service.BackgroundTaskEnded.Task;
+
+            Assert.True(service.BackgroundTaskWasCancelled);
+            Assert.Null(service.BackgroundTaskException);
+        }
+
+        [Fact]
+        public async Task DisposeWaitsForTheBackgroundTask()
+        {
+            // Arrange
+            object lockObj = new();
+            int callOrderMarker = 1;
+            TaskCompletionSource backgroundTaskCompletion = new();
+
+            // If Dispose() completes first, callOrderMarker will be 10
+            // Otherwise, callOrderMarker will be 20
+            void OnDisposeCompleted()
+            {
+                lock (lockObj)
+                {
+                    callOrderMarker *= 10;
+                }
+            }
+
+            async Task BackgroundWork()
+            {
+                await backgroundTaskCompletion.Task;
+                lock (lockObj)
+                {
+                    callOrderMarker += 1;
+                }
+            }
+
+            MockBackgroundService service = new MockBackgroundService(BackgroundWork(), OnDisposeCompleted);
+
+            // Act
+            service.Start();
+            await service.BackgroundTaskStarted.Task;
+            service.Stop();
+
+            Task disposeTask = Task.Run(service.Dispose);
+
+            await service.DisposeStarted.Task;
+            await Task.Delay(100); // Ensure that Dispose is waiting for the background task to complete
+            backgroundTaskCompletion.SetResult();
+
+            await disposeTask;
+
+            // Assert
+            Assert.Null(service.BackgroundTaskException);
+            Assert.Equal(20, callOrderMarker);
+        }
+
+        [Fact]
+        public async Task BackgroundTaskExceptionIsCaptured()
+        {
+            // Arrange
+            static async Task BackgroundWork()
+            {
+                await Task.Yield();
+                throw new NotImplementedException();
+            }
+            MockBackgroundService service = new MockBackgroundService(BackgroundWork());
+
+            // Act
+            service.Start();
+            await service.BackgroundTaskStarted.Task;
+
+            service.Stop();
+            service.Dispose();
+
+            // Assert
+            Assert.IsType<NotImplementedException>(service.BackgroundTaskException);
+        }
+    }
+}

src/Tests/Microsoft.Diagnostics.Monitoring.StartupHook.UnitTests/BackgroundService/MockBackgroundService.cs

@@ -0,0 +1,64 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Diagnostics.Monitoring.StartupHook
+{
+    internal sealed class MockBackgroundService : BackgroundService, IDisposable
+    {
+        private readonly Task _backgroundTask;
+        private readonly Action _postDisposeAction;
+
+        public MockBackgroundService()
+        {
+            _backgroundTask = Task.CompletedTask;
+            _postDisposeAction = () => { };
+        }
+
+        public MockBackgroundService(Task backgroundTaskInput)
+        {
+            _backgroundTask = backgroundTaskInput;
+            _postDisposeAction = () => { };
+        }
+
+        public MockBackgroundService(Task backgroundTaskInput, Action postDisposeAction)
+        {
+            _backgroundTask = backgroundTaskInput;
+            _postDisposeAction = postDisposeAction;
+        }
+
+        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
+        {
+            BackgroundTaskStarted.SetResult();
+
+            await _backgroundTask;
+
+            if (stoppingToken.IsCancellationRequested)
+            {
+                BackgroundTaskWasCancelled = true;
+            }
+
+            BackgroundTaskEnded.SetResult();
+        }
+
+        public override void Dispose()
+        {
+            DisposeStarted.SetResult();
+
+            base.Dispose();
+
+            _postDisposeAction();
+        }
+
+        public TaskCompletionSource BackgroundTaskStarted { get; } = new();
+
+        public TaskCompletionSource BackgroundTaskEnded { get; } = new();
+
+        public TaskCompletionSource DisposeStarted { get; } = new();
+
+        public bool BackgroundTaskWasCancelled { get; private set; }
+    }
+}