Allow html requests to be properly cancelled

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/HtmlDocumentServices/HtmlDocumentSynchronizer.SynchronizationRequest.cs

@@ -21,18 +21,20 @@ private class SynchronizationRequest(RazorDocumentVersion requestedVersion) : ID
 
         public RazorDocumentVersion RequestedVersion => _requestedVersion;
 
-        internal static SynchronizationRequest CreateAndStart(TextDocument document, RazorDocumentVersion requestedVersion, Func<TextDocument, RazorDocumentVersion, CancellationToken, Task<SynchronizationResult>> syncFunction)
+        internal static SynchronizationRequest CreateAndStart(TextDocument document, RazorDocumentVersion requestedVersion, Func<TextDocument, RazorDocumentVersion, CancellationToken, Task<SynchronizationResult>> syncFunction, CancellationToken cancellationToken)
         {
             var request = new SynchronizationRequest(requestedVersion);
-            request.Start(document, syncFunction);
+            request.Start(document, syncFunction, cancellationToken);
             return request;
         }
 
-        private void Start(TextDocument document, Func<TextDocument, RazorDocumentVersion, CancellationToken, Task<SynchronizationResult>> syncFunction)
+        private void Start(TextDocument document, Func<TextDocument, RazorDocumentVersion, CancellationToken, Task<SynchronizationResult>> syncFunction, CancellationToken cancellationToken)
         {
-            _cts = new(TimeSpan.FromMinutes(1));
+            _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+            var token = _cts.Token;
+            _cts.CancelAfter(TimeSpan.FromMilliseconds(500));
             _cts.Token.Register(Dispose);
-            _ = syncFunction.Invoke(document, _requestedVersion, _cts.Token).ContinueWith((t, state) =>
+            _ = syncFunction.Invoke(document, _requestedVersion, token).ContinueWith((t, state) =>
             {
                 var tcs = (TaskCompletionSource<SynchronizationResult>)state.AssumeNotNull();
                 if (t.IsCanceled)
@@ -50,7 +52,7 @@ private void Start(TextDocument document, Func<TextDocument, RazorDocumentVersio
 
                 _cts?.Dispose();
                 _cts = null;
-            }, _tcs, _cts.Token, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
+            }, _tcs, token, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
         }
 
         public void Dispose()

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/HtmlDocumentServices/HtmlDocumentSynchronizer.cs

@@ -11,6 +11,7 @@
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Remote;
+using Microsoft.CodeAnalysis.Razor.Threading;
 
 namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
@@ -27,18 +28,20 @@ internal sealed partial class HtmlDocumentSynchronizer(
     private readonly ILogger _logger = loggerFactory.GetOrCreateLogger<HtmlDocumentSynchronizer>();
 
     private readonly Dictionary<Uri, SynchronizationRequest> _synchronizationRequests = [];
-    private readonly object _gate = new();
+    // Semaphore to lock access to the dictionary above
+#pragma warning disable RS0030 // Do not use banned APIs
+    private readonly SemaphoreSlim _semaphore = new(initialCount: 1);
+#pragma warning restore RS0030 // Do not use banned APIs
 
-    public void DocumentRemoved(Uri razorFileUri)
+    public void DocumentRemoved(Uri razorFileUri, CancellationToken cancellationToken)
     {
-        lock (_gate)
+        using var _ = _semaphore.DisposableWait(cancellationToken);
+
+        if (_synchronizationRequests.TryGetValue(razorFileUri, out var request))
         {
-            if (_synchronizationRequests.TryGetValue(razorFileUri, out var request))
-            {
-                _logger.LogDebug($"Document {razorFileUri} removed, so we're disposing and clearing out the sync request for it");
-                request.Dispose();
-                _synchronizationRequests.Remove(razorFileUri);
-            }
+            _logger.LogDebug($"Document {razorFileUri} removed, so we're disposing and clearing out the sync request for it");
+            request.Dispose();
+            _synchronizationRequests.Remove(razorFileUri);
         }
     }
 
@@ -48,17 +51,39 @@ public async Task<SynchronizationResult> TrySynchronizeAsync(TextDocument docume
 
         _logger.LogDebug($"TrySynchronize for {document.FilePath} as at {requestedVersion}");
 
-        // We are not passing on the cancellation token through to the actual task that does the generation, because
-        // we do the actual work on whatever happens to be the first request that needs Html, without knowing how important
-        // that request is, nor why it might be cancelled. If that request is cancelled because of a document update,
-        // then the next request will cancel any work we start anyway, and if that request was cancelled for some other reason,
-        // then the next request probably wants the same version of the document, so we've got a head start.
-        return await GetSynchronizationRequestTaskAsync(document, requestedVersion).ConfigureAwait(false);
+        return await GetSynchronizationRequestTaskAsync(document, requestedVersion, cancellationToken).ConfigureAwait(false);
     }
 
-    private Task<SynchronizationResult> GetSynchronizationRequestTaskAsync(TextDocument document, RazorDocumentVersion requestedVersion)
+    /// <summary>
+    /// Returns a task that will complete when the html document for the Razor document has been made available
+    /// </summary>
+    /// <remarks>
+    /// Whilst this is a task-returning method, really its not is to manage the <see cref="_synchronizationRequests" /> dictionary.
+    /// When this method is called, one of 3 things could happen:
+    /// <list type="number">
+    /// <item>Nobody has asked for that document before, or they asked but the task failed, so a new task is started and returned</item>
+    /// <item>Somebody else already asked for that document, so you get the task they were given</item>
+    /// <item>Somebody else already asked for a future version of that document, so you get nothing</item>
+    /// </list>
+    /// If option 1 is taken, any pending tasks for older versions of the document will be cancelled.
+    /// </remarks>
+    private async Task<SynchronizationResult> GetSynchronizationRequestTaskAsync(TextDocument document, RazorDocumentVersion requestedVersion, CancellationToken cancellationToken)
     {
-        lock (_gate)
+        Task<SynchronizationResult> taskToGetResult;
+        using (var _ = await _semaphore.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                _logger.LogDebug($"Not synchronizing Html text for {document.FilePath} as the request was cancelled.");
+                return default;
+            }
+
+            taskToGetResult = GetOrAddResultTask_CallUnderLockAsync();
+        }
+
+        return await taskToGetResult.ConfigureAwait(false);
+
+        Task<SynchronizationResult> GetOrAddResultTask_CallUnderLockAsync()
         {
             var documentUri = document.CreateUri();
             if (_synchronizationRequests.TryGetValue(documentUri, out var request))
@@ -69,20 +94,18 @@ private Task<SynchronizationResult> GetSynchronizationRequestTaskAsync(TextDocum
                     // Html documents don't require semantic information. WorkspaceVersion changed too often to be used as a measure
                     // of equality for this purpose.
 
-#pragma warning disable VSTHRD103 // Use await instead of .Result
-#pragma warning disable VSTHRD003 // Avoid awaiting foreign Tasks
-                    if (request.Task.IsCompleted && request.Task.Result.Equals(default))
+                    if (request.Task.IsCompleted && !request.Task.VerifyCompleted().Synchronized)
                     {
                         _logger.LogDebug($"Already finished that version for {document.FilePath}, but was unsuccessful, so will recompute");
                         request.Dispose();
                     }
                     else
                     {
                         _logger.LogDebug($"Already {(request.Task.IsCompleted ? "finished" : "working on")} that version for {document.FilePath}");
+#pragma warning disable VSTHRD003 // Avoid awaiting foreign Tasks
                         return request.Task;
-                    }
 #pragma warning restore VSTHRD003 // Avoid awaiting foreign Tasks
-#pragma warning restore VSTHRD103 // Use await instead of .Result
+                    }
                 }
                 else if (requestedVersion.WorkspaceVersion < request.RequestedVersion.WorkspaceVersion)
                 {
@@ -104,7 +127,7 @@ private Task<SynchronizationResult> GetSynchronizationRequestTaskAsync(TextDocum
 
             _logger.LogDebug($"Going to start working on Html for {document.FilePath} as at {requestedVersion}");
 
-            var newRequest = SynchronizationRequest.CreateAndStart(document, requestedVersion, PublishHtmlDocumentAsync);
+            var newRequest = SynchronizationRequest.CreateAndStart(document, requestedVersion, PublishHtmlDocumentAsync, cancellationToken);
             _synchronizationRequests[documentUri] = newRequest;
             return newRequest.Task;
         }
@@ -119,15 +142,15 @@ private async Task<SynchronizationResult> PublishHtmlDocumentAsync(TextDocument
                 (service, solutionInfo, ct) => service.GetHtmlDocumentTextAsync(solutionInfo, document.Id, ct),
                 cancellationToken).ConfigureAwait(false);
         }
-        catch (Exception ex)
+        catch (Exception ex) when (ex is not OperationCanceledException)
         {
             _logger.LogError(ex, $"Error getting Html text for {document.FilePath}. Html document contents will be stale");
             return default;
         }
 
         if (cancellationToken.IsCancellationRequested)
         {
-            // Checking cancellation before logging, as a new request coming in doesn't count as "Couldn't get Html"
+            _logger.LogDebug($"Not publishing Html text for {document.FilePath} as the request was cancelled.");
             return default;
         }
 
@@ -141,6 +164,14 @@ private async Task<SynchronizationResult> PublishHtmlDocumentAsync(TextDocument
         {
             var result = new SynchronizationResult(true, requestedVersion.Checksum);
             await _htmlDocumentPublisher.PublishAsync(document, result, htmlText, cancellationToken).ConfigureAwait(false);
+
+            // If we were cancelled, we can't trust that the publish worked.
+            if (cancellationToken.IsCancellationRequested)
+            {
+                _logger.LogDebug($"Not publishing Html text for {document.FilePath} as the request was cancelled.");
+                return default;
+            }
+
             return result;
         }
         catch (Exception ex)
@@ -164,7 +195,7 @@ internal TestAccessor(HtmlDocumentSynchronizer instance)
             _instance = instance;
         }
 
-        public Task<SynchronizationResult> GetSynchronizationRequestTaskAsync(TextDocument document, RazorDocumentVersion requestedVersion)
-            => _instance.GetSynchronizationRequestTaskAsync(document, requestedVersion);
+        public Task<SynchronizationResult> GetSynchronizationRequestTaskAsync(TextDocument document, RazorDocumentVersion requestedVersion, CancellationToken cancellationToken)
+            => _instance.GetSynchronizationRequestTaskAsync(document, requestedVersion, cancellationToken);
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/HtmlDocumentServices/IHtmlDocumentSynchronizer.cs

@@ -10,6 +10,6 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
 internal interface IHtmlDocumentSynchronizer
 {
-    void DocumentRemoved(Uri uri);
+    void DocumentRemoved(Uri uri, CancellationToken cancellationToken);
     Task<SynchronizationResult> TrySynchronizeAsync(TextDocument document, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Cohost/HtmlDocumentRemoveListener.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System.ComponentModel.Composition;
+using System.Threading;
 using Microsoft.VisualStudio.LanguageServer.ContainedLanguage;
 using Microsoft.VisualStudio.Utilities;
 
@@ -20,7 +21,7 @@ public override void Changed(LSPDocumentSnapshot? old, LSPDocumentSnapshot? @new
     {
         if (kind == LSPDocumentChangeKind.Removed && old is not null)
         {
-            _htmlDocumentSynchronizer.DocumentRemoved(old.Uri);
+            _htmlDocumentSynchronizer.DocumentRemoved(old.Uri, CancellationToken.None);
         }
     }
 }

src/Razor/src/Microsoft.VisualStudioCode.RazorExtension/Endpoints/RazorDocumentClosedEndpoint.cs

@@ -29,7 +29,7 @@ internal class RazorDocumentClosedEndpoint(IHtmlDocumentSynchronizer htmlDocumen
 
     protected override Task<VoidResult> HandleRequestAsync(TextDocumentIdentifier textDocument, RazorCohostRequestContext requestContext, CancellationToken cancellationToken)
     {
-        _htmlDocumentSynchronizer.DocumentRemoved(requestContext.Uri.AssumeNotNull());
+        _htmlDocumentSynchronizer.DocumentRemoved(requestContext.Uri.AssumeNotNull(), cancellationToken);
         return SpecializedTasks.Default<VoidResult>();
     }
 }

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/HtmlDocumentSynchronizerTest.cs

@@ -8,6 +8,7 @@
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.Test.Common.VisualStudio;
+using Microsoft.AspNetCore.Razor.Threading;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.Remote;
@@ -63,7 +64,7 @@ public async Task TrySynchronize_ReopenedDocument_Generates()
         Assert.True(syncResult.Synchronized);
 
         // "Close" the document
-        synchronizer.DocumentRemoved(document.CreateUri());
+        synchronizer.DocumentRemoved(document.CreateUri(), DisposalToken);
 
         Assert.True((await synchronizer.TrySynchronizeAsync(document, DisposalToken)).Synchronized);
 
@@ -230,7 +231,6 @@ public async Task TrySynchronize_RequestOldVersion_ImmediateFail()
 
         Assert.True((await task).Synchronized);
 
-        // We should have two publishes
         Assert.Collection(publisher.Publishes,
             i =>
             {
@@ -285,9 +285,7 @@ public async Task TrySynchronize_RequestNewVersion_CancelOldTask()
         tcs.SetResult(true);
 
         await Task.WhenAll(task1, task2);
-#pragma warning disable xUnit1031 // Do not use blocking task operations in test method
-        Assert.False(task1.Result.Synchronized);
-#pragma warning restore xUnit1031 // Do not use blocking task operations in test method
+        Assert.False(task1.VerifyCompleted().Synchronized);
 
         Assert.Collection(publisher.Publishes,
             i =>
@@ -298,22 +296,28 @@ public async Task TrySynchronize_RequestNewVersion_CancelOldTask()
     }
 
     [Fact]
-    public async Task GetSynchronizationRequestTask_RequestSameVersion_ReturnsSameTask()
+    public async Task GetSynchronizationRequestTask_RequestSameVersion_InvokedRemoteOnce()
     {
         var document = Workspace.CurrentSolution.GetAdditionalDocument(_documentId).AssumeNotNull();
 
         var tcs = new TaskCompletionSource<bool>();
         var publisher = new TestHtmlDocumentPublisher();
-        var remoteServiceInvoker = new RemoteServiceInvoker(document, () => tcs.Task);
+
+        var remoteInvocations = 0;
+        var remoteServiceInvoker = new RemoteServiceInvoker(document, () =>
+        {
+            remoteInvocations++;
+            return tcs.Task;
+        });
         var synchronizer = new HtmlDocumentSynchronizer(remoteServiceInvoker, publisher, LoggerFactory);
 
         var version = await RazorDocumentVersion.CreateAsync(document, DisposalToken);
 
         var accessor = synchronizer.GetTestAccessor();
-        var task1 = accessor.GetSynchronizationRequestTaskAsync(document, version);
-        var task2 = accessor.GetSynchronizationRequestTaskAsync(document, version);
+        var task1 = accessor.GetSynchronizationRequestTaskAsync(document, version, DisposalToken);
+        var task2 = accessor.GetSynchronizationRequestTaskAsync(document, version, DisposalToken);
 
-        Assert.Same(task1, task2);
+        Assert.Equal(1, remoteInvocations);
     }
 
     private class RemoteServiceInvoker(TextDocument document, Func<Task>? generateTask = null) : IRemoteServiceInvoker
@@ -349,7 +353,7 @@ private class RemoteServiceInvoker(TextDocument document, Func<Task>? generateTa
 
     private class TestHtmlDocumentPublisher : IHtmlDocumentPublisher
     {
-        private List<(TextDocument, string, ChecksumWrapper)> _publishes = [];
+        private readonly List<(TextDocument, string, ChecksumWrapper)> _publishes = [];
 
         public List<(TextDocument Document, string Text, ChecksumWrapper Checksum)> Publishes => _publishes;