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,62 @@
+// 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 long _disposedState;
+
+        public Task? ExecutingTask { get; private set; }
+
+        public void Start()
+        {
+            ExecutingTask = Task.Run(async () =>
+            {
+                await ExecuteAsync(_cts.Token).ConfigureAwait(false);
+            }, _cts.Token);
+        }
+
+        public void Stop()
+        {
+            SafeCancel();
+
+            try
+            {
+                ExecutingTask?.Wait(TimeSpan.FromSeconds(1));
+            }
+            catch
+            {
+                // ignore
+            }
+        }
+
+        public virtual void Dispose()
+        {
+            if (!DisposableHelper.CanDispose(ref _disposedState))
+                return;
+
+            SafeCancel();
+            _cts.Dispose();
+        }
+
+        private void SafeCancel()
+        {
+            try
+            {
+                _cts.Cancel();
+            }
+            catch (AggregateException)
+            {
+                //  Ignore all exceptions thrown by registered callbacks on the associated CancellationToken.
+            }
+        }
+
+        protected abstract Task ExecuteAsync(CancellationToken stoppingToken);
+    }
+}

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

@@ -0,0 +1,129 @@
+// 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;
+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 Start_RunsBackgroundTask()
+        {
+            // Arrange
+            using CancellationTokenSource cts = new(CommonTestTimeouts.GeneralTimeout);
+            using MockBackgroundService service = new MockBackgroundService();
+
+            // Act
+            service.Start();
+
+            // Assert
+            await service.BackgroundTaskStarted.Task.WaitAsync(cts.Token);
+        }
+
+        [Fact]
+        public async Task Stop_TriggersCancellation()
+        {
+            // Arrange
+            using CancellationTokenSource cts = new(CommonTestTimeouts.GeneralTimeout);
+            using MockBackgroundService service = new MockBackgroundService(async (CancellationToken stoppingToken) =>
+            {
+                await Task.Delay(Timeout.Infinite, stoppingToken);
+            });
+
+            // Act
+            service.Start();
+            await service.BackgroundTaskStarted.Task.WaitAsync(cts.Token);
+            service.Stop();
+
+            // Assert
+            Assert.NotNull(service.ExecutingTask);
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(() => service.ExecutingTask);
+        }
+
+        [Fact]
+        public async Task Stop_WaitsForTheBackgroundTask()
+        {
+            // Arrange
+            using CancellationTokenSource cts = new(CommonTestTimeouts.GeneralTimeout);
+            object lockObj = new();
+            bool stopCompleted = false;
+            bool taskCompleted = false;
+            TaskCompletionSource backgroundTaskCompletion = new(TaskCreationOptions.RunContinuationsAsynchronously);
+            TaskCompletionSource beforeStopCompletion = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            MockBackgroundService service = new MockBackgroundService(async _ =>
+            {
+                await backgroundTaskCompletion.Task.WaitAsync(cts.Token);
+                lock (lockObj)
+                {
+                    Assert.False(stopCompleted, "Stop completed before the background task.");
+                    taskCompleted = true;
+                }
+            });
+
+            // Act
+            service.Start();
+            await service.BackgroundTaskStarted.Task.WaitAsync(cts.Token);
+
+            Task stopTask = Task.Run(async () =>
+            {
+                await Task.Yield();
+                beforeStopCompletion.SetResult();
+                service.Stop();
+
+                lock (lockObj)
+                {
+                    Assert.True(taskCompleted, "Stop completed before the background task.");
+                    stopCompleted = true;
+                }
+            });
+
+            await beforeStopCompletion.Task.WaitAsync(cts.Token);
+            // Wait a bit to ensure Stop() is waiting for the background task to complete
+            await Task.Delay(TimeSpan.FromMilliseconds(100));
+
+            backgroundTaskCompletion.SetResult();
+
+            await stopTask.WaitAsync(cts.Token);
+
+            // Assert
+            Assert.NotNull(service.ExecutingTask);
+            Assert.False(service.ExecutingTask.IsFaulted);
+        }
+
+        [Fact]
+        public async Task WorkerThrows_TaskExceptionIsCaptured()
+        {
+            // Arrange
+            using CancellationTokenSource cts = new(CommonTestTimeouts.GeneralTimeout);
+            MockBackgroundService service = new MockBackgroundService(async _ =>
+            {
+                await Task.Yield();
+                throw new NotImplementedException();
+            });
+
+            // Act
+            service.Start();
+            await service.BackgroundTaskStarted.Task.WaitAsync(cts.Token);
+
+            service.Stop();
+            service.Dispose();
+
+            // Assert
+            Assert.NotNull(service.ExecutingTask);
+            await Assert.ThrowsAsync<NotImplementedException>(() => service.ExecutingTask);
+        }
+    }
+}

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

@@ -0,0 +1,38 @@
+// 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 Func<CancellationToken, Task> _backgroundFunc;
+
+        public MockBackgroundService()
+        {
+            _backgroundFunc = _ => Task.CompletedTask;
+        }
+
+        public MockBackgroundService(Func<CancellationToken, Task> backgroundFunc)
+        {
+            _backgroundFunc = backgroundFunc;
+        }
+
+        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
+        {
+            BackgroundTaskStarted.SetResult();
+
+            await _backgroundFunc(stoppingToken);
+        }
+
+        public override void Dispose()
+        {
+            base.Dispose();
+        }
+
+        public TaskCompletionSource BackgroundTaskStarted { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    }
+}