Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Apr 25, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
084437c
*** WIP: Switch Project data structures from ImmutableDictionary => D…
ToddGrun Apr 23, 2025
014d868
move lock inside enumerator
ToddGrun Apr 23, 2025
9381844
cleanup
ToddGrun Apr 23, 2025
5677bc5
Change Solution._projectIdToProjectMap to a guarded dictionary too.
ToddGrun Apr 24, 2025
2b26fe6
1) Fix multi-threaded bug. Accessing dictionary outside of lock. Bad …
ToddGrun Apr 25, 2025
64c72c2
Be a bit more clear in the comment around usage of Interlocked.Compar…
ToddGrun Apr 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 42 additions & 12 deletions src/Workspaces/Core/Portable/Workspace/Solution/Project.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand Down Expand Up @@ -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);
}
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

small nit. consider extracting helper for this pattern. we often have triplicate methods for docs/analyzerdocs/additionaldocs. i lke it when they all call into a single helper.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Extracted both GetOrAdd and TryGetValue calls into their own methods that can be called for any of the dictionaries.


/// <summary>
/// Gets a document or a source generated document in this solution with the specified document ID.
Expand Down Expand Up @@ -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)
Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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)
Expand Down
11 changes: 6 additions & 5 deletions src/Workspaces/Core/Portable/Workspace/Solution/Solution.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [];
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since you do initialize this, consider having hte dict be its own lock.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

are you auditing all usages of _projectIdToProjectMap to make sure it is locked?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I did check all five dictionaries converted, just didn't notice the TryGetValue calls on _idToSourceGeneratedDocumentMap


/// <summary>
/// Result of calling <see cref="WithFrozenPartialCompilationsAsync"/>.
Expand All @@ -46,7 +45,6 @@ private Solution(
SolutionCompilationState compilationState,
AsyncLazy<Solution>? cachedFrozenSolution = null)
{
_projectIdToProjectMap = ImmutableDictionary<ProjectId, Project>.Empty;
CompilationState = compilationState;

_cachedFrozenSolution = cachedFrozenSolution ??
Expand Down Expand Up @@ -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;
Expand Down
Loading