Merge pull request #9571 from mavasani/AnalyzerStateAllocations

Port analyzer driver pooling fix to future-stabilization

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

@@ -15,6 +15,7 @@ internal partial class AnalysisState
     {
         private class PerAnalyzerState
         {
+            private const int SymbolLimitForPooling = 100;
             private readonly object _gate = new object();
             private readonly Dictionary<CompilationEvent, AnalyzerStateData> _pendingEvents = new Dictionary<CompilationEvent, AnalyzerStateData>();
             private readonly Dictionary<ISymbol, AnalyzerStateData> _pendingSymbols = new Dictionary<ISymbol, AnalyzerStateData>();
@@ -37,6 +38,48 @@ private class PerAnalyzerState
                 _currentlyAnalyzingDeclarationsMapPool = currentlyAnalyzingDeclarationsMapPool;
             }
 
+            /// <summary>
+            /// Returns true if the object should be returned to the pool.
+            /// </summary>
+            public bool Free()
+            {
+                lock (_gate)
+                {
+                    foreach (var analyzerStateData in _pendingEvents.Values)
+                    {
+                        FreeState_NoLock(analyzerStateData, _analyzerStateDataPool);
+                    }
+
+                    foreach (var analyzerStateData in _pendingSymbols.Values)
+                    {
+                        FreeState_NoLock(analyzerStateData, _analyzerStateDataPool);
+                    }
+
+                    foreach (var declarationDataMap in _pendingDeclarations.Values)
+                    {
+                        foreach (var declarationStateData in declarationDataMap.Values)
+                        {
+                            FreeDeclarationAnalyzerState_NoLock(declarationStateData);
+                        }
+
+                        FreeDeclarationDataMap_NoLock(declarationDataMap);
+                    }
+
+                    // If we have too many symbols then just discard the state object from the pool - we don't want to hold onto really large dictionaries.
+                    if (_pendingSymbols.Count > SymbolLimitForPooling)
+                    {
+                        return false;
+                    }
+
+                    _pendingEvents.Clear();
+                    _pendingSymbols.Clear();
+                    _pendingDeclarations.Clear();
+                    _lazySyntaxTreesWithAnalysisData = null;
+                    _pendingSyntaxAnalysisTreesCount = 0;
+                    return true;
+                }
+            }
+
             public void AddPendingEvents(HashSet<CompilationEvent> uniqueEvents)
             {
                 lock (_gate)
@@ -275,7 +318,7 @@ private void FreeDeclarationAnalyzerState_NoLock(DeclarationAnalyzerStateData st
             private static void FreeState_NoLock<TAnalyzerStateData>(TAnalyzerStateData state, ObjectPool<TAnalyzerStateData> pool)
                 where TAnalyzerStateData : AnalyzerStateData
             {
-                if (state != null)
+                if (state != null && !ReferenceEquals(state, AnalyzerStateData.FullyProcessedInstance))
                 {
                     state.Free();
                     pool.Free(state);

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

@@ -63,6 +63,13 @@ internal partial class AnalysisState
         private readonly ObjectPool<HashSet<CompilationEvent>> _compilationEventsPool;
         private readonly HashSet<CompilationEvent> _pooledEventsWithAnyActionsSet;
 
+        // Create static pools for heavily allocated per-analyzer state objects - this helps in reducing allocations across CompilationWithAnalyzer instances.
+        private const int PoolSize = 5000;
+        private static readonly ObjectPool<AnalyzerStateData> s_analyzerStateDataPool = new ObjectPool<AnalyzerStateData>(() => new AnalyzerStateData(), PoolSize);
+        private static readonly ObjectPool<DeclarationAnalyzerStateData> s_declarationAnalyzerStateDataPool = new ObjectPool<DeclarationAnalyzerStateData>(() => new DeclarationAnalyzerStateData(), PoolSize);
+        private static readonly ObjectPool<Dictionary<int, DeclarationAnalyzerStateData>> s_currentlyAnalyzingDeclarationsMapPool = new ObjectPool<Dictionary<int, DeclarationAnalyzerStateData>>(() => new Dictionary<int, DeclarationAnalyzerStateData>(), PoolSize);
+        private static readonly ObjectPool<PerAnalyzerState> s_perAnalyzerStatePool = new ObjectPool<PerAnalyzerState>(() => new PerAnalyzerState(s_analyzerStateDataPool, s_declarationAnalyzerStateDataPool, s_currentlyAnalyzingDeclarationsMapPool), PoolSize);
+
         public AnalysisState(ImmutableArray<DiagnosticAnalyzer> analyzers, CompilationData compilationData)
         {
             _gate = new object();
@@ -80,19 +87,33 @@ public AnalysisState(ImmutableArray<DiagnosticAnalyzer> analyzers, CompilationDa
             _pooledEventsWithAnyActionsSet = new HashSet<CompilationEvent>();
         }
 
-        private static ImmutableDictionary<DiagnosticAnalyzer, int> CreateAnalyzerStateMap(ImmutableArray<DiagnosticAnalyzer> analyzers, out ImmutableArray<PerAnalyzerState> analyzerStates)
+        ~AnalysisState()
         {
-            var analyzerStateDataPool = new ObjectPool<AnalyzerStateData>(() => new AnalyzerStateData());
-            var declarationAnalyzerStateDataPool = new ObjectPool<DeclarationAnalyzerStateData>(() => new DeclarationAnalyzerStateData());
-            var currentlyAnalyzingDeclarationsMapPool = new ObjectPool<Dictionary<int, DeclarationAnalyzerStateData>>(
-                () => new Dictionary<int, DeclarationAnalyzerStateData>());
+            // Free the per-analyzer state tracking objects.
+            foreach (var analyzerState in _analyzerStates)
+            {
+                var shouldReturnToPool = analyzerState.Free();
+
+                // If we have too many symbols then just discard the state object from the pool - we don't want to hold onto really large dictionaries.
+                if (shouldReturnToPool)
+                {
+                    s_perAnalyzerStatePool.Free(analyzerState);
+                }
+                else
+                {
+                    s_perAnalyzerStatePool.ForgetTrackedObject(analyzerState);
+                }
+            }
+        }
 
+        private static ImmutableDictionary<DiagnosticAnalyzer, int> CreateAnalyzerStateMap(ImmutableArray<DiagnosticAnalyzer> analyzers, out ImmutableArray<PerAnalyzerState> analyzerStates)
+        {
             var statesBuilder = ImmutableArray.CreateBuilder<PerAnalyzerState>();
             var map = ImmutableDictionary.CreateBuilder<DiagnosticAnalyzer, int>();
             var index = 0;
             foreach (var analyzer in analyzers)
             {
-                statesBuilder.Add(new PerAnalyzerState(analyzerStateDataPool, declarationAnalyzerStateDataPool, currentlyAnalyzingDeclarationsMapPool));
+                statesBuilder.Add(s_perAnalyzerStatePool.Allocate());
                 map[analyzer] = index;
                 index++;
             }