Switch Project data structures from ImmutableDictionary => Dictionary and lock (#78287)

Solution/Project classes keep several immutable dictionaries for mapping (Document/Project)Id to various data. ImmutableDictionaries are a fine (but slow) data structure when data between versions is to be shared, but none of the values in these ImmutableDictionaries persist between solution/project transformations. In fact, the only operation performed on these dictionaries is GetOrAdd calls. As immutable dictionaries are quite poor performing wrt lookups/adds, I thought I would try out two approaches to see if they showed improvements:

Switch to ConcurrentDictionary (see *** WIP: Switch Project data structures from ImmutableDictionary => ConcurrentDictionary #78285).
Switch to simple Dictionaries guarded by locks (this PR)
I created test insertions for each of these PRs to gather speedometer numbers, particularly during the solution load of the c# editing speedometer test. The measurements below are from GC Heap Allocs and CPU samples for those tests under (Solution/Project).GetDocument calls

ConcurrentDictionary change test insertion: https://dev.azure.com/devdiv/DevDiv/_git/VS/pullrequest/631277
Dictionary/lock test insertion: https://dev.azure.com/devdiv/DevDiv/_git/VS/pullrequest/631278
*** Baseline ***
Alloc: 120 MB (1.0%)
CPU: 560 ms (0.3%)

*** Concurrent Dictionary ***
Alloc: 59 MB
CPU: 258 ms (0.1%)

*** Dictionary with lock ***
Alloc: 28 MB
CPU: 172 ms (< 0.1%)

So, it looks like the best performing of the bunch is this PR, which uses simple dictionaries and locks, thus why I'm elevating this PR out of draft mode. The ConcurrentDictionary PR is interesting as the code changes are simple and do give a nice perf improvement.

Author: ToddGrun

src/Workspaces/Core/Portable/Workspace/Solution/Project.cs

@@ -26,10 +26,11 @@ namespace Microsoft.CodeAnalysis;
 [DebuggerDisplay("{GetDebuggerDisplay(),nq}")]
 public partial class Project
 {
-    private ImmutableDictionary<DocumentId, Document?> _idToDocumentMap = ImmutableDictionary<DocumentId, Document?>.Empty;
-    private ImmutableDictionary<DocumentId, SourceGeneratedDocument> _idToSourceGeneratedDocumentMap = ImmutableDictionary<DocumentId, SourceGeneratedDocument>.Empty;
-    private ImmutableDictionary<DocumentId, AdditionalDocument?> _idToAdditionalDocumentMap = ImmutableDictionary<DocumentId, AdditionalDocument?>.Empty;
-    private ImmutableDictionary<DocumentId, AnalyzerConfigDocument?> _idToAnalyzerConfigDocumentMap = ImmutableDictionary<DocumentId, AnalyzerConfigDocument?>.Empty;
+    // Only access these dictionaries when holding them as a lock
+    private Dictionary<DocumentId, Document?>? _idToDocumentMap;
+    private Dictionary<DocumentId, SourceGeneratedDocument>? _idToSourceGeneratedDocumentMap;
+    private Dictionary<DocumentId, AdditionalDocument?>? _idToAdditionalDocumentMap;
+    private Dictionary<DocumentId, AnalyzerConfigDocument?>? _idToAnalyzerConfigDocumentMap;
 
     internal Project(Solution solution, ProjectState projectState)
     {
@@ -239,19 +240,48 @@ public bool ContainsAnalyzerConfigDocument(DocumentId documentId)
     /// Get the document in this project with the specified document Id.
     /// </summary>
     public Document? GetDocument(DocumentId documentId)
-        => ImmutableInterlocked.GetOrAdd(ref _idToDocumentMap, documentId, s_tryCreateDocumentFunction, this);
+        => GetOrAddDocumentUnderLock(documentId, ref _idToDocumentMap, s_tryCreateDocumentFunction, this);
 
     /// <summary>
     /// Get the additional document in this project with the specified document Id.
     /// </summary>
     public TextDocument? GetAdditionalDocument(DocumentId documentId)
-        => ImmutableInterlocked.GetOrAdd(ref _idToAdditionalDocumentMap, documentId, s_tryCreateAdditionalDocumentFunction, this);
+        => GetOrAddDocumentUnderLock(documentId, ref _idToAdditionalDocumentMap, s_tryCreateAdditionalDocumentFunction, this);
 
     /// <summary>
     /// Get the analyzer config document in this project with the specified document Id.
     /// </summary>
     public AnalyzerConfigDocument? GetAnalyzerConfigDocument(DocumentId documentId)
-        => ImmutableInterlocked.GetOrAdd(ref _idToAnalyzerConfigDocumentMap, documentId, s_tryCreateAnalyzerConfigDocumentFunction, this);
+        => GetOrAddDocumentUnderLock(documentId, ref _idToAnalyzerConfigDocumentMap, s_tryCreateAnalyzerConfigDocumentFunction, this);
+
+    private static TDocument GetOrAddDocumentUnderLock<TDocument, TArg>(DocumentId documentId, ref Dictionary<DocumentId, TDocument>? idMap, Func<DocumentId, TArg, TDocument> tryCreate, TArg arg)
+    {
+        if (idMap == null)
+        {
+            // First call assigns a new dictionary. Any other simultaneous requests will not assign the
+            // dictionary they created to idMap. At that point, normal locking rules apply.
+            Interlocked.CompareExchange(ref idMap, new Dictionary<DocumentId, TDocument>(), null);
+        }
+
+        lock (idMap)
+        {
+            return idMap.GetOrAdd(documentId, tryCreate, arg);
+        }
+    }
+
+    private static bool TryGetDocumentValueUnderLock<TDocument>(DocumentId documentId, ref Dictionary<DocumentId, TDocument>? idMap, out TDocument? document)
+    {
+        if (idMap == null)
+        {
+            document = default;
+            return false;
+        }
+
+        lock (idMap)
+        {
+            return idMap.TryGetValue(documentId, out document);
+        }
+    }
 
     /// <summary>
     /// Gets a document or a source generated document in this solution with the specified document ID.
@@ -288,7 +318,7 @@ public async ValueTask<IEnumerable<SourceGeneratedDocument>> GetSourceGeneratedD
 
         // return an iterator to avoid eagerly allocating all the document instances
         return generatedDocumentStates.States.Values.Select(state =>
-            ImmutableInterlocked.GetOrAdd(ref _idToSourceGeneratedDocumentMap, state.Id, s_createSourceGeneratedDocumentFunction, (state, this)));
+            GetOrAddDocumentUnderLock(state.Id, ref _idToSourceGeneratedDocumentMap, s_createSourceGeneratedDocumentFunction, (state, this)));
     }
 
     internal async IAsyncEnumerable<Document> GetAllRegularAndSourceGeneratedDocumentsAsync([EnumeratorCancellation] CancellationToken cancellationToken)
@@ -312,7 +342,7 @@ internal async IAsyncEnumerable<Document> GetAllRegularAndSourceGeneratedDocumen
             return null;
 
         // Quick check first: if we already have created a SourceGeneratedDocument wrapper, we're good
-        if (_idToSourceGeneratedDocumentMap.TryGetValue(documentId, out var sourceGeneratedDocument))
+        if (TryGetDocumentValueUnderLock(documentId, ref _idToSourceGeneratedDocumentMap, out var sourceGeneratedDocument))
             return sourceGeneratedDocument;
 
         // We'll have to run generators if we haven't already and now try to find it.
@@ -325,7 +355,7 @@ internal async IAsyncEnumerable<Document> GetAllRegularAndSourceGeneratedDocumen
     }
 
     internal SourceGeneratedDocument GetOrCreateSourceGeneratedDocument(SourceGeneratedDocumentState state)
-        => ImmutableInterlocked.GetOrAdd(ref _idToSourceGeneratedDocumentMap, state.Id, s_createSourceGeneratedDocumentFunction, (state, this));
+        => GetOrAddDocumentUnderLock(state.Id, ref _idToSourceGeneratedDocumentMap, s_createSourceGeneratedDocumentFunction, (state, this));
 
     /// <summary>
     /// Returns the <see cref="SourceGeneratedDocumentState"/> for a source generated document that has already been generated and observed.
@@ -346,7 +376,7 @@ internal SourceGeneratedDocument GetOrCreateSourceGeneratedDocument(SourceGenera
             return null;
 
         // Easy case: do we already have the SourceGeneratedDocument created?
-        if (_idToSourceGeneratedDocumentMap.TryGetValue(documentId, out var document))
+        if (TryGetDocumentValueUnderLock(documentId, ref _idToSourceGeneratedDocumentMap, out var document))
             return document;
 
         // Trickier case now: it's possible we generated this, but we don't actually have the SourceGeneratedDocument for it, so let's go
@@ -355,7 +385,7 @@ internal SourceGeneratedDocument GetOrCreateSourceGeneratedDocument(SourceGenera
         if (documentState == null)
             return null;
 
-        return ImmutableInterlocked.GetOrAdd(ref _idToSourceGeneratedDocumentMap, documentId, s_createSourceGeneratedDocumentFunction, (documentState, this));
+        return GetOrAddDocumentUnderLock(documentId, ref _idToSourceGeneratedDocumentMap, s_createSourceGeneratedDocumentFunction, (documentState, this));
     }
 
     internal ValueTask<ImmutableArray<Diagnostic>> GetSourceGeneratorDiagnosticsAsync(CancellationToken cancellationToken)

src/Workspaces/Core/Portable/Workspace/Solution/Solution.cs

@@ -27,9 +27,8 @@ namespace Microsoft.CodeAnalysis;
 /// </summary>
 public partial class Solution
 {
-
-    // Values for all these are created on demand.
-    private ImmutableDictionary<ProjectId, Project> _projectIdToProjectMap;
+    // Values for all these are created on demand. Only access when holding the dictionary as a lock.
+    private readonly Dictionary<ProjectId, Project> _projectIdToProjectMap = [];
 
     /// <summary>
     /// Result of calling <see cref="WithFrozenPartialCompilationsAsync"/>.
@@ -46,7 +45,6 @@ private Solution(
         SolutionCompilationState compilationState,
         AsyncLazy<Solution>? cachedFrozenSolution = null)
     {
-        _projectIdToProjectMap = ImmutableDictionary<ProjectId, Project>.Empty;
         CompilationState = compilationState;
 
         _cachedFrozenSolution = cachedFrozenSolution ??
@@ -152,7 +150,10 @@ public bool ContainsProject([NotNullWhen(returnValue: true)] ProjectId? projectI
     {
         if (this.ContainsProject(projectId))
         {
-            return ImmutableInterlocked.GetOrAdd(ref _projectIdToProjectMap, projectId, s_createProjectFunction, this);
+            lock (_projectIdToProjectMap)
+            {
+                return _projectIdToProjectMap.GetOrAdd(projectId, s_createProjectFunction, this);
+            }
         }
 
         return null;