Merge pull request #67982 from sharwell/clear-request

Avoid retaining memory while waiting for changes

Author: sharwell

src/Features/LanguageServer/Protocol/Handler/Diagnostics/AbstractPullDiagnosticHandler.cs

@@ -191,6 +191,10 @@ await ComputeAndReportCurrentDiagnosticsAsync(
                 }
             }
 
+            // Clear out the solution context to avoid retaining memory
+            // https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1809058
+            context.ClearSolutionContext();
+
             // Some implementations of the spec will re-open requests as soon as we close them, spamming the server.
             // In those cases, we wait for the implementation to indicate that changes have occurred, then we close the connection
             // so that the client asks us again.

src/Features/LanguageServer/Protocol/Handler/RequestContext.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Text;
@@ -44,22 +45,79 @@ internal readonly struct RequestContext
 
     private readonly ILspServices _lspServices;
 
+    /// <summary>
+    /// Provides backing storage for the LSP workspace used by this RequestContext instance, allowing it to be cleared
+    /// on demand from all copies that may exist of this value type.
+    /// </summary>
+    /// <remarks>
+    /// This field is only initialized for handlers that request solution context.
+    /// </remarks>
+    private readonly StrongBox<(Workspace Workspace, Solution Solution, Document? Document)>? _lspSolution;
+
     /// <summary>
     /// The workspace this request is for, if applicable.  This will be present if <see cref="Document"/> is
     /// present.  It will be <see langword="null"/> if <c>requiresLSPSolution</c> is false.
     /// </summary>
-    public readonly Workspace? Workspace;
+    public Workspace? Workspace
+    {
+        get
+        {
+            if (_lspSolution is null)
+            {
+                // This request context never had a workspace instance
+                return null;
+            }
+
+            // The workspace is available unless it has been cleared by a call to ClearSolutionContext. Explicitly throw
+            // for attempts to access this property after it has been manually cleared.
+            return _lspSolution.Value.Workspace ?? throw new InvalidOperationException();
+        }
+    }
 
     /// <summary>
     /// The solution state that the request should operate on, if the handler requires an LSP solution, or <see langword="null"/> otherwise
     /// </summary>
-    public readonly Solution? Solution;
+    public Solution? Solution
+    {
+        get
+        {
+            if (_lspSolution is null)
+            {
+                // This request context never had a solution instance
+                return null;
+            }
+
+            // The solution is available unless it has been cleared by a call to ClearSolutionContext. Explicitly throw
+            // for attempts to access this property after it has been manually cleared.
+            return _lspSolution.Value.Solution ?? throw new InvalidOperationException();
+        }
+    }
 
     /// <summary>
     /// The document that the request is for, if applicable. This comes from the <see cref="TextDocumentIdentifier"/> returned from the handler itself via a call to 
     /// <see cref="ITextDocumentIdentifierHandler{RequestType, TextDocumentIdentifierType}.GetTextDocumentIdentifier(RequestType)"/>.
     /// </summary>
-    public readonly Document? Document;
+    public Document? Document
+    {
+        get
+        {
+            if (_lspSolution is null)
+            {
+                // This request context never had a solution instance
+                return null;
+            }
+
+            // The solution is available unless it has been cleared by a call to ClearSolutionContext. Explicitly throw
+            // for attempts to access this property after it has been manually cleared. Note that we can't rely on
+            // Document being null for this check, because it is not always provided as part of the solution context.
+            if (_lspSolution.Value.Workspace is null)
+            {
+                throw new InvalidOperationException();
+            }
+
+            return _lspSolution.Value.Document;
+        }
+    }
 
     /// <summary>
     /// The LSP server handling the request.
@@ -97,9 +155,18 @@ public RequestContext(
         ILspServices lspServices,
         CancellationToken queueCancellationToken)
     {
-        Workspace = workspace;
-        Document = document;
-        Solution = solution;
+        if (workspace is not null)
+        {
+            RoslynDebug.Assert(solution is not null);
+            _lspSolution = new StrongBox<(Workspace Workspace, Solution Solution, Document? Document)>((workspace, solution, document));
+        }
+        else
+        {
+            RoslynDebug.Assert(solution is null);
+            RoslynDebug.Assert(document is null);
+            _lspSolution = null;
+        }
+
         _clientCapabilities = clientCapabilities;
         ServerKind = serverKind;
         SupportedLanguages = supportedLanguages;
@@ -228,6 +295,14 @@ public ValueTask StopTrackingAsync(Uri uri, CancellationToken cancellationToken)
     public bool IsTracking(Uri documentUri)
         => _trackedDocuments.ContainsKey(documentUri);
 
+    public void ClearSolutionContext()
+    {
+        if (_lspSolution is null)
+            return;
+
+        _lspSolution.Value = default;
+    }
+
     /// <summary>
     /// Logs an informational message.
     /// </summary>