diff --git a/src/Workspaces/Core/Portable/Workspace/Solution/Project.cs b/src/Workspaces/Core/Portable/Workspace/Solution/Project.cs index f8f4981ead216..d20005fd748bb 100644 --- a/src/Workspaces/Core/Portable/Workspace/Solution/Project.cs +++ b/src/Workspaces/Core/Portable/Workspace/Solution/Project.cs @@ -26,10 +26,11 @@ namespace Microsoft.CodeAnalysis; [DebuggerDisplay("{GetDebuggerDisplay(),nq}")] public partial class Project { - private ImmutableDictionary _idToDocumentMap = ImmutableDictionary.Empty; - private ImmutableDictionary _idToSourceGeneratedDocumentMap = ImmutableDictionary.Empty; - private ImmutableDictionary _idToAdditionalDocumentMap = ImmutableDictionary.Empty; - private ImmutableDictionary _idToAnalyzerConfigDocumentMap = ImmutableDictionary.Empty; + // Only access these dictionaries when holding them as a lock + private Dictionary? _idToDocumentMap; + private Dictionary? _idToSourceGeneratedDocumentMap; + private Dictionary? _idToAdditionalDocumentMap; + private Dictionary? _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. /// public Document? GetDocument(DocumentId documentId) - => ImmutableInterlocked.GetOrAdd(ref _idToDocumentMap, documentId, s_tryCreateDocumentFunction, this); + => GetOrAddDocumentUnderLock(documentId, ref _idToDocumentMap, s_tryCreateDocumentFunction, this); /// /// Get the additional document in this project with the specified document Id. /// public TextDocument? GetAdditionalDocument(DocumentId documentId) - => ImmutableInterlocked.GetOrAdd(ref _idToAdditionalDocumentMap, documentId, s_tryCreateAdditionalDocumentFunction, this); + => GetOrAddDocumentUnderLock(documentId, ref _idToAdditionalDocumentMap, s_tryCreateAdditionalDocumentFunction, this); /// /// Get the analyzer config document in this project with the specified document Id. /// public AnalyzerConfigDocument? GetAnalyzerConfigDocument(DocumentId documentId) - => ImmutableInterlocked.GetOrAdd(ref _idToAnalyzerConfigDocumentMap, documentId, s_tryCreateAnalyzerConfigDocumentFunction, this); + => GetOrAddDocumentUnderLock(documentId, ref _idToAnalyzerConfigDocumentMap, s_tryCreateAnalyzerConfigDocumentFunction, this); + + private static TDocument GetOrAddDocumentUnderLock(DocumentId documentId, ref Dictionary? idMap, Func 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(), null); + } + + lock (idMap) + { + return idMap.GetOrAdd(documentId, tryCreate, arg); + } + } + + private static bool TryGetDocumentValueUnderLock(DocumentId documentId, ref Dictionary? idMap, out TDocument? document) + { + if (idMap == null) + { + document = default; + return false; + } + + lock (idMap) + { + return idMap.TryGetValue(documentId, out document); + } + } /// /// Gets a document or a source generated document in this solution with the specified document ID. @@ -288,7 +318,7 @@ public async ValueTask> 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 GetAllRegularAndSourceGeneratedDocumentsAsync([EnumeratorCancellation] CancellationToken cancellationToken) @@ -312,7 +342,7 @@ internal async IAsyncEnumerable 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 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)); /// /// Returns the 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> GetSourceGeneratorDiagnosticsAsync(CancellationToken cancellationToken) diff --git a/src/Workspaces/Core/Portable/Workspace/Solution/Solution.cs b/src/Workspaces/Core/Portable/Workspace/Solution/Solution.cs index 0ad1c652735c5..20f5a637864e8 100644 --- a/src/Workspaces/Core/Portable/Workspace/Solution/Solution.cs +++ b/src/Workspaces/Core/Portable/Workspace/Solution/Solution.cs @@ -27,9 +27,8 @@ namespace Microsoft.CodeAnalysis; /// public partial class Solution { - - // Values for all these are created on demand. - private ImmutableDictionary _projectIdToProjectMap; + // Values for all these are created on demand. Only access when holding the dictionary as a lock. + private readonly Dictionary _projectIdToProjectMap = []; /// /// Result of calling . @@ -46,7 +45,6 @@ private Solution( SolutionCompilationState compilationState, AsyncLazy? cachedFrozenSolution = null) { - _projectIdToProjectMap = ImmutableDictionary.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;