[LSP] Fix race in document change tracker when used in non-mutating requests.

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

@@ -3,12 +3,8 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
-using System.Collections.Generic;
 using System.Collections.Immutable;
-using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using Microsoft.CodeAnalysis.Options;
-using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using Roslyn.Utilities;
@@ -22,14 +18,19 @@ namespace Microsoft.CodeAnalysis.LanguageServer.Handler
     internal readonly struct RequestContext
     {
         /// <summary>
-        /// This will be null for non-mutating requests because they're not allowed to change documents
+        /// This will be the <see cref="NonMutatingDocumentChangeTracker"/> for non-mutating requests because they're not allowed to change documents
         /// </summary>
         private readonly IDocumentChangeTracker _documentChangeTracker;
 
         /// <summary>
-        /// Manages the workspaces registered for LSP and handles updates from both LSP text sync and workspace updates.
+        /// Contains the LSP text for all opened LSP documents from when this request was processed in the queue.
         /// </summary>
-        private readonly LspWorkspaceManager _lspWorkspaceManager;
+        /// <remarks>
+        /// This is a snapshot of the source text that reflects the LSP text based on the order of this request in the queue.
+        /// It contains text that is consistent with all prior LSP text sync notifications, but LSP text sync requests
+        /// which are ordered after this one in the queue are not reflected here.
+        /// </remarks>
+        private readonly ImmutableDictionary<Uri, SourceText> _trackedDocuments;
 
         /// <summary>
         /// The solution state that the request should operate on, if the handler requires an LSP solution, or <see langword="null"/> otherwise
@@ -70,7 +71,7 @@ public RequestContext(
             string? clientName,
             Document? document,
             IDocumentChangeTracker documentChangeTracker,
-            LspWorkspaceManager lspWorkspaceManager,
+            ImmutableDictionary<Uri, SourceText> trackedDocuments,
             ImmutableArray<string> supportedLanguages,
             IGlobalOptionService globalOptions)
         {
@@ -82,7 +83,7 @@ public RequestContext(
             GlobalOptions = globalOptions;
             _documentChangeTracker = documentChangeTracker;
             _traceInformation = traceInformation;
-            _lspWorkspaceManager = lspWorkspaceManager;
+            _trackedDocuments = trackedDocuments;
         }
 
         public static RequestContext? Create(
@@ -96,13 +97,17 @@ public RequestContext(
             ImmutableArray<string> supportedLanguages,
             IGlobalOptionService globalOptions)
         {
+            // Retrieve the current LSP tracked text as of this request.
+            // This is safe as all creation of request contexts cannot happen concurrently.
+            var trackedDocuments = lspWorkspaceManager.GetTrackedLspText();
+
             // If the handler doesn't need an LSP solution we do two important things:
             // 1. We don't bother building the LSP solution for perf reasons
             // 2. We explicitly don't give the handler a solution or document, even if we could
             //    so they're not accidentally operating on stale solution state.
             if (!requiresLSPSolution)
             {
-                return new RequestContext(solution: null, logger.TraceInformation, clientCapabilities, clientName, document: null, documentChangeTracker, lspWorkspaceManager, supportedLanguages, globalOptions);
+                return new RequestContext(solution: null, logger.TraceInformation, clientCapabilities, clientName, document: null, documentChangeTracker, trackedDocuments, supportedLanguages, globalOptions);
             }
 
             // Go through each registered workspace, find the solution that contains the document that
@@ -133,7 +138,7 @@ public RequestContext(
                 clientName,
                 document,
                 documentChangeTracker,
-                lspWorkspaceManager,
+                trackedDocuments,
                 supportedLanguages,
                 globalOptions);
             return context;
@@ -144,36 +149,30 @@ public RequestContext(
         /// Mutating requests are serialized by the execution queue in order to prevent concurrent access.
         /// </summary>
         public void StartTracking(Uri uri, SourceText initialText)
-        {
-            _documentChangeTracker.StartTracking(uri, initialText);
-            _lspWorkspaceManager.TrackLspDocument(uri, initialText);
-        }
+            => _documentChangeTracker.StartTracking(uri, initialText);
 
         /// <summary>
         /// Allows a mutating request to update the contents of a tracked document.
         /// Mutating requests are serialized by the execution queue in order to prevent concurrent access.
         /// </summary>
         public void UpdateTrackedDocument(Uri uri, SourceText changedText)
-        {
-            _documentChangeTracker.UpdateTrackedDocument(uri, changedText);
-            _lspWorkspaceManager.UpdateLspDocument(uri, changedText);
-        }
+            => _documentChangeTracker.UpdateTrackedDocument(uri, changedText);
 
         public SourceText GetTrackedDocumentSourceText(Uri documentUri)
-            => _documentChangeTracker.GetTrackedDocumentSourceText(documentUri);
+        {
+            Contract.ThrowIfFalse(_trackedDocuments.ContainsKey(documentUri), $"Attempted to get text for {documentUri} which is not open.");
+            return _trackedDocuments[documentUri];
+        }
 
         /// <summary>
         /// Allows a mutating request to close a document and stop it being tracked.
         /// Mutating requests are serialized by the execution queue in order to prevent concurrent access.
         /// </summary>
         public void StopTracking(Uri uri)
-        {
-            _documentChangeTracker.StopTracking(uri);
-            _lspWorkspaceManager.StopTrackingLspDocument(uri);
-        }
+            => _documentChangeTracker.StopTracking(uri);
 
         public bool IsTracking(Uri documentUri)
-            => _documentChangeTracker.IsTracking(documentUri);
+            => _trackedDocuments.ContainsKey(documentUri);
 
         /// <summary>
         /// Logs an informational message.

src/Features/LanguageServer/Protocol/Handler/RequestExecutionQueue.DocumentChangeTracker.cs

@@ -6,7 +6,7 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
-using Microsoft.CodeAnalysis.Host;
+using Microsoft.CodeAnalysis.LanguageServer.Handler.DocumentChanges;
 using Microsoft.CodeAnalysis.Text;
 using Roslyn.Utilities;
 
@@ -16,102 +16,31 @@ internal partial class RequestExecutionQueue
     {
         /// <summary>
         /// Associates LSP document URIs with the roslyn source text containing the LSP document text.
+        /// Called via <see cref="DidOpenHandler"/>, <see cref="DidChangeHandler"/> and <see cref="DidCloseHandler"/>
         /// </summary>
         internal interface IDocumentChangeTracker
         {
             void StartTracking(Uri documentUri, SourceText initialText);
             void UpdateTrackedDocument(Uri documentUri, SourceText text);
             void StopTracking(Uri documentUri);
-            bool IsTracking(Uri documentUri);
-            ImmutableArray<(Uri DocumentUri, SourceText Text)> GetTrackedDocuments();
-            SourceText GetTrackedDocumentSourceText(Uri documentUri);
         }
 
         private class NonMutatingDocumentChangeTracker : IDocumentChangeTracker
         {
-            private readonly DocumentChangeTracker _tracker;
-
-            public NonMutatingDocumentChangeTracker(DocumentChangeTracker tracker)
-            {
-                _tracker = tracker;
-            }
-
-            public ImmutableArray<(Uri DocumentUri, SourceText Text)> GetTrackedDocuments()
-                => _tracker.GetTrackedDocuments();
-
-            public SourceText GetTrackedDocumentSourceText(Uri documentUri)
-            {
-                Contract.Fail("Mutating documents not allowed in a non-mutating request handler");
-                throw new NotImplementedException();
-            }
-
-            public bool IsTracking(Uri documentUri)
-                => _tracker.IsTracking(documentUri);
-
             public void StartTracking(Uri documentUri, SourceText initialText)
             {
-                Contract.Fail("Mutating documents not allowed in a non-mutating request handler");
-                throw new NotImplementedException();
+                throw new InvalidOperationException("Mutating documents not allowed in a non-mutating request handler");
             }
 
             public void StopTracking(Uri documentUri)
             {
-                Contract.Fail("Mutating documents not allowed in a non-mutating request handler");
-                throw new NotImplementedException();
-            }
-
-            public void UpdateTrackedDocument(Uri documentUri, SourceText text)
-            {
-                Contract.Fail("Mutating documents not allowed in a non-mutating request handler");
-                throw new NotImplementedException();
-            }
-        }
-
-        /// <summary>
-        /// Keeps track of changes to documents that are opened in the LSP client. Calls MUST not overlap, so this
-        /// should be called from a mutating request handler. See <see cref="RequestExecutionQueue"/> for more details.
-        /// </summary>
-        internal class DocumentChangeTracker : IDocumentChangeTracker
-        {
-            private readonly Dictionary<Uri, SourceText> _trackedDocuments = new();
-
-            public DocumentChangeTracker()
-            {
-            }
-
-            public bool IsTracking(Uri documentUri)
-                => _trackedDocuments.ContainsKey(documentUri);
-
-            public void StartTracking(Uri documentUri, SourceText initialText)
-            {
-                Contract.ThrowIfTrue(_trackedDocuments.ContainsKey(documentUri), $"didOpen received for {documentUri} which is already open.");
-
-                _trackedDocuments.Add(documentUri, initialText);
+                throw new InvalidOperationException("Mutating documents not allowed in a non-mutating request handler");
             }
 
             public void UpdateTrackedDocument(Uri documentUri, SourceText text)
             {
-                Contract.ThrowIfFalse(_trackedDocuments.ContainsKey(documentUri), $"didChange received for {documentUri} which is not open.");
-
-                _trackedDocuments[documentUri] = text;
+                throw new InvalidOperationException("Mutating documents not allowed in a non-mutating request handler");
             }
-
-            public SourceText GetTrackedDocumentSourceText(Uri documentUri)
-            {
-                Contract.ThrowIfFalse(_trackedDocuments.ContainsKey(documentUri), "didChange received for a document that isn't open.");
-
-                return _trackedDocuments[documentUri];
-            }
-
-            public void StopTracking(Uri documentUri)
-            {
-                Contract.ThrowIfFalse(_trackedDocuments.ContainsKey(documentUri), $"didClose received for {documentUri} which is not open.");
-
-                _trackedDocuments.Remove(documentUri);
-            }
-
-            public ImmutableArray<(Uri DocumentUri, SourceText Text)> GetTrackedDocuments()
-                => _trackedDocuments.Select(k => (k.Key, k.Value)).ToImmutableArray();
         }
 
         internal TestAccessor GetTestAccessor()
@@ -124,8 +53,8 @@ internal readonly struct TestAccessor
             public TestAccessor(RequestExecutionQueue queue)
                 => _queue = queue;
 
-            public List<SourceText> GetTrackedTexts()
-                => _queue._documentChangeTracker.GetTrackedDocuments().Select(i => i.Text).ToList();
+            public ImmutableArray<SourceText> GetTrackedTexts()
+                => _queue._lspWorkspaceManager.GetTrackedLspText().Select(i => i.Value).ToImmutableArray();
 
             public LspWorkspaceManager GetLspWorkspaceManager() => _queue._lspWorkspaceManager;
 

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

@@ -58,7 +58,6 @@ internal partial class RequestExecutionQueue
 
         private readonly AsyncQueue<QueueItem> _queue;
         private readonly CancellationTokenSource _cancelSource;
-        private readonly DocumentChangeTracker _documentChangeTracker;
         private readonly RequestTelemetryLogger _requestTelemetryLogger;
         private readonly IGlobalOptionService _globalOptions;
 
@@ -93,15 +92,14 @@ public RequestExecutionQueue(
 
             _queue = new AsyncQueue<QueueItem>();
             _cancelSource = new CancellationTokenSource();
-            _documentChangeTracker = new DocumentChangeTracker();
 
             // Pass the language client instance type name to the telemetry logger to ensure we can
             // differentiate between the different C# LSP servers that have the same client name.
             // We also don't use the language client's name property as it is a localized user facing string
             // which is difficult to write telemetry queries for.
             _requestTelemetryLogger = new RequestTelemetryLogger(serverTypeName);
 
-            _lspWorkspaceManager = new LspWorkspaceManager(lspWorkspaceRegistrationService, lspMiscellaneousFilesWorkspace, _documentChangeTracker, logger, _requestTelemetryLogger);
+            _lspWorkspaceManager = new LspWorkspaceManager(logger, lspMiscellaneousFilesWorkspace, lspWorkspaceRegistrationService, _requestTelemetryLogger);
 
             // Start the queue processing
             _ = ProcessQueueAsync();
@@ -307,8 +305,8 @@ private void DrainQueue()
         private RequestContext? CreateRequestContext(QueueItem queueItem)
         {
             var trackerToUse = queueItem.MutatesSolutionState
-                ? (IDocumentChangeTracker)_documentChangeTracker
-                : new NonMutatingDocumentChangeTracker(_documentChangeTracker);
+                ? (IDocumentChangeTracker)_lspWorkspaceManager
+                : new NonMutatingDocumentChangeTracker();
 
             return RequestContext.Create(
                 queueItem.RequiresLSPSolution,