Merge pull request dotnet#74506 from dibarbet/cherry_pick_74189

Cherry pick project system workspace update fix to 17.11

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/HostWorkspace/MetadataService.cs

@@ -24,21 +24,19 @@ public IWorkspaceService CreateService(HostWorkspaceServices workspaceServices)
 
     internal sealed class MetadataService : IMetadataService
     {
-        private readonly IDocumentationProviderService _documentationProviderService;
+        private readonly MetadataReferenceCache _metadataCache;
 
         public MetadataService(IDocumentationProviderService documentationProviderService)
         {
-            _documentationProviderService = documentationProviderService;
+            _metadataCache = new MetadataReferenceCache((path, properties) =>
+                MetadataReference.CreateFromFile(path, properties, documentationProviderService.GetDocumentationProvider(path)));
         }
 
         public PortableExecutableReference GetReference(string resolvedPath, MetadataReferenceProperties properties)
         {
-            // HACK: right now the FileWatchedPortableExecutableReferenceFactory in Roslyn presumes that each time it calls IMetadataService
-            // it gets a unique instance back; the default MetadataService at the workspace layer has a cache to encourage sharing, but that
-            // breaks the assumption.
             try
             {
-                return MetadataReference.CreateFromFile(resolvedPath, properties, _documentationProviderService.GetDocumentationProvider(resolvedPath));
+                return (PortableExecutableReference)_metadataCache.GetReference(resolvedPath, properties);
             }
             catch (IOException ex)
             {

src/Workspaces/Core/Portable/Workspace/Host/Metadata/MetadataServiceFactory.cs

@@ -20,9 +20,9 @@ public MetadataServiceFactory()
     }
 
     public IWorkspaceService CreateService(HostWorkspaceServices workspaceServices)
-        => new Service(workspaceServices.GetService<IDocumentationProviderService>());
+        => new DefaultMetadataService(workspaceServices.GetService<IDocumentationProviderService>());
 
-    private sealed class Service(IDocumentationProviderService documentationService) : IMetadataService
+    private sealed class DefaultMetadataService(IDocumentationProviderService documentationService) : IMetadataService
     {
         private readonly MetadataReferenceCache _metadataCache = new MetadataReferenceCache((path, properties) =>
                 MetadataReference.CreateFromFile(path, properties, documentationService.GetDocumentationProvider(path)));

src/Workspaces/Core/Portable/Workspace/ProjectSystem/FileWatchedPortableExecutableReferenceFactory.cs

@@ -19,8 +19,6 @@ internal sealed class FileWatchedPortableExecutableReferenceFactory
 {
     private readonly object _gate = new();
 
-    private readonly SolutionServices _solutionServices;
-
     /// <summary>
     /// A file change context used to watch metadata references. This is lazy to avoid creating this immediately during our LSP process startup, when we
     /// don't yet know the LSP client's capabilities.
@@ -31,7 +29,7 @@ internal sealed class FileWatchedPortableExecutableReferenceFactory
     /// File watching tokens from <see cref="_fileReferenceChangeContext"/> that are watching metadata references. These are only created once we are actually applying a batch because
     /// we don't determine until the batch is applied if the file reference will actually be a file reference or it'll be a converted project reference.
     /// </summary>
-    private readonly Dictionary<PortableExecutableReference, IWatchedFile> _metadataReferenceFileWatchingTokens = [];
+    private readonly Dictionary<PortableExecutableReference, (IWatchedFile Token, int RefCount)> _metadataReferenceFileWatchingTokens = [];
 
     /// <summary>
     /// Stores the caller for a previous disposal of a reference produced by this class, to track down a double-dispose issue.
@@ -49,11 +47,8 @@ internal sealed class FileWatchedPortableExecutableReferenceFactory
     private readonly Dictionary<string, CancellationTokenSource> _metadataReferenceRefreshCancellationTokenSources = [];
 
     public FileWatchedPortableExecutableReferenceFactory(
-        SolutionServices solutionServices,
         IFileChangeWatcher fileChangeWatcher)
     {
-        _solutionServices = solutionServices;
-
         _fileReferenceChangeContext = new Lazy<IFileChangeContext>(() =>
         {
             var referenceDirectories = new HashSet<string>();
@@ -101,34 +96,56 @@ public FileWatchedPortableExecutableReferenceFactory(
 
     public event EventHandler<string>? ReferenceChanged;
 
-    public PortableExecutableReference CreateReferenceAndStartWatchingFile(string fullFilePath, MetadataReferenceProperties properties)
+    /// <summary>
+    /// Starts watching a particular PortableExecutableReference for changes to the file.
+    /// If this is already being watched , the reference count will be incremented.
+    /// This is *not* safe to attempt to call multiple times for the same project and reference (e.g. in applying workspace updates)
+    /// </summary>
+    public void StartWatchingReference(PortableExecutableReference reference, string fullFilePath)
     {
         lock (_gate)
         {
-            var reference = _solutionServices.GetRequiredService<IMetadataService>().GetReference(fullFilePath, properties);
-            var fileWatchingToken = _fileReferenceChangeContext.Value.EnqueueWatchingFile(fullFilePath);
-
-            _metadataReferenceFileWatchingTokens.Add(reference, fileWatchingToken);
-
-            return reference;
+            var (token, count) = _metadataReferenceFileWatchingTokens.GetOrAdd(reference, _ =>
+            {
+                var fileToken = _fileReferenceChangeContext.Value.EnqueueWatchingFile(fullFilePath);
+                return (fileToken, RefCount: 0);
+            });
+            _metadataReferenceFileWatchingTokens[reference] = (token, RefCount: count + 1);
         }
     }
 
+    /// <summary>
+    /// Decrements the reference count for the given PortableExecutableReference.
+    /// When the reference count reaches 0, the file watcher will be stopped.
+    /// This is *not* safe to attempt to call multiple times for the same project and reference (e.g. in applying workspace updates)
+    /// </summary>
     public void StopWatchingReference(PortableExecutableReference reference, [CallerFilePath] string callerFilePath = "", [CallerLineNumber] int callerLineNumber = 0)
     {
         lock (_gate)
         {
             var disposalLocation = callerFilePath + ", line " + callerLineNumber;
-
-            if (!_metadataReferenceFileWatchingTokens.TryGetValue(reference, out var watchedFile))
+            if (!_metadataReferenceFileWatchingTokens.TryGetValue(reference, out var watchedFileReference))
             {
+                // We're attempting to stop watching a file that we never started watching. This is a bug.
                 var existingDisposalStackTrace = _previousDisposalLocations.TryGetValue(reference, out var previousDisposalLocation);
                 throw new ArgumentException("The reference was already disposed at " + previousDisposalLocation);
             }
 
-            watchedFile.Dispose();
-            _metadataReferenceFileWatchingTokens.Remove(reference);
-            _previousDisposalLocations.Add(reference, disposalLocation);
+            var newRefCount = watchedFileReference.RefCount - 1;
+            Contract.ThrowIfFalse(newRefCount >= 0, "Ref count cannot be negative");
+            if (newRefCount == 0)
+            {
+                // No one else is watching this file, so stop watching it and remove from our map.
+                watchedFileReference.Token.Dispose();
+                _metadataReferenceFileWatchingTokens.Remove(reference);
+
+                _previousDisposalLocations.Remove(reference);
+                _previousDisposalLocations.Add(reference, disposalLocation);
+            }
+            else
+            {
+                _metadataReferenceFileWatchingTokens[reference] = (watchedFileReference.Token, newRefCount);
+            }
 
             // Note we still potentially have an outstanding change that we haven't raised a notification
             // for due to the delay we use. We could cancel the notification for that file path,

src/Workspaces/Core/Portable/Workspace/ProjectSystem/ProjectSystemProject.BatchingDocumentCollection.cs

@@ -416,7 +416,7 @@ public async ValueTask ProcessRegularFileChangesAsync(ImmutableSegmentedList<str
                     return;
                 }
 
-                await _project._projectSystemProjectFactory.ApplyBatchChangeToWorkspaceAsync(solutionChanges =>
+                await _project._projectSystemProjectFactory.ApplyBatchChangeToWorkspaceAsync((solutionChanges, projectUpdateState) =>
                 {
                     foreach (var (documentId, textLoader) in documentsToChange)
                     {
@@ -428,7 +428,8 @@ await _project._projectSystemProjectFactory.ApplyBatchChangeToWorkspaceAsync(sol
                                 [documentId]);
                         }
                     }
-                }).ConfigureAwait(false);
+                    return projectUpdateState;
+                }, onAfterUpdateAlways: null).ConfigureAwait(false);
 
                 documentsToChange.Free();
             }
@@ -528,51 +529,83 @@ public void ReorderFiles(ImmutableArray<string> filePaths)
             }
         }
 
-        internal void UpdateSolutionForBatch(
+        /// <summary>
+        /// Updates the solution for a set of batch changes.
+        /// While it is OK for this method to *read* local state, it cannot *modify* it as this may
+        /// be called multiple times (when the workspace update fails due to interceding updates).
+        /// </summary>
+        internal ImmutableArray<(DocumentId documentId, SourceTextContainer textContainer)> UpdateSolutionForBatch(
             SolutionChangeAccumulator solutionChanges,
             ImmutableArray<string>.Builder documentFileNamesAdded,
-            List<(DocumentId documentId, SourceTextContainer textContainer)> documentsToOpen,
             Func<Solution, ImmutableArray<DocumentInfo>, Solution> addDocuments,
             WorkspaceChangeKind addDocumentChangeKind,
             Func<Solution, ImmutableArray<DocumentId>, Solution> removeDocuments,
             WorkspaceChangeKind removeDocumentChangeKind)
         {
-            // Document adding...
-            solutionChanges.UpdateSolutionForDocumentAction(
-                newSolution: addDocuments(solutionChanges.Solution, _documentsAddedInBatch.ToImmutable()),
-                changeKind: addDocumentChangeKind,
-                documentIds: _documentsAddedInBatch.Select(d => d.Id));
-
-            foreach (var documentInfo in _documentsAddedInBatch)
+            // Intentionally making copies to pass into the static update function.
+            // State is cleared at the end once the solution changes are actually applied via ClearBatchState.
+            return UpdateSolutionForBatch(solutionChanges, documentFileNamesAdded, addDocuments,
+                addDocumentChangeKind, removeDocuments, removeDocumentChangeKind, _project.Id, _documentsAddedInBatch.ToImmutableArray(),
+                _documentsRemovedInBatch.ToImmutableArray(), _orderedDocumentsInBatch,
+                documentId => _sourceTextContainersToDocumentIds.GetKeyOrDefault(documentId));
+
+            static ImmutableArray<(DocumentId documentId, SourceTextContainer textContainer)> UpdateSolutionForBatch(
+                SolutionChangeAccumulator solutionChanges,
+                ImmutableArray<string>.Builder documentFileNamesAdded,
+                Func<Solution, ImmutableArray<DocumentInfo>, Solution> addDocuments,
+                WorkspaceChangeKind addDocumentChangeKind,
+                Func<Solution, ImmutableArray<DocumentId>, Solution> removeDocuments,
+                WorkspaceChangeKind removeDocumentChangeKind,
+                ProjectId projectId,
+                ImmutableArray<DocumentInfo> documentsAddedInBatch,
+                ImmutableArray<DocumentId> documentsRemovedInBatch,
+                ImmutableList<DocumentId>? orderedDocumentsInBatch,
+                Func<DocumentId, SourceTextContainer?> getContainer)
             {
-                Contract.ThrowIfNull(documentInfo.FilePath, "We shouldn't be adding documents without file paths.");
-                documentFileNamesAdded.Add(documentInfo.FilePath);
+                using var _ = ArrayBuilder<(DocumentId documentId, SourceTextContainer textContainer)>.GetInstance(out var documentsToOpen);
 
-                if (_sourceTextContainersToDocumentIds.TryGetKey(documentInfo.Id, out var textContainer))
+                // Document adding...
+                solutionChanges.UpdateSolutionForDocumentAction(
+                    newSolution: addDocuments(solutionChanges.Solution, documentsAddedInBatch),
+                    changeKind: addDocumentChangeKind,
+                    documentIds: documentsAddedInBatch.Select(d => d.Id));
+
+                foreach (var documentInfo in documentsAddedInBatch)
                 {
-                    documentsToOpen.Add((documentInfo.Id, textContainer));
-                }
-            }
+                    Contract.ThrowIfNull(documentInfo.FilePath, "We shouldn't be adding documents without file paths.");
+                    documentFileNamesAdded.Add(documentInfo.FilePath);
 
-            ClearAndZeroCapacity(_documentsAddedInBatch);
+                    var textContainer = getContainer(documentInfo.Id);
+                    if (textContainer != null)
+                    {
+                        documentsToOpen.Add((documentInfo.Id, textContainer));
+                    }
+                }
 
-            // Document removing...
-            solutionChanges.UpdateSolutionForRemovedDocumentAction(removeDocuments(solutionChanges.Solution, [.. _documentsRemovedInBatch]),
-                removeDocumentChangeKind,
-                _documentsRemovedInBatch);
+                // Document removing...
+                solutionChanges.UpdateSolutionForRemovedDocumentAction(removeDocuments(solutionChanges.Solution, documentsRemovedInBatch),
+                    removeDocumentChangeKind,
+                    documentsRemovedInBatch);
 
-            ClearAndZeroCapacity(_documentsRemovedInBatch);
+                // Update project's order of documents.
+                if (orderedDocumentsInBatch != null)
+                {
+                    solutionChanges.UpdateSolutionForProjectAction(
+                        projectId,
+                        solutionChanges.Solution.WithProjectDocumentsOrder(projectId, orderedDocumentsInBatch));
+                }
 
-            // Update project's order of documents.
-            if (_orderedDocumentsInBatch != null)
-            {
-                solutionChanges.UpdateSolutionForProjectAction(
-                    _project.Id,
-                    solutionChanges.Solution.WithProjectDocumentsOrder(_project.Id, _orderedDocumentsInBatch));
-                _orderedDocumentsInBatch = null;
+                return documentsToOpen.ToImmutable();
             }
         }
 
+        internal void ClearBatchState()
+        {
+            ClearAndZeroCapacity(_documentsAddedInBatch);
+            ClearAndZeroCapacity(_documentsRemovedInBatch);
+            _orderedDocumentsInBatch = null;
+        }
+
         private DocumentInfo CreateDocumentInfoFromFileInfo(DynamicFileInfo fileInfo, ImmutableArray<string> folders)
         {
             Contract.ThrowIfTrue(folders.IsDefault);