Fix ArgumentNullException under a race condition in analyzer driver

src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalysisState.AnalyzerStateData.cs

@@ -4,7 +4,9 @@
 
 #nullable disable
 
+using System;
 using System.Collections.Generic;
+using Microsoft.CodeAnalysis.PooledObjects;
 
 namespace Microsoft.CodeAnalysis.Diagnostics
 {
@@ -25,6 +27,16 @@ internal class AnalyzerStateData
             /// </summary>
             public HashSet<AnalyzerAction> ProcessedActions { get; }
 
+            /// <summary>
+            /// Deferred callback for freeing this state object and returning it back to pool of state objects.
+            /// This callback is initialized to a non-null value when
+            /// <see cref="FreeAndReturnToPool{TAnalyzerStateData}(ObjectPool{TAnalyzerStateData})"/>
+            /// is invoked while the state object is still being used by another thread.
+            /// Once the state object is reset back to <see cref="StateKind.ReadyToProcess"/>, we perform
+            /// this deferred free operation and return the object back to the pool.
+            /// </summary>
+            private Action _callbackForDeferredFreeOperation;
+
             public static readonly AnalyzerStateData FullyProcessedInstance = CreateFullyProcessedInstance();
 
             public AnalyzerStateData()
@@ -52,13 +64,49 @@ public virtual void SetStateKind(StateKind stateKind)
             public void ResetToReadyState()
             {
                 SetStateKind(StateKind.ReadyToProcess);
+
+                if (_callbackForDeferredFreeOperation != null)
+                {
+                    _callbackForDeferredFreeOperation();
+                    _callbackForDeferredFreeOperation = null;
+                }
             }
 
             public virtual void Free()
             {
                 this.StateKind = StateKind.ReadyToProcess;
                 this.ProcessedActions.Clear();
             }
+
+            public void FreeAndReturnToPool<TAnalyzerStateData>(ObjectPool<TAnalyzerStateData> pool)
+                where TAnalyzerStateData : AnalyzerStateData
+            {
+                var state = (TAnalyzerStateData)this;
+                if (StateKind == StateKind.InProcess)
+                {
+                    // Do not free the state here if it is still being processed for analysis,
+                    // i.e. StateKind is 'InProcess'. Instead, marked it for a deferred free operation to
+                    // execute after the StateKind is reset to 'ReadyToProcess'.
+                    // This can happen in rare cases when multiple threads are trying to analyze the same
+                    // symbol/syntax/operation, such that one thread has finished all the analyzer callbacks,
+                    // but before this thread marks the state as complete and invokes this Free call,
+                    // second thread starts processing the same symbol/syntax/operation and sets the StateKind to
+                    // 'InProcess' for the same state object. If we free the state here in the first thread,
+                    // then this can lead to data corruption/exceptions in the second thread.
+                    // For example, see https://github.com/dotnet/roslyn/issues/59988.
+                    _callbackForDeferredFreeOperation = () => FreeAndReturnToPoolCore(state, pool);
+                }
+                else
+                {
+                    FreeAndReturnToPoolCore(state, pool);
+                }
+
+                static void FreeAndReturnToPoolCore(TAnalyzerStateData state, ObjectPool<TAnalyzerStateData> pool)
+                {
+                    state.Free();
+                    pool.Free(state);
+                }
+            }
         }
     }
 }

src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalysisState.PerAnalyzerState.cs

@@ -319,8 +319,7 @@ private static void FreeState_NoLock<TAnalyzerStateData>(TAnalyzerStateData? sta
             {
                 if (state != null && !ReferenceEquals(state, AnalyzerStateData.FullyProcessedInstance))
                 {
-                    state.Free();
-                    pool.Free(state);
+                    state.FreeAndReturnToPool(pool);
                 }
             }