Do not attempt to prune processes during tracing

src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/DiagController.cs

@@ -46,6 +46,7 @@ public partial class DiagController : ControllerBase
         private readonly IOptionsMonitor<GlobalCounterOptions> _counterOptions;
         private readonly EgressOperationStore _operationsStore;
         private readonly IDumpService _dumpService;
+        private readonly OperationTrackerService _operationTrackerService;
 
         public DiagController(ILogger<DiagController> logger,
             IServiceProvider serviceProvider)
@@ -56,6 +57,7 @@ public DiagController(ILogger<DiagController> logger,
             _operationsStore = serviceProvider.GetRequiredService<EgressOperationStore>();
             _dumpService = serviceProvider.GetRequiredService<IDumpService>();
             _counterOptions = serviceProvider.GetRequiredService<IOptionsMonitor<GlobalCounterOptions>>();
+            _operationTrackerService = serviceProvider.GetRequiredService<OperationTrackerService>();
         }
 
         /// <summary>
@@ -537,7 +539,22 @@ private Task<ActionResult> StartTrace(
             return Result(
                 Utilities.ArtifactType_Trace,
                 egressProvider,
-                (outputStream, token) => TraceUtilities.CaptureTraceAsync(null, processInfo.EndpointInfo, configuration, duration, outputStream, token),
+                async (outputStream, token) =>
+                {
+                    IDisposable operationRegistration = null;
+                    try
+                    {
+                        if (_diagnosticPortOptions.Value.ConnectionMode == DiagnosticPortConnectionMode.Listen)
+                        {
+                            operationRegistration = _operationTrackerService.Register(processInfo.EndpointInfo);
+                        }
+                        await TraceUtilities.CaptureTraceAsync(null, processInfo.EndpointInfo, configuration, duration, outputStream, token);
+                    }
+                    finally
+                    {
+                        operationRegistration?.Dispose();
+                    }
+                },
                 fileName,
                 ContentTypes.ApplicationOctetStream,
                 processInfo.EndpointInfo);

src/Microsoft.Diagnostics.Monitoring.WebApi/DumpService.cs

@@ -17,16 +17,18 @@ namespace Microsoft.Diagnostics.Monitoring.WebApi
 {
     internal class DumpService : IDumpService
     {
-        private readonly ConcurrentDictionary<IEndpointInfo, OperationsTracker> _operations = new();
+        private readonly OperationTrackerService _operationTrackerService;
         private readonly IOptionsMonitor<DiagnosticPortOptions> _portOptions;
         private readonly IOptionsMonitor<StorageOptions> _storageOptions;
 
         public DumpService(
             IOptionsMonitor<StorageOptions> storageOptions,
-            IOptionsMonitor<DiagnosticPortOptions> portOptions)
+            IOptionsMonitor<DiagnosticPortOptions> portOptions,
+            OperationTrackerService operationTrackerService)
         {
             _portOptions = portOptions ?? throw new ArgumentNullException(nameof(portOptions));
             _storageOptions = storageOptions ?? throw new ArgumentNullException(nameof(StorageOptions));
+            _operationTrackerService = operationTrackerService ?? throw new ArgumentNullException(nameof(operationTrackerService));
         }
 
         public async Task<Stream> DumpAsync(IEndpointInfo endpointInfo, Models.DumpType mode, CancellationToken token)
@@ -50,7 +52,7 @@ public async Task<Stream> DumpAsync(IEndpointInfo endpointInfo, Models.DumpType
                 // a dump operation causing the runtime to temporarily be unresponsive. Long term, this
                 // concept should be folded into RequestLimitTracker with registered endpoint information
                 // and allowing to query the status of an endpoint for a given artifact.
-                operationRegistration = GetOperationsTracker(endpointInfo).Register();
+                operationRegistration = _operationTrackerService.Register(endpointInfo);
             }
 
             try
@@ -75,25 +77,6 @@ public async Task<Stream> DumpAsync(IEndpointInfo endpointInfo, Models.DumpType
             return new AutoDeleteFileStream(dumpFilePath);
         }
 
-        public bool IsExecutingOperation(IEndpointInfo endpointInfo)
-        {
-            if (IsListenMode)
-            {
-                return GetOperationsTracker(endpointInfo).IsExecutingOperation;
-            }
-            return false;
-        }
-
-        public void EndpointRemoved(IEndpointInfo endpointInfo)
-        {
-            _operations.TryRemove(endpointInfo, out _);
-        }
-
-        private OperationsTracker GetOperationsTracker(IEndpointInfo endpointInfo)
-        {
-            return _operations.GetOrAdd(endpointInfo, _ => new OperationsTracker());
-        }
-
         private bool IsListenMode => _portOptions.CurrentValue.GetConnectionMode() == DiagnosticPortConnectionMode.Listen;
 
         /// <summary>
@@ -134,19 +117,6 @@ private static DumpType MapDumpType(Models.DumpType dumpType)
             }
         }
 
-        private sealed class OperationsTracker : IDisposable
-        {
-            private int _count = 0;
 
-            public IDisposable Register()
-            {
-                Interlocked.Increment(ref _count);
-                return this;
-            }
-
-            public bool IsExecutingOperation => 0 != Interlocked.CompareExchange(ref _count, 0, 0);
-
-            public void Dispose() => Interlocked.Decrement(ref _count);
-        }
     }
 }

src/Microsoft.Diagnostics.Monitoring.WebApi/IDumpService.cs

@@ -11,9 +11,5 @@ namespace Microsoft.Diagnostics.Monitoring.WebApi
     internal interface IDumpService
     {
         Task<Stream> DumpAsync(IEndpointInfo endpointInfo, Models.DumpType mode, CancellationToken token);
-
-        bool IsExecutingOperation(IEndpointInfo endpointInfo);
-
-        void EndpointRemoved(IEndpointInfo endpointInfo);
     }
 }

src/Microsoft.Diagnostics.Monitoring.WebApi/OperationTrackerService.cs

@@ -0,0 +1,50 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Concurrent;
+using System.Threading;
+
+namespace Microsoft.Diagnostics.Monitoring.WebApi
+{
+    internal sealed class OperationTrackerService
+    {
+        private readonly ConcurrentDictionary<IEndpointInfo, OperationsTracker> _operations = new();
+
+        private sealed class OperationsTracker : IDisposable
+        {
+            private int _count = 0;
+
+            public IDisposable Register()
+            {
+                Interlocked.Increment(ref _count);
+                return this;
+            }
+
+            public bool IsExecutingOperation => 0 != Interlocked.CompareExchange(ref _count, 0, 0);
+
+            public void Dispose() => Interlocked.Decrement(ref _count);
+        }
+
+        public bool IsExecutingOperation(IEndpointInfo endpointInfo)
+        {
+            return GetOperationsTracker(endpointInfo).IsExecutingOperation;
+        }
+
+        public IDisposable Register(IEndpointInfo endpointInfo)
+        {
+            return GetOperationsTracker(endpointInfo).Register();
+        }
+
+        public void EndpointRemoved(IEndpointInfo endpointInfo)
+        {
+            _operations.TryRemove(endpointInfo, out _);
+        }
+
+        private OperationsTracker GetOperationsTracker(IEndpointInfo endpointInfo)
+        {
+            return _operations.GetOrAdd(endpointInfo, _ => new OperationsTracker());
+        }
+    }
+}

src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/EndpointInfoSourceTests.cs

@@ -161,8 +161,9 @@ await runners.ExecuteAsync(async () =>
         public async Task ServerSourceNoPruneDuringDumpTest(TargetFrameworkMoniker appTfm)
         {
             EndpointInfoSourceCallback callback = new(_outputHelper);
-            MockDumpService dumpService = new();
-            await using ServerSourceHolder sourceHolder = await _endpointUtilities.StartServerAsync(callback, dumpService);
+            var operationTrackerService = new OperationTrackerService();
+            MockDumpService dumpService = new(operationTrackerService);
+            await using ServerSourceHolder sourceHolder = await _endpointUtilities.StartServerAsync(callback, dumpService, operationTrackerService);
 
             AppRunner runner = _endpointUtilities.CreateAppRunner(sourceHolder.TransportName, appTfm);
 
@@ -226,30 +227,27 @@ private sealed class MockDumpService : IDumpService
             private readonly TaskCompletionSource<Stream> _dumpCompletionSource =
                 new(TaskCreationOptions.RunContinuationsAsynchronously);
 
-            private int _operationCount = 0;
+            private readonly OperationTrackerService _operationTrackerService;
+
+            public MockDumpService(OperationTrackerService operationTrackerService)
+            {
+                _operationTrackerService = operationTrackerService;
+            }
 
             public async Task<Stream> DumpAsync(IEndpointInfo endpointInfo, DumpType mode, CancellationToken token)
             {
-                Interlocked.Increment(ref _operationCount);
+                IDisposable operationRegistration = null;
                 try
                 {
+                    operationRegistration = _operationTrackerService.Register(endpointInfo);
                     return await _dumpCompletionSource.Task;
                 }
                 finally
                 {
-                    Interlocked.Decrement(ref _operationCount);
+                    operationRegistration?.Dispose();
                 }
             }
 
-            public void EndpointRemoved(IEndpointInfo endpointInfo)
-            {
-            }
-
-            public bool IsExecutingOperation(IEndpointInfo endpointInfo)
-            {
-                return 0 != Interlocked.CompareExchange(ref _operationCount, 0, 0);
-            }
-
             public void CompleteOperation()
             {
                 _dumpCompletionSource.SetResult(new MemoryStream());

src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/EndpointUtilities.cs

@@ -28,7 +28,8 @@ public EndpointUtilities(ITestOutputHelper outputHelper)
             _outputHelper = outputHelper;
         }
 
-        public async Task<ServerSourceHolder> StartServerAsync(EndpointInfoSourceCallback sourceCallback = null, IDumpService dumpService = null)
+        public async Task<ServerSourceHolder> StartServerAsync(EndpointInfoSourceCallback sourceCallback = null, IDumpService dumpService = null,
+            OperationTrackerService operationTrackerService = null)
         {
             DiagnosticPortHelper.Generate(DiagnosticPortConnectionMode.Listen, out _, out string transportName);
             _outputHelper.WriteLine("Starting server endpoint info source at '" + transportName + "'.");
@@ -39,7 +40,7 @@ public async Task<ServerSourceHolder> StartServerAsync(EndpointInfoSourceCallbac
                 callbacks.Add(sourceCallback);
                 if (null != dumpService)
                 {
-                    callbacks.Add(new DumpServiceEndpointInfoSourceCallback(dumpService));
+                    callbacks.Add(new OperationTrackerServiceEndpointInfoSourceCallback(operationTrackerService));
                 }
             }
 
@@ -50,7 +51,7 @@ public async Task<ServerSourceHolder> StartServerAsync(EndpointInfoSourceCallbac
                     EndpointName = transportName
                 });
 
-            ServerEndpointInfoSource source = new(portOptions, callbacks, dumpService);
+            ServerEndpointInfoSource source = new(portOptions, callbacks, operationTrackerService);
 
             await source.StartAsync(CancellationToken.None);
 

src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/TestHostHelper.cs

@@ -86,6 +86,7 @@ public static IHost CreateHost(
                 .ConfigureServices((HostBuilderContext context, IServiceCollection services) =>
                 {
                     services.ConfigureGlobalCounter(context.Configuration);
+                    services.AddSingleton<OperationTrackerService>();
                     services.ConfigureCollectionRules();
                     services.ConfigureEgress();
 

src/Tools/dotnet-monitor/CollectionRules/Actions/CollectTraceAction.cs

@@ -45,12 +45,14 @@ private sealed class CollectTraceAction :
         {
             private readonly IServiceProvider _serviceProvider;
             private readonly IOptionsMonitor<GlobalCounterOptions> _counterOptions;
+            private readonly OperationTrackerService _operationTrackerService;
 
             public CollectTraceAction(IServiceProvider serviceProvider, IEndpointInfo endpointInfo, CollectTraceOptions options)
                 : base(endpointInfo, options)
             {
                 _serviceProvider = serviceProvider ?? throw new ArgumentNullException(nameof(serviceProvider));
                 _counterOptions = _serviceProvider.GetRequiredService<IOptionsMonitor<GlobalCounterOptions>>();
+                _operationTrackerService = _serviceProvider.GetRequiredService<OperationTrackerService>();
             }
 
             protected override async Task<CollectionRuleActionResult> ExecuteCoreAsync(
@@ -83,7 +85,11 @@ protected override async Task<CollectionRuleActionResult> ExecuteCoreAsync(
                 KeyValueLogScope scope = Utils.CreateArtifactScope(Utils.ArtifactType_Trace, EndpointInfo);
 
                 EgressOperation egressOperation = new EgressOperation(
-                    (outputStream, token) => TraceUtilities.CaptureTraceAsync(startCompletionSource, EndpointInfo, configuration, duration, outputStream, token),
+                    async (outputStream, token) =>
+                    {
+                        using IDisposable operationRegistration = _operationTrackerService.Register(EndpointInfo);
+                        await TraceUtilities.CaptureTraceAsync(startCompletionSource, EndpointInfo, configuration, duration, outputStream, token);
+                    },
                     egressProvider,
                     fileName,
                     EndpointInfo,

src/Tools/dotnet-monitor/DiagnosticsMonitorCommandHandler.cs

@@ -243,6 +243,7 @@ public static IHostBuilder CreateHostBuilder(IConsole console, string[] urls, st
 
                     services.Configure<DiagnosticPortOptions>(context.Configuration.GetSection(ConfigurationKeys.DiagnosticPort));
                     services.AddSingleton<IValidateOptions<DiagnosticPortOptions>, DiagnosticPortValidateOptions>();
+                    services.AddSingleton<OperationTrackerService>();
 
                     services.ConfigureGlobalCounter(context.Configuration);
 
@@ -251,7 +252,7 @@ public static IHostBuilder CreateHostBuilder(IConsole console, string[] urls, st
                     services.AddHostedServiceForwarder<ServerEndpointInfoSource>();
                     services.AddSingleton<IDiagnosticServices, DiagnosticServices>();
                     services.AddSingleton<IDumpService, DumpService>();
-                    services.AddSingleton<IEndpointInfoSourceCallbacks, DumpServiceEndpointInfoSourceCallback>();
+                    services.AddSingleton<IEndpointInfoSourceCallbacks, OperationTrackerServiceEndpointInfoSourceCallback>();
                     services.AddSingleton<RequestLimitTracker>();
                     services.ConfigureOperationStore();
                     services.ConfigureEgress();

src/Tools/dotnet-monitor/EndpointInfo/ServerEndpointInfoSource.cs

@@ -5,6 +5,7 @@
 using Microsoft.Diagnostics.Monitoring.WebApi;
 using Microsoft.Diagnostics.NETCore.Client;
 using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using System;
 using System.Collections.Generic;
@@ -42,7 +43,8 @@ internal sealed class ServerEndpointInfoSource : BackgroundService, IEndpointInf
         private readonly IEnumerable<IEndpointInfoSourceCallbacks> _callbacks;
         private readonly DiagnosticPortOptions _portOptions;
 
-        private readonly IDumpService _dumpService;
+        private readonly OperationTrackerService _operationTrackerService;
+        private readonly ILogger<ServerEndpointInfoSource> _logger;
 
         /// <summary>
         /// Constructs a <see cref="ServerEndpointInfoSource"/> that aggregates diagnostic endpoints
@@ -51,11 +53,13 @@ internal sealed class ServerEndpointInfoSource : BackgroundService, IEndpointInf
         public ServerEndpointInfoSource(
             IOptions<DiagnosticPortOptions> portOptions,
             IEnumerable<IEndpointInfoSourceCallbacks> callbacks = null,
-            IDumpService dumpService = null)
+            OperationTrackerService operationTrackerService = null,
+            ILogger<ServerEndpointInfoSource> logger = null)
         {
             _callbacks = callbacks ?? Enumerable.Empty<IEndpointInfoSourceCallbacks>();
-            _dumpService = dumpService;
+            _operationTrackerService = operationTrackerService;
             _portOptions = portOptions.Value;
+            _logger = logger;
 
             BoundedChannelOptions channelOptions = new(PendingRemovalChannelCapacity)
             {
@@ -253,7 +257,7 @@ private async Task<bool> CheckEndpointAsync(EndpointInfo info, CancellationToken
             // If a dump operation is in progress, the runtime is likely to not respond to
             // diagnostic requests. Do not check for responsiveness while the dump operation
             // is in progress.
-            if (_dumpService?.IsExecutingOperation(info) == true)
+            if (_operationTrackerService?.IsExecutingOperation(info) == true)
             {
                 return true;
             }
@@ -269,6 +273,7 @@ private async Task<bool> CheckEndpointAsync(EndpointInfo info, CancellationToken
             }
             catch (OperationCanceledException) when (timeoutSource.IsCancellationRequested)
             {
+                _logger?.EndpointTimeout(info.ProcessId.ToString(System.Globalization.CultureInfo.InvariantCulture));
                 return false;
             }
             catch (Exception ex) when (ex is not OperationCanceledException)

src/Tools/dotnet-monitor/LoggingEventIds.cs

@@ -59,5 +59,6 @@ internal static class LoggingEventIds
         public const int InvalidActionReference = 49;
         public const int InvalidActionResultReference = 50;
         public const int ActionSettingsTokenizationNotSupported = 51;
+        public const int EndpointTimeout = 52;
     }
 }

src/Tools/dotnet-monitor/LoggingExtensions.cs

@@ -277,6 +277,12 @@ internal static class LoggingExtensions
             logLevel: LogLevel.Error,
             formatString: Strings.LogFormatString_ActionSettingsTokenizationNotSupported);
 
+        private static readonly Action<ILogger, string, Exception> _endpointTimeout =
+            LoggerMessage.Define<string>(
+                eventId: new EventId(LoggingEventIds.EndpointTimeout, "EndpointTimeout"),
+                logLevel: LogLevel.Warning,
+                formatString: Strings.LogFormatString_EndpointTimeout);
+
         public static void EgressProviderInvalidOptions(this ILogger logger, string providerName)
         {
             _egressProviderInvalidOptions(logger, providerName, null);
@@ -500,5 +506,10 @@ public static void ActionSettingsTokenizationNotSupported(this ILogger logger, s
         {
             _actionSettingsTokenizationNotSupported(logger, settingsType, null);
         }
+
+        public static void EndpointTimeout(this ILogger logger, string processId)
+        {
+            _endpointTimeout(logger, processId, null);
+        }
     }
 }