Support resetting state when connecting to a previously monitored app (#8154)

* Support resetting state when connecting to a previously monitored app

* Fix switch build issue

* Fix unit tests

* PR Feedback

* Add exception startup collection enablement

* Fix env check

* PR feedback

* PR feedback

* Use an interlocked many reader single writer to synchronize exceptions threads and pipeline stop calls

* Fixup PR tests

* PR feedback

* PR feedback

Co-authored-by: Justin Anderson <[email protected]>

* Fixup schema

* Unit test

* (test)

---------

Co-authored-by: Justin Anderson <[email protected]>

Author: wiktork

cspell.json

@@ -55,6 +55,7 @@
         "hresult",
         "Hsts",
         "Impls",
+        "inproc",
         "JITID",
         "JWTs",
         "LCID",
@@ -107,6 +108,7 @@
         "ukwn",
         "uninitialize",
         "Uninstallation",
+        "uninstrumented",
         "uniquifier",
         "Unlocalized",
         "Unredacted",

documentation/schema.json

@@ -1123,6 +1123,14 @@
               "$ref": "#/definitions/ExceptionsConfiguration"
             }
           ]
+        },
+        "CollectOnStartup": {
+          "type": [
+            "boolean",
+            "null"
+          ],
+          "description": "[Experimental] Determines if exception collection should begin immediately",
+          "default": true
         }
       }
     },

src/Microsoft.Diagnostics.Monitoring.Options/ExceptionsOptions.cs

@@ -27,5 +27,11 @@ public sealed class ExceptionsOptions :
             ResourceType = typeof(OptionsDisplayStrings),
             Description = nameof(OptionsDisplayStrings.DisplayAttributeDescription_ExceptionsOptions_CollectionFilters))]
         public ExceptionsConfiguration? CollectionFilters { get; set; }
+
+        [Display(
+            ResourceType = typeof(OptionsDisplayStrings),
+            Description = nameof(OptionsDisplayStrings.DisplayAttributeDescription_ExceptionsOptions_CollectOnStartup))]
+        [DefaultValue(ExceptionsOptionsDefaults.CollectOnStartup)]
+        public bool? CollectOnStartup { get; set; }
     }
 }

src/Microsoft.Diagnostics.Monitoring.Options/ExceptionsOptionsDefaults.cs

@@ -8,5 +8,7 @@ internal static class ExceptionsOptionsDefaults
         public const bool Enabled = true;
 
         public const int TopLevelLimit = 20;
+
+        public const bool CollectOnStartup = true;
     }
 }

src/Microsoft.Diagnostics.Monitoring.Options/ExceptionsOptionsExtensions.cs

@@ -14,5 +14,10 @@ public static int GetTopLevelLimit(this ExceptionsOptions options)
         {
             return options.TopLevelLimit.GetValueOrDefault(ExceptionsOptionsDefaults.TopLevelLimit);
         }
+
+        public static bool CollectOnStartup(this ExceptionsOptions options)
+        {
+            return options.CollectOnStartup.GetValueOrDefault(ExceptionsOptionsDefaults.CollectOnStartup);
+        }
     }
 }

src/Microsoft.Diagnostics.Monitoring.Options/OptionsDisplayStrings.Designer.cs

@@ -780,4 +780,7 @@
     <value>Performs a match based on the name of the managed entry point assembly of the process.</value>
     <comment>The description provided for the ManagedEntryPointAssemblyName parameter on ProcessFilterDescriptor.</comment>
   </data>
+  <data name="DisplayAttributeDescription_ExceptionsOptions_CollectOnStartup" xml:space="preserve">
+    <value>[Experimental] Determines if exception collection should begin immediately</value>
+  </data>
 </root>

src/Microsoft.Diagnostics.Monitoring.Options/OptionsDisplayStrings.resx

@@ -23,7 +23,13 @@ internal sealed class DiagnosticsBootstrapper :
         public DiagnosticsBootstrapper()
         {
             _exceptionProcessor = new(ToolIdentifiers.IsEnvVarEnabled(InProcessFeaturesIdentifiers.EnvironmentVariables.Exceptions.IncludeMonitorExceptions));
-            _exceptionProcessor.Start();
+
+            // If collectOnStart is not set, assume the default value (true).
+            string? collectOnStart = Environment.GetEnvironmentVariable(InProcessFeaturesIdentifiers.EnvironmentVariables.Exceptions.CollectOnStartup);
+            if (collectOnStart == null || ToolIdentifiers.IsEnvVarValueEnabled(collectOnStart))
+            {
+                _exceptionProcessor.Start();
+            }
 
             using IDisposable _ = MonitorExecutionContextTracker.MonitorScope();
 
@@ -35,6 +41,17 @@ public DiagnosticsBootstrapper()
                     SharedInternals.MessageDispatcher = new MessageDispatcher.MonitorMessageDispatcher(
                         new MessageDispatcher.ProfilerMessageSource(CommandSet.StartupHook));
                     ToolIdentifiers.EnableEnvVar(InProcessFeaturesIdentifiers.EnvironmentVariables.AvailableInfrastructure.ManagedMessaging);
+
+                    SharedInternals.MessageDispatcher.RegisterCallback<EmptyPayload>(StartupHookCommand.StopAllFeatures, (IpcMessage) =>
+                    {
+                        _exceptionProcessor.Stop();
+                        _parameterCapturingService?.RequestStopAll();
+                    });
+
+                    SharedInternals.MessageDispatcher.RegisterCallback<EmptyPayload>(StartupHookCommand.StartAllFeatures, (IpcMessage) =>
+                    {
+                        _exceptionProcessor.Start();
+                    });
                 }
 
                 if (ToolIdentifiers.IsEnvVarEnabled(InProcessFeaturesIdentifiers.EnvironmentVariables.ParameterCapturing.Enable))

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

@@ -14,10 +14,10 @@ internal sealed class CurrentAppDomainExceptionProcessor : IDisposable
         private readonly ExceptionIdSource _idSource = new();
 
         private readonly CurrentAppDomainFirstChanceExceptionSource _firstChanceSource;
-        private readonly ExceptionPipeline _firstChancePipeline;
+        private ExceptionPipeline _firstChancePipeline;
 
         private readonly CurrentAppDomainUnhandledExceptionSource _unhandledSource;
-        private readonly ExceptionPipeline _unhandledPipeline;
+        private ExceptionPipeline _unhandledPipeline;
 
         private readonly bool _includeMonitorExceptions;
 
@@ -35,7 +35,23 @@ public CurrentAppDomainExceptionProcessor(bool includeMonitorExceptions)
         public void Start()
         {
             _firstChancePipeline.Start();
+            _firstChanceSource.Start();
+
             _unhandledPipeline.Start();
+            _unhandledSource.Start();
+        }
+
+        public void Stop()
+        {
+            // Stop all exception flow from the sources. This will allow draining the pipelines.
+            _firstChanceSource.Stop();
+            _unhandledSource.Stop();
+            _firstChancePipeline.Stop();
+            _unhandledPipeline.Stop();
+
+            // We reset the pipelines entirely to allow their cache to be cleared
+            _firstChancePipeline = new(_firstChanceSource, ConfigureFirstChancePipeline);
+            _unhandledPipeline = new(_unhandledSource, ConfigureUnhandledPipeline);
         }
 
         private void ConfigureFirstChancePipeline(ExceptionPipelineBuilder builder)

src/Microsoft.Diagnostics.Monitoring.StartupHook/Exceptions/CurrentAppDomainExceptionProcessor.cs

@@ -15,11 +15,16 @@ internal sealed class CurrentAppDomainFirstChanceExceptionSource :
     {
         private long _disposedState;
 
-        public CurrentAppDomainFirstChanceExceptionSource()
+        public override void Start()
         {
             AppDomain.CurrentDomain.FirstChanceException += CurrentDomain_FirstChanceException;
         }
 
+        public override void Stop()
+        {
+            AppDomain.CurrentDomain.FirstChanceException -= CurrentDomain_FirstChanceException;
+        }
+
         private void CurrentDomain_FirstChanceException(object? sender, FirstChanceExceptionEventArgs e)
         {
             RaiseExceptionGuarded(e.Exception);

src/Microsoft.Diagnostics.Monitoring.StartupHook/Exceptions/CurrentAppDomainFirstChanceExceptionSource.cs

@@ -14,11 +14,16 @@ internal sealed class CurrentAppDomainUnhandledExceptionSource :
     {
         private long _disposedState;
 
-        public CurrentAppDomainUnhandledExceptionSource()
+        public override void Start()
         {
             AppDomain.CurrentDomain.UnhandledException += CurrentDomain_UnhandledException;
         }
 
+        public override void Stop()
+        {
+            AppDomain.CurrentDomain.UnhandledException -= CurrentDomain_UnhandledException;
+        }
+
         private void CurrentDomain_UnhandledException(object? sender, UnhandledExceptionEventArgs e)
         {
             if (e.ExceptionObject is Exception exception)

src/Microsoft.Diagnostics.Monitoring.StartupHook/Exceptions/CurrentAppDomainUnhandledExceptionSource.cs

@@ -16,6 +16,10 @@ protected void RaiseException(Exception ex, DateTime timestamp, string? activity
             ExceptionAvailable?.Invoke(this, new ExceptionAvailableEventArgs(ex, timestamp, activityId, format));
         }
 
+        public virtual void Start() { }
+
+        public virtual void Stop() { }
+
         /// <summary>
         /// Event that is raised each time an exception is made available.
         /// </summary>

src/Microsoft.Diagnostics.Monitoring.StartupHook/Exceptions/ExceptionSourceBase.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Threading;
 
 namespace Microsoft.Diagnostics.Monitoring.StartupHook.Exceptions.Pipeline
 {
@@ -13,6 +14,9 @@ internal sealed class ExceptionPipeline :
 
         private long _disposedState;
 
+        private long _processingState;
+        private const long WriteFlag = unchecked((long)0x8000000000000000);
+
         public ExceptionPipeline(ExceptionSourceBase exceptionSource, Action<ExceptionPipelineBuilder> configure)
         {
             ArgumentNullException.ThrowIfNull(exceptionSource);
@@ -28,10 +32,25 @@ public ExceptionPipeline(ExceptionSourceBase exceptionSource, Action<ExceptionPi
         public void Start()
         {
             DisposableHelper.ThrowIfDisposed<ExceptionPipeline>(ref _disposedState);
+            _processingState = 0;
 
             _exceptionSource.ExceptionAvailable += ExceptionSource_ExceptionAvailable;
         }
 
+        public void Stop()
+        {
+            // We must do this first to prevent any new exception callbacks.
+            _exceptionSource.ExceptionAvailable -= ExceptionSource_ExceptionAvailable;
+
+            // Wait until all outstanding exception callbacks are finished. Once _processingState reaches 0, we will set the WriteFlag.
+            // This will prevent any further exception handlers from processing.
+            while (0 != Interlocked.CompareExchange(ref _processingState, WriteFlag, 0) && !DisposableHelper.IsDisposed(ref _disposedState))
+            {
+                // Wait for all exceptions to be processed
+                Thread.Sleep(100);
+            }
+        }
+
         private void ExceptionSource_ExceptionAvailable(object? sender, ExceptionAvailableEventArgs args)
         {
             // DESIGN: While async patterns are typically favored over synchronous patterns,
@@ -40,12 +59,45 @@ private void ExceptionSource_ExceptionAvailable(object? sender, ExceptionAvailab
             // (e.g. EventSource provides events but diagnostic pipe events are queued and asynchronously emitted).
             // Synchronous execution is required for scenarios where the exception needs to be held
             // at the site of where it is thrown before allowing it to unwind (e.g. capturing a dump of the exception).
-            _exceptionHandler.Invoke(
-                args.Exception,
-                new ExceptionPipelineExceptionContext(
-                    args.Timestamp,
-                    args.ActivityId,
-                    args.ActivityIdFormat));
+
+            SpinWait spinWait = new();
+
+            while (true)
+            {
+                long state = Interlocked.Read(ref _processingState);
+                if ((state & WriteFlag) == WriteFlag)
+                {
+                    // Stop has been requested but the event handler has already occurred. Return early.
+                    return;
+                }
+                if (DisposableHelper.IsDisposed(ref _disposedState))
+                {
+                    return;
+                }
+                // Increment the value by 1 to indicate that we are doing work.
+                // We cannot do a simple increment since it's possible for a Stop call to attempt to set the write bit
+                // at this stage.
+                if (Interlocked.CompareExchange(ref _processingState, state + 1, state) == state)
+                {
+                    break;
+                }
+                spinWait.SpinOnce();
+            }
+
+            try
+            {
+                _exceptionHandler.Invoke(
+                    args.Exception,
+                    new ExceptionPipelineExceptionContext(
+                        args.Timestamp,
+                        args.ActivityId,
+                        args.ActivityIdFormat));
+
+            }
+            finally
+            {
+                Interlocked.Decrement(ref _processingState);
+            }
         }
 
 

src/Microsoft.Diagnostics.Monitoring.StartupHook/Exceptions/Pipeline/ExceptionPipeline.cs

@@ -220,5 +220,10 @@ public override void Dispose()
 
             base.Dispose();
         }
+
+        public void RequestStopAll()
+        {
+            _pipeline?.RequestStopAll();
+        }
     }
 }

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

@@ -206,6 +206,17 @@ public void SubmitRequest(StartCapturingParametersPayload payload, IFunctionProb
             }
         }
 
+        public void RequestStopAll()
+        {
+            foreach (var key in _allRequests.Keys)
+            {
+                if (_allRequests.TryRemove(key, out CapturingRequest? request))
+                {
+                    request?.StopRequest?.TrySetResult();
+                }
+            }
+        }
+
         public void RequestStop(Guid requestId)
         {
             if (!_allRequests.TryGetValue(requestId, out CapturingRequest? request))

src/Microsoft.Diagnostics.Monitoring.StartupHook/ParameterCapturing/Pipeline/ParameterCapturingPipeline.cs

@@ -11,5 +11,7 @@ public interface IInProcessFeatures
         bool IsStartupHookRequired { get; }
 
         bool IsLibrarySharingRequired { get; }
+
+        bool CollectExceptionsOnStartup { get; }
     }
 }

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

@@ -27,6 +27,13 @@ public ProfilerChannel(IOptionsMonitor<StorageOptions> storageOptions)
             _storageOptions = storageOptions;
         }
 
+        public async Task SendMessage(IEndpointInfo endpointInfo, IProfilerMessage message, CancellationToken token, TimeSpan timeout)
+        {
+            using CancellationTokenSource cancellationTokenSource = new(timeout);
+            using CancellationTokenSource linkedCts = CancellationTokenSource.CreateLinkedTokenSource(token, cancellationTokenSource.Token);
+            await SendMessage(endpointInfo, message, linkedCts.Token);
+        }
+
         public async Task SendMessage(IEndpointInfo endpointInfo, IProfilerMessage message, CancellationToken token)
         {
             if (message.Payload.Length > MaxPayloadSize)

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

@@ -26,13 +26,17 @@ public enum ServerResponseCommand : ushort
 
     public enum ProfilerCommand : ushort
     {
-        Callstack
+        Callstack,
+        StopAllFeatures,
+        StartAllFeatures,
     };
 
     public enum StartupHookCommand : ushort
     {
         StartCapturingParameters,
-        StopCapturingParameters
+        StopCapturingParameters,
+        StopAllFeatures,
+        StartAllFeatures,
     };
 
     public interface IProfilerMessage