Fix memory leaks

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/DocumentSnapshot.cs

@@ -7,15 +7,25 @@
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem.Legacy;
+using Microsoft.CodeAnalysis.Razor.ProjectSystem.Sources;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
 
-internal sealed class DocumentSnapshot(ProjectSnapshot project, DocumentState state) : IDocumentSnapshot, ILegacyDocumentSnapshot, IDesignTimeCodeGenerator
+internal sealed class DocumentSnapshot : IDocumentSnapshot, ILegacyDocumentSnapshot, IDesignTimeCodeGenerator
 {
-    public ProjectSnapshot Project { get; } = project;
+    private readonly GeneratedOutputSource _generatedOutputSource;
 
-    private readonly DocumentState _state = state;
+    public ProjectSnapshot Project { get; }
+
+    private readonly DocumentState _state;
+
+    public DocumentSnapshot(ProjectSnapshot project, DocumentState state)
+    {
+        Project = project;
+        _state = state;
+        _generatedOutputSource = new(this);
+    }
 
     public HostDocument HostDocument => _state.HostDocument;
 
@@ -40,10 +50,10 @@ public ValueTask<VersionStamp> GetTextVersionAsync(CancellationToken cancellatio
         => _state.GetTextVersionAsync(cancellationToken);
 
     public bool TryGetGeneratedOutput([NotNullWhen(true)] out RazorCodeDocument? result)
-        => _state.TryGetGeneratedOutput(out result);
+        => _generatedOutputSource.TryGetValue(out result);
 
     public ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(CancellationToken cancellationToken)
-        => _state.GetGeneratedOutputAsync(this, cancellationToken);
+        => _generatedOutputSource.GetValueAsync(cancellationToken);
 
     public IDocumentSnapshot WithText(SourceText text)
     {

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/DocumentState.cs

@@ -4,7 +4,6 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem.Sources;
 using Microsoft.CodeAnalysis.Text;
 
@@ -16,22 +15,19 @@ internal sealed partial class DocumentState
     public int Version { get; }
 
     private readonly ITextAndVersionSource _textAndVersionSource;
-    private readonly GeneratedOutputSource _generatedOutputSource;
 
     private DocumentState(HostDocument hostDocument, ITextAndVersionSource textAndVersionSource)
     {
         HostDocument = hostDocument;
         Version = 1;
         _textAndVersionSource = textAndVersionSource;
-        _generatedOutputSource = new();
     }
 
     private DocumentState(DocumentState oldState, ITextAndVersionSource textAndVersionSource)
     {
         HostDocument = oldState.HostDocument;
         Version = oldState.Version + 1;
         _textAndVersionSource = textAndVersionSource;
-        _generatedOutputSource = new();
     }
 
     public static DocumentState Create(HostDocument hostDocument, SourceText text)
@@ -46,12 +42,6 @@ private static ConstantTextAndVersionSource CreateTextAndVersionSource(SourceTex
     private static LoadableTextAndVersionSource CreateTextAndVersionSource(TextLoader textLoader)
         => new(textLoader);
 
-    public bool TryGetGeneratedOutput([NotNullWhen(true)] out RazorCodeDocument? result)
-        => _generatedOutputSource.TryGetValue(out result);
-
-    public ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(DocumentSnapshot document, CancellationToken cancellationToken)
-        => _generatedOutputSource.GetValueAsync(document, cancellationToken);
-
     public bool TryGetTextAndVersion([NotNullWhen(true)] out TextAndVersion? result)
         => _textAndVersionSource.TryGetValue(out result);
 
@@ -110,13 +100,7 @@ async ValueTask<VersionStamp> GetTextVersionCoreAsync(CancellationToken cancella
         }
     }
 
-    public DocumentState WithConfigurationChange()
-        => new(this, _textAndVersionSource);
-
-    public DocumentState WithImportsChange()
-        => new(this, _textAndVersionSource);
-
-    public DocumentState WithProjectWorkspaceStateChange()
+    public DocumentState UpdateVersion()
         => new(this, _textAndVersionSource);
 
     public DocumentState WithText(SourceText text, VersionStamp textVersion)

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/ProjectSnapshot.cs

@@ -20,7 +20,7 @@ internal sealed class ProjectSnapshot(ProjectState state) : IProjectSnapshot, IL
     private readonly ProjectState _state = state;
 
     private readonly object _gate = new();
-    private readonly Dictionary<string, DocumentSnapshot> _filePathToDocumentMap = new(FilePathNormalizingComparer.Instance);
+    private Dictionary<string, DocumentSnapshot>? _filePathToDocumentMap;
 
     public HostProject HostProject => _state.HostProject;
     public RazorCompilerOptions CompilerOptions => _state.CompilerOptions;
@@ -53,8 +53,12 @@ public bool ContainsDocument(string filePath)
             // ImmutableDictionary<,>, which has O(log n) lookup. So, checking _filePathToDocumentMap
             // first is faster if the DocumentSnapshot has already been created.
 
-            return _filePathToDocumentMap.ContainsKey(filePath) ||
-                   _state.Documents.ContainsKey(filePath);
+            if (_filePathToDocumentMap is not null && _filePathToDocumentMap.ContainsKey(filePath))
+            {
+                return true;
+            }
+
+            return _state.Documents.ContainsKey(filePath);
         }
     }
 
@@ -63,7 +67,8 @@ public bool TryGetDocument(string filePath, [NotNullWhen(true)] out DocumentSnap
         lock (_gate)
         {
             // Have we already seen this document? If so, return it!
-            if (_filePathToDocumentMap.TryGetValue(filePath, out var snapshot))
+            if (_filePathToDocumentMap is not null &&
+                _filePathToDocumentMap.TryGetValue(filePath, out var snapshot))
             {
                 document = snapshot;
                 return true;
@@ -78,6 +83,8 @@ public bool TryGetDocument(string filePath, [NotNullWhen(true)] out DocumentSnap
 
             // If we have DocumentState, go ahead and create a new DocumentSnapshot.
             snapshot = new DocumentSnapshot(this, state);
+
+            _filePathToDocumentMap ??= new(capacity: _state.Documents.Count, FilePathNormalizingComparer.Instance);
             _filePathToDocumentMap.Add(filePath, snapshot);
 
             document = snapshot;

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/ProjectState.cs

@@ -259,7 +259,7 @@ public ProjectState WithHostProject(HostProject hostProject)
             return this;
         }
 
-        var documents = UpdateDocuments(static x => x.WithConfigurationChange());
+        var documents = UpdateDocuments(static x => x.UpdateVersion());
 
         // If the host project has changed then we need to recompute the imports map
         var importsToRelatedDocuments = BuildImportsMap(documents.Values, ProjectEngine);
@@ -277,7 +277,7 @@ public ProjectState WithProjectWorkspaceState(ProjectWorkspaceState projectWorks
             return this;
         }
 
-        var documents = UpdateDocuments(static x => x.WithProjectWorkspaceStateChange());
+        var documents = UpdateDocuments(static x => x.UpdateVersion());
 
         return new(this, HostProject, projectWorkspaceState, documents, ImportsToRelatedDocuments, retainProjectEngine: true);
     }
@@ -378,7 +378,7 @@ private ImmutableDictionary<string, DocumentState> UpdateRelatedDocumentsIfNeces
             return documents;
         }
 
-        var updates = relatedDocuments.Select(x => KeyValuePair.Create(x, documents[x].WithImportsChange()));
+        var updates = relatedDocuments.Select(x => KeyValuePair.Create(x, documents[x].UpdateVersion()));
         return documents.SetItems(updates);
     }
 

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/Sources/GeneratedOutputSource.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
+using System;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using System.Threading.Tasks;
@@ -9,19 +10,29 @@
 
 namespace Microsoft.CodeAnalysis.Razor.ProjectSystem.Sources;
 
-internal sealed class GeneratedOutputSource
+internal sealed class GeneratedOutputSource(DocumentSnapshot document)
 {
+    private readonly DocumentSnapshot _document = document;
     private readonly SemaphoreSlim _gate = new(initialCount: 1);
 
-    private RazorCodeDocument? _output;
+    // Hold the output in a WeakReference to avoid memory leaks in the case of a long-lived
+    // document snapshots. In particular, the DynamicFileInfo system results in the Roslyn
+    // workspace holding onto document snapshots.
+    private WeakReference<RazorCodeDocument>? _output;
 
     public bool TryGetValue([NotNullWhen(true)] out RazorCodeDocument? result)
     {
-        result = _output;
-        return result is not null;
+        var output = _output;
+        if (output is null)
+        {
+            result = null;
+            return false;
+        }
+
+        return output.TryGetTarget(out result);
     }
 
-    public async ValueTask<RazorCodeDocument> GetValueAsync(DocumentSnapshot document, CancellationToken cancellationToken)
+    public async ValueTask<RazorCodeDocument> GetValueAsync(CancellationToken cancellationToken)
     {
         if (TryGetValue(out var result))
         {
@@ -35,15 +46,24 @@ public async ValueTask<RazorCodeDocument> GetValueAsync(DocumentSnapshot documen
                 return result;
             }
 
-            var project = document.Project;
+            var project = _document.Project;
             var projectEngine = project.ProjectEngine;
             var compilerOptions = project.CompilerOptions;
 
-            _output = await CompilationHelpers
-                .GenerateCodeDocumentAsync(document, projectEngine, compilerOptions, cancellationToken)
+            result = await CompilationHelpers
+                .GenerateCodeDocumentAsync(_document, projectEngine, compilerOptions, cancellationToken)
                 .ConfigureAwait(false);
 
-            return _output;
+            if (_output is null)
+            {
+                _output = new(result);
+            }
+            else
+            {
+                _output.SetTarget(result);
+            }
+
+            return result;
         }
     }
 }

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/DynamicFiles/BackgroundDocumentGenerator.cs

@@ -86,11 +86,11 @@ public void Dispose()
     protected Task WaitUntilCurrentBatchCompletesAsync()
         => _workQueue.WaitUntilCurrentBatchCompletesAsync();
 
-    protected virtual async Task ProcessDocumentAsync(DocumentSnapshot document, CancellationToken cancellationToken)
+    protected virtual Task ProcessDocumentAsync(DocumentSnapshot document, CancellationToken cancellationToken)
     {
-        await document.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
-
         UpdateFileInfo(document);
+
+        return Task.CompletedTask;
     }
 
     public virtual void EnqueueIfNecessary(DocumentKey documentKey)