moved cached metadata refs to instance state

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

@@ -14,23 +14,46 @@
 
 namespace Microsoft.CodeAnalysis
 {
- internal class MetadataOnlyReference
+ internal class CachedMetadataOnlyReferences
  {
- // version based cache
- private static readonly ConditionalWeakTable<ProjectId, MetadataOnlyReferenceSet> s_projectIdToReferenceMap = new();
-
- // snapshot based cache
+ /// <summary>
+ /// Mapping from compilation instance to metadata-references for it. Safe to use as a static CWT as anyone 
+ /// with a reference to this compilation (and the same <see cref="MetadataReferenceProperties"/> is allowed
+ /// to reference it using the same <see cref="MetadataReference"/>.
+ /// </summary>
  private static readonly ConditionalWeakTable<Compilation, MetadataOnlyReferenceSet> s_compilationToReferenceMap = new();
 
- internal static MetadataReference GetOrBuildReference(
+ private readonly object _gate = new();
+
+ /// <summary>
+ /// Mapping from <see cref="ProjectId"/> to the skeleton <see cref="MetadataReference"/>s for it.
+ /// </summary>
+ private ImmutableDictionary<ProjectId, MetadataOnlyReferenceSet> _projectIdToReferenceMap;
+
+ public static readonly CachedMetadataOnlyReferences Empty = new(ImmutableDictionary<ProjectId, MetadataOnlyReferenceSet>.Empty);
+
+ private CachedMetadataOnlyReferences(
+ ImmutableDictionary<ProjectId, MetadataOnlyReferenceSet> projectIdToReferenceMap)
+ {
+ _projectIdToReferenceMap = projectIdToReferenceMap;
+ }
+
+ public CachedMetadataOnlyReferences Clone()
+ => new(_projectIdToReferenceMap);
+
+ internal MetadataReference GetOrBuildReference(
  SolutionState solution,
  ProjectReference projectReference,
  Compilation finalCompilation,
  VersionStamp version,
  CancellationToken cancellationToken)
  {
+ // First see if we already have a cached reference for either finalCompilation or for projectReference.
+ // If we have one for the latter, we'll make sure that it's version matches what we're asking for before
+ // returning it.
  solution.Workspace.LogTestMessage($"Looking to see if we already have a skeleton assembly for {projectReference.ProjectId} before we build one...");
- if (TryGetReference(solution, projectReference, finalCompilation, version, out var reference))
+ if (TryGetReference(
+ solution, projectReference, finalCompilation, version, cancellationToken, out var reference))
  {
  solution.Workspace.LogTestMessage($"A reference was found {projectReference.ProjectId} so we're skipping the build.");
  return reference;
@@ -45,23 +68,16 @@ internal static MetadataReference GetOrBuildReference(
 
  if (image.IsEmpty)
  {
- // unfortunately, we couldn't create one. do best effort
- if (TryGetReference(solution, projectReference, finalCompilation, VersionStamp.Default, out reference))
+ // unfortunately, we couldn't create one. see if we have one from previous compilation., it might be
+ // out-of-date big time, but better than nothing.
+ if (TryGetReference(
+ solution, projectReference, finalCompilation, version: null, cancellationToken, out reference))
  {
  solution.Workspace.LogTestMessage($"We failed to create metadata so we're using the one we just found from an earlier version.");
-
- // we have one from previous compilation!!, it might be out-of-date big time, but better than nothing.
- // re-use it
  return reference;
  }
  }
 
- // okay, proceed with whatever image we have
-
- // now, remove existing set
- s_projectIdToReferenceMap.Remove(projectReference.ProjectId);
-
- // create new one
  var newReferenceSet = new MetadataOnlyReferenceSet(version, image);
  var referenceSet = s_compilationToReferenceMap.GetValue(finalCompilation, _ => newReferenceSet);
  if (newReferenceSet != referenceSet)
@@ -71,7 +87,8 @@ internal static MetadataReference GetOrBuildReference(
  image.Cleanup();
 
  // return new reference
- return referenceSet.GetMetadataReference(finalCompilation, projectReference.Aliases, projectReference.EmbedInteropTypes);
+ return referenceSet.GetMetadataReference(
+ finalCompilation, projectReference.Aliases, projectReference.EmbedInteropTypes, cancellationToken);
  }
  else
  {
@@ -80,34 +97,61 @@ internal static MetadataReference GetOrBuildReference(
 
  // record it to version based cache as well. snapshot cache always has a higher priority. we don't need to check returned set here
  // since snapshot based cache will take care of same compilation for us.
-#if NETCOREAPP
- s_projectIdToReferenceMap.AddOrUpdate(projectReference.ProjectId, referenceSet);
-#else
- lock (s_projectIdToReferenceMap)
+
+ lock (_gate)
  {
- s_projectIdToReferenceMap.Remove(projectReference.ProjectId);
- s_projectIdToReferenceMap.Add(projectReference.ProjectId, referenceSet);
+ _projectIdToReferenceMap = _projectIdToReferenceMap.Remove(projectReference.ProjectId);
+ _projectIdToReferenceMap = _projectIdToReferenceMap.Add(projectReference.ProjectId, referenceSet);
  }
-#endif
 
  // return new reference
- return referenceSet.GetMetadataReference(finalCompilation, projectReference.Aliases, projectReference.EmbedInteropTypes);
+ return referenceSet.GetMetadataReference(
+ finalCompilation, projectReference.Aliases, projectReference.EmbedInteropTypes, cancellationToken);
  }
 
- internal static bool TryGetReference(
- SolutionState solution, ProjectReference projectReference, Compilation finalOrDeclarationCompilation, VersionStamp version, out MetadataReference reference)
+ /// <summary>
+ /// Tries to get the <see cref="MetadataReference"/> associated with the provided <paramref name="projectReference"/>
+ /// that produced the <see cref="Compilation"/> <paramref name="finalOrDeclarationCompilation"/>.
+ /// </summary>
+ internal bool TryGetReference(
+ SolutionState solution,
+ ProjectReference projectReference,
+ Compilation finalOrDeclarationCompilation,
+ VersionStamp version,
+ CancellationToken cancellationToken,
+ out MetadataReference reference)
+ {
+ return TryGetReference(solution, projectReference, finalOrDeclarationCompilation, (VersionStamp?)version, cancellationToken, out reference);
+ }
+
+ /// <remarks>
+ /// <inheritdoc cref="TryGetReference(SolutionState, ProjectReference, Compilation, VersionStamp, CancellationToken, out MetadataReference)"/>.
+ /// If <paramref name="version"/> is <see langword="null"/>, any <see cref="MetadataReference"/> for that
+ /// project may be returned, even if it doesn't correspond to that compilation. This is useful in error tolerance
+ /// cases as building a skeleton assembly may easily fail. In that case it's better to use the last successfully
+ /// built skeleton than just have no semantic information for that project at all.
+ /// </remarks> 
+ private bool TryGetReference(
+ SolutionState solution,
+ ProjectReference projectReference,
+ Compilation finalOrDeclarationCompilation,
+ VersionStamp? version,
+ CancellationToken cancellationToken,
+ out MetadataReference reference)
  {
  // if we have one from snapshot cache, use it. it will make sure same compilation will get same metadata reference always.
  if (s_compilationToReferenceMap.TryGetValue(finalOrDeclarationCompilation, out var referenceSet))
  {
  solution.Workspace.LogTestMessage($"Found already cached metadata in {nameof(s_compilationToReferenceMap)} for the exact compilation");
- reference = referenceSet.GetMetadataReference(finalOrDeclarationCompilation, projectReference.Aliases, projectReference.EmbedInteropTypes);
+ reference = referenceSet.GetMetadataReference(
+ finalOrDeclarationCompilation, projectReference.Aliases, projectReference.EmbedInteropTypes, cancellationToken);
  return true;
  }
 
  // okay, now use version based cache that can live multiple compilation as long as there is no semantic changes.
 
- if (TryGetReference(projectReference, finalOrDeclarationCompilation, version, out reference))
+ if (TryGetReference(
+ projectReference, finalOrDeclarationCompilation, version, cancellationToken, out reference))
  {
  solution.Workspace.LogTestMessage($"Found already cached metadata for the branch and version {version}");
  return true;
@@ -118,16 +162,20 @@ internal static bool TryGetReference(
  return false;
  }
 
- private static bool TryGetReference(
- ProjectReference projectReference, Compilation finalOrDeclarationCompilation, VersionStamp version, out MetadataReference reference)
+ private bool TryGetReference(
+ ProjectReference projectReference,
+ Compilation finalOrDeclarationCompilation,
+ VersionStamp? version,
+ CancellationToken cancellationToken,
+ out MetadataReference reference)
  {
- if (s_projectIdToReferenceMap.TryGetValue(projectReference.ProjectId, out var referenceSet) &&
- (version == VersionStamp.Default || referenceSet.Version == version))
+ if (_projectIdToReferenceMap.TryGetValue(projectReference.ProjectId, out var referenceSet) &&
+ (version == null || referenceSet.Version == version))
  {
  // record it to snapshot based cache.
  var newReferenceSet = s_compilationToReferenceMap.GetValue(finalOrDeclarationCompilation, _ => referenceSet);
-
- reference = newReferenceSet.GetMetadataReference(finalOrDeclarationCompilation, projectReference.Aliases, projectReference.EmbedInteropTypes);
+ reference = newReferenceSet.GetMetadataReference(
+  finalOrDeclarationCompilation, projectReference.Aliases, projectReference.EmbedInteropTypes, cancellationToken);
  return true;
  }
 
@@ -140,31 +188,26 @@ private class MetadataOnlyReferenceSet
  // use WeakReference so we don't keep MetadataReference's alive if they are not being consumed
  private readonly NonReentrantLock _gate = new(useThisInstanceForSynchronization: true);
 
- // here, there is a very small chance of leaking Tuple and WeakReference, but it is so small chance,
- // I don't believe it will actually happen in real life situation. basically, for leak to happen, 
- // every image creation except the first one has to fail so that we end up re-use old reference set.
- // and the user creates many different metadata references with multiple combination of the key (tuple).
- private readonly Dictionary<MetadataReferenceProperties, WeakReference<MetadataReference>> _metadataReferences
- = new();
-
- private readonly VersionStamp _version;
+ private readonly Dictionary<MetadataReferenceProperties, WeakReference<MetadataReference>> _metadataReferences = new();
  private readonly MetadataOnlyImage _image;
 
+ public readonly VersionStamp Version;
+
  public MetadataOnlyReferenceSet(VersionStamp version, MetadataOnlyImage image)
  {
- _version = version;
+ Version = version;
  _image = image;
  }
 
- public VersionStamp Version => _version;
-
- public MetadataReference GetMetadataReference(Compilation compilation, ImmutableArray<string> aliases, bool embedInteropTypes)
+ public MetadataReference GetMetadataReference(
+ Compilation compilation, ImmutableArray<string> aliases, bool embedInteropTypes, CancellationToken cancellationToken)
  {
  var key = new MetadataReferenceProperties(MetadataImageKind.Assembly, aliases, embedInteropTypes);
 
- using (_gate.DisposableWait())
+ using (_gate.DisposableWait(cancellationToken))
  {
- if (!_metadataReferences.TryGetValue(key, out var weakMetadata) || !weakMetadata.TryGetTarget(out var metadataReference))
+ if (!_metadataReferences.TryGetValue(key, out var weakMetadata) ||
+ !weakMetadata.TryGetTarget(out var metadataReference))
  {
  // here we give out strong reference to compilation. so there is possibility that we end up making 2 compilations for same project alive.
  // one for final compilation and one for declaration only compilation. but the final compilation will be eventually kicked out from compilation cache

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

@@ -44,22 +44,26 @@ private partial class CompilationTracker : ICompilationTracker
  // guarantees only one thread is building at a time
  private readonly SemaphoreSlim _buildLock = new(initialCount: 1);
 
+ public CachedMetadataOnlyReferences CachedMetadataOnlyReferences { get; }
+
  private CompilationTracker(
  ProjectState project,
- CompilationTrackerState state)
+ CompilationTrackerState state,
+ CachedMetadataOnlyReferences cachedmetadataOnlyReferences)
  {
  Contract.ThrowIfNull(project);
 
  this.ProjectState = project;
  _stateDoNotAccessDirectly = state;
+ CachedMetadataOnlyReferences = cachedmetadataOnlyReferences;
  }
 
  /// <summary>
  /// Creates a tracker for the provided project. The tracker will be in the 'empty' state
  /// and will have no extra information beyond the project itself.
  /// </summary>
  public CompilationTracker(ProjectState project)
- : this(project, CompilationTrackerState.Empty)
+ : this(project, CompilationTrackerState.Empty, CachedMetadataOnlyReferences.Empty)
  {
  }
 
@@ -143,13 +147,13 @@ public ICompilationTracker Fork(
  var newState = CompilationTrackerState.Create(
  solutionServices, baseCompilation, state.GeneratorInfo, state.FinalCompilationWithGeneratedDocuments?.GetValueOrNull(cancellationToken), intermediateProjects);
 
- return new CompilationTracker(newProject, newState);
+ return new CompilationTracker(newProject, newState, this.CachedMetadataOnlyReferences.Clone());
  }
  else
  {
  // We have no compilation, but we might have information about generated docs.
  var newState = new NoCompilationState(state.GeneratorInfo.WithDocumentsAreFinal(false));
- return new CompilationTracker(newProject, newState);
+ return new CompilationTracker(newProject, newState, this.CachedMetadataOnlyReferences.Clone());
  }
  }
 
@@ -193,7 +197,7 @@ public ICompilationTracker FreezePartialStateWithTree(SolutionState solution, Do
  this.ProjectState.Id,
  metadataReferenceToProjectId);
 
- return new CompilationTracker(inProgressProject, finalState);
+ return new CompilationTracker(inProgressProject, finalState, this.CachedMetadataOnlyReferences.Clone());
  }
 
  /// <summary>
@@ -1012,7 +1016,8 @@ private async Task<MetadataReference> GetMetadataOnlyImageReferenceAsync(
  var declarationCompilation = await this.GetOrBuildDeclarationCompilationAsync(solution.Services, cancellationToken: cancellationToken).ConfigureAwait(false);
  solution.Workspace.LogTestMessage($"Looking for a cached skeleton assembly for {projectReference.ProjectId} before taking the lock...");
 
- if (!MetadataOnlyReference.TryGetReference(solution, projectReference, declarationCompilation, version, out var reference))
+ if (!this.CachedMetadataOnlyReferences.TryGetReference(
+ solution, projectReference, declarationCompilation, version, cancellationToken, out var reference))
  {
  // using async build lock so we don't get multiple consumers attempting to build metadata-only images for the same compilation.
  using (await _buildLock.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
@@ -1021,7 +1026,8 @@ private async Task<MetadataReference> GetMetadataOnlyImageReferenceAsync(
 
  // okay, we still don't have one. bring the compilation to final state since we are going to use it to create skeleton assembly
  var compilationInfo = await this.GetOrBuildCompilationInfoAsync(solution, lockGate: false, cancellationToken: cancellationToken).ConfigureAwait(false);
- reference = MetadataOnlyReference.GetOrBuildReference(solution, projectReference, compilationInfo.Compilation, version, cancellationToken);
+ reference = this.CachedMetadataOnlyReferences.GetOrBuildReference(
+ solution, projectReference, compilationInfo.Compilation, version, cancellationToken);
  }
  }
  else

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

@@ -38,6 +38,7 @@ public GeneratedFileReplacingCompilationTracker(ICompilationTracker underlyingTr
  }
 
  public ProjectState ProjectState => _underlyingTracker.ProjectState;
+ public CachedMetadataOnlyReferences CachedMetadataOnlyReferences => _underlyingTracker.CachedMetadataOnlyReferences;
 
  public bool ContainsAssemblyOrModuleOrDynamic(ISymbol symbol, bool primary)
  {
@@ -142,7 +143,7 @@ public async Task<MetadataReference> GetMetadataReferenceAsync(SolutionState sol
  else
  {
  var version = await GetDependentSemanticVersionAsync(solution, cancellationToken).ConfigureAwait(false);
- return MetadataOnlyReference.GetOrBuildReference(solution, projectReference, compilation, version, cancellationToken);
+ return _underlyingTracker.CachedMetadataOnlyReferences.GetOrBuildReference(solution, projectReference, compilation, version, cancellationToken);
  }
  }
 

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

@@ -13,6 +13,7 @@ internal partial class SolutionState
  private interface ICompilationTracker
  {
  ProjectState ProjectState { get; }
+ CachedMetadataOnlyReferences CachedMetadataOnlyReferences { get; }
 
  /// <summary>
  /// Returns <see langword="true"/> if this <see cref="Project"/>/<see cref="Compilation"/> could produce the