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

Author: wiktork

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

@@ -14,7 +14,8 @@ internal sealed class ExceptionPipeline :
 
         private long _disposedState;
 
-        private int _processingState;
+        private long _processingState;
+        private const long WriteFlag = unchecked((long)0x8000000000000000);
 
         public ExceptionPipeline(ExceptionSourceBase exceptionSource, Action<ExceptionPipelineBuilder> configure)
         {
@@ -31,15 +32,19 @@ 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;
 
-            while (0 != Interlocked.CompareExchange(ref _processingState, 0, 0) && !DisposableHelper.IsDisposed(ref _disposedState))
+            // 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);
@@ -55,9 +60,32 @@ private void ExceptionSource_ExceptionAvailable(object? sender, ExceptionAvailab
             // 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).
 
+            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
             {
-                Interlocked.Increment(ref _processingState);
                 _exceptionHandler.Invoke(
                     args.Exception,
                     new ExceptionPipelineExceptionContext(