IDocumentSnapshot finally gets ValueTasks and CancellationTokens

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorCSharpFormattingBenchmark.cs

@@ -58,7 +58,7 @@ public async Task InitializeRazorCSharpFormattingAsync()
 
         DocumentUri = new Uri(_filePath);
         DocumentSnapshot = await GetDocumentSnapshotAsync(projectFilePath, _filePath, targetPath);
-        DocumentText = await DocumentSnapshot.GetTextAsync();
+        DocumentText = await DocumentSnapshot.GetTextAsync(CancellationToken.None);
     }
 
     private static void WriteSampleFormattingFile(string filePath, bool preformatted, int blocks)

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorCodeActionsBenchmark.cs

@@ -77,7 +77,7 @@ public async Task SetupAsync()
 
         DocumentUri = new Uri(_filePath);
         DocumentSnapshot = await GetDocumentSnapshotAsync(projectFilePath, _filePath, targetPath);
-        DocumentText = await DocumentSnapshot.GetTextAsync();
+        DocumentText = await DocumentSnapshot.GetTextAsync(CancellationToken.None);
 
         RazorCodeActionRange = DocumentText.GetZeroWidthRange(razorCodeActionIndex);
         CSharpCodeActionRange = DocumentText.GetZeroWidthRange(csharpCodeActionIndex);

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorCompletionBenchmark.cs

@@ -75,7 +75,7 @@ public async Task SetupAsync()
 
         DocumentUri = new Uri(_filePath);
         DocumentSnapshot = await GetDocumentSnapshotAsync(projectFilePath, _filePath, targetPath);
-        DocumentText = await DocumentSnapshot.GetTextAsync();
+        DocumentText = await DocumentSnapshot.GetTextAsync(CancellationToken.None);
 
         RazorPosition = DocumentText.GetPosition(razorCodeActionIndex);
 

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorDocumentMappingBenchmark.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Threading;
 using System.Threading.Tasks;
 using BenchmarkDotNet.Attributes;
 using Microsoft.AspNetCore.Razor.Language;
@@ -50,7 +51,7 @@ public async Task InitializeRazorCSharpFormattingAsync()
 
         DocumentSnapshot = await GetDocumentSnapshotAsync(projectFilePath, _filePath, targetPath);
 
-        var codeDocument = await DocumentSnapshot.GetGeneratedOutputAsync();
+        var codeDocument = await DocumentSnapshot.GetGeneratedOutputAsync(CancellationToken.None);
         CSharpDocument = codeDocument.GetCSharpDocument();
     }
 

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorSemanticTokensScrollingBenchmark.cs

@@ -51,7 +51,7 @@ public async Task InitializeRazorSemanticAsync()
         var documentSnapshot = await GetDocumentSnapshotAsync(ProjectFilePath, filePath, TargetPath);
         DocumentContext = new DocumentContext(documentUri, documentSnapshot, projectContext: null);
 
-        var text = await DocumentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var text = await DocumentSnapshot.GetTextAsync(CancellationToken.None).ConfigureAwait(false);
         Range = VsLspFactory.CreateRange(
             start: (0, 0),
             end: (text.Lines.Count - 1, 0));

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/ProjectSystem/BackgroundCodeGenerationBenchmark.cs

@@ -59,6 +59,6 @@ private void SnapshotManager_Changed(object sender, ProjectChangeEventArgs e)
         var project = ProjectManager.GetLoadedProject(e.ProjectKey);
         var document = project.GetDocument(e.DocumentFilePath);
 
-        Tasks.Add(document.GetGeneratedOutputAsync());
+        Tasks.Add(document.GetGeneratedOutputAsync(CancellationToken.None).AsTask());
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeActions/CSharp/UnformattedRemappingCSharpCodeActionResolver.cs

@@ -70,7 +70,7 @@ public async override Task<CodeAction> ResolveAsync(
             return codeAction;
         }
 
-        var codeDocument = await documentContext.Snapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentContext.Snapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (codeDocument.IsUnsupported())
         {
             return codeAction;

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeActions/CodeActionEndpoint.cs

@@ -65,7 +65,7 @@ internal sealed class CodeActionEndpoint(
             return null;
         }
 
-        var razorCodeActionContext = await GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot).ConfigureAwait(false);
+        var razorCodeActionContext = await GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot, cancellationToken).ConfigureAwait(false);
         if (razorCodeActionContext is null)
         {
             return null;
@@ -134,15 +134,18 @@ public void ApplyCapabilities(VSInternalServerCapabilities serverCapabilities, V
     }
 
     // internal for testing
-    internal async Task<RazorCodeActionContext?> GenerateRazorCodeActionContextAsync(VSCodeActionParams request, IDocumentSnapshot documentSnapshot)
+    internal async Task<RazorCodeActionContext?> GenerateRazorCodeActionContextAsync(
+        VSCodeActionParams request,
+        IDocumentSnapshot documentSnapshot,
+        CancellationToken cancellationToken)
     {
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (codeDocument.IsUnsupported())
         {
             return null;
         }
 
-        var sourceText = await documentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var sourceText = codeDocument.Source.Text;
 
         // VS Provides `CodeActionParams.Context.SelectionRange` in addition to
         // `CodeActionParams.Range`. The `SelectionRange` is relative to where the

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeActions/Html/DefaultHtmlCodeActionProvider.cs

@@ -49,7 +49,7 @@ public static async Task RemapAndFixHtmlCodeActionEditAsync(IEditMappingService
 
         if (codeAction.Edit.TryGetTextDocumentEdits(out var documentEdits))
         {
-            var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+            var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
             var htmlSourceText = codeDocument.GetHtmlSourceText();
 
             foreach (var edit in documentEdits)

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeActions/Razor/AddUsingsCodeActionResolver.cs

@@ -45,7 +45,7 @@ internal sealed class AddUsingsCodeActionResolver(IDocumentContextFactory docume
 
         var documentSnapshot = documentContext.Snapshot;
 
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (codeDocument.IsUnsupported())
         {
             return null;

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Diagnostics/DocumentPullDiagnosticsEndpoint.cs

@@ -80,7 +80,7 @@ public TextDocumentIdentifier GetTextDocumentIdentifier(VSInternalDocumentDiagno
 
         var documentSnapshot = documentContext.Snapshot;
 
-        var razorDiagnostics = await GetRazorDiagnosticsAsync(documentSnapshot).ConfigureAwait(false);
+        var razorDiagnostics = await GetRazorDiagnosticsAsync(documentSnapshot, cancellationToken).ConfigureAwait(false);
 
         await ReportRZ10012TelemetryAsync(documentContext, razorDiagnostics, cancellationToken).ConfigureAwait(false);
 
@@ -106,7 +106,9 @@ public TextDocumentIdentifier GetTextDocumentIdentifier(VSInternalDocumentDiagno
             {
                 if (report.Diagnostics is not null)
                 {
-                    var mappedDiagnostics = await _translateDiagnosticsService.TranslateAsync(RazorLanguageKind.CSharp, report.Diagnostics, documentSnapshot).ConfigureAwait(false);
+                    var mappedDiagnostics = await _translateDiagnosticsService
+                        .TranslateAsync(RazorLanguageKind.CSharp, report.Diagnostics, documentSnapshot, cancellationToken)
+                        .ConfigureAwait(false);
                     report.Diagnostics = mappedDiagnostics;
                 }
 
@@ -120,7 +122,9 @@ public TextDocumentIdentifier GetTextDocumentIdentifier(VSInternalDocumentDiagno
             {
                 if (report.Diagnostics is not null)
                 {
-                    var mappedDiagnostics = await _translateDiagnosticsService.TranslateAsync(RazorLanguageKind.Html, report.Diagnostics, documentSnapshot).ConfigureAwait(false);
+                    var mappedDiagnostics = await _translateDiagnosticsService
+                        .TranslateAsync(RazorLanguageKind.Html, report.Diagnostics, documentSnapshot, cancellationToken)
+                        .ConfigureAwait(false);
                     report.Diagnostics = mappedDiagnostics;
                 }
 
@@ -131,9 +135,9 @@ public TextDocumentIdentifier GetTextDocumentIdentifier(VSInternalDocumentDiagno
         return allDiagnostics.ToArray();
     }
 
-    private static async Task<VSInternalDiagnosticReport[]?> GetRazorDiagnosticsAsync(IDocumentSnapshot documentSnapshot)
+    private static async Task<VSInternalDiagnosticReport[]?> GetRazorDiagnosticsAsync(IDocumentSnapshot documentSnapshot, CancellationToken cancellationToken)
     {
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         var sourceText = codeDocument.Source.Text;
         var csharpDocument = codeDocument.GetCSharpDocument();
         var diagnostics = csharpDocument.Diagnostics;

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Diagnostics/RazorDiagnosticsPublisher.cs

@@ -113,10 +113,10 @@ private async ValueTask ProcessBatchAsync(ImmutableArray<IDocumentSnapshot> item
         }
     }
 
-    private async Task PublishDiagnosticsAsync(IDocumentSnapshot document, CancellationToken token)
+    private async Task PublishDiagnosticsAsync(IDocumentSnapshot document, CancellationToken cancellationToken)
     {
-        var result = await document.GetGeneratedOutputAsync().ConfigureAwait(false);
-        var csharpDiagnostics = await GetCSharpDiagnosticsAsync(document, token).ConfigureAwait(false);
+        var result = await document.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
+        var csharpDiagnostics = await GetCSharpDiagnosticsAsync(document, cancellationToken).ConfigureAwait(false);
         var razorDiagnostics = result.GetCSharpDocument().Diagnostics;
 
         lock (_publishedDiagnostics)
@@ -188,7 +188,7 @@ .. csharpDiagnostics ?? []
                     if (_documentContextFactory.Value.TryCreate(delegatedParams.TextDocument.Uri, projectContext: null, out var documentContext))
                     {
                         return await _translateDiagnosticsService.Value
-                            .TranslateAsync(RazorLanguageKind.CSharp, fullDiagnostics.Items, documentContext.Snapshot)
+                            .TranslateAsync(RazorLanguageKind.CSharp, fullDiagnostics.Items, documentContext.Snapshot, cancellationToken)
                             .ConfigureAwait(false);
                     }
                 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/DocumentSnapshotTextLoader.cs

@@ -9,23 +9,13 @@
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer;
 
-internal class DocumentSnapshotTextLoader : TextLoader
+internal class DocumentSnapshotTextLoader(IDocumentSnapshot documentSnapshot) : TextLoader
 {
-    private readonly IDocumentSnapshot _documentSnapshot;
-
-    public DocumentSnapshotTextLoader(IDocumentSnapshot documentSnapshot)
-    {
-        if (documentSnapshot is null)
-        {
-            throw new ArgumentNullException(nameof(documentSnapshot));
-        }
-
-        _documentSnapshot = documentSnapshot;
-    }
+    private readonly IDocumentSnapshot _documentSnapshot = documentSnapshot;
 
     public override async Task<TextAndVersion> LoadTextAndVersionAsync(LoadTextOptions options, CancellationToken cancellationToken)
     {
-        var sourceText = await _documentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var sourceText = await _documentSnapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
         var textAndVersion = TextAndVersion.Create(sourceText, VersionStamp.Default);
 
         return textAndVersion;

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Formatting/HtmlFormatter.cs

@@ -47,7 +47,7 @@ public async Task<ImmutableArray<TextChange>> GetDocumentFormattingEditsAsync(
             return [];
         }
 
-        var sourceText = await documentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var sourceText = await documentSnapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
         return result.Edits.SelectAsArray(sourceText.GetTextChange);
     }
 
@@ -78,7 +78,7 @@ public async Task<ImmutableArray<TextChange>> GetOnTypeFormattingEditsAsync(
             return [];
         }
 
-        var sourceText = await documentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var sourceText = await documentSnapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
         return result.Edits.SelectAsArray(sourceText.GetTextChange);
     }
 

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Formatting/LspFormattingCodeDocumentProvider.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.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.Formatting;
@@ -10,9 +11,9 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer.Formatting;
 
 internal sealed class LspFormattingCodeDocumentProvider : IFormattingCodeDocumentProvider
 {
-    public Task<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot)
+    public ValueTask<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot, CancellationToken cancellationToken)
     {
         // Formatting always uses design time
-        return snapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: true);
+        return snapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: true, cancellationToken);
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/MapCode/MapCodeEndpoint.cs

@@ -105,7 +105,7 @@ public void ApplyCapabilities(VSInternalServerCapabilities serverCapabilities, V
 
                 // We create a new Razor file based on each content in each mapping order to get the syntax tree that we'll later use to map.
                 var newSnapshot = snapshot.WithText(SourceText.From(content));
-                var codeToMap = await newSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+                var codeToMap = await newSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
 
                 var mappingSuccess = await TryMapCodeAsync(
                     codeToMap, mapping.FocusLocations, changes, mapCodeCorrelationId, documentContext, cancellationToken).ConfigureAwait(false);
@@ -228,7 +228,7 @@ private async Task<bool> TryMapCodeAsync(
                     razorNodesToMap.Add(nodeToMap);
                 }
 
-                var sourceText = await documentContext.Snapshot.GetTextAsync().ConfigureAwait(false);
+                var sourceText = await documentContext.Snapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
 
                 foreach (var nodeToMap in razorNodesToMap)
                 {

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Mapping/RazorLanguageQueryEndpoint.cs

@@ -42,8 +42,8 @@ public TextDocumentIdentifier GetTextDocumentIdentifier(RazorLanguageQueryParams
         var documentSnapshot = documentContext.Snapshot;
         var documentVersion = documentContext.Snapshot.Version;
 
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
-        var sourceText = await documentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
+        var sourceText = codeDocument.Source.Text;
         var hostDocumentIndex = sourceText.GetPosition(request.Position);
         var responsePosition = request.Position;
 

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/OpenDocumentGenerator.cs

@@ -75,7 +75,7 @@ private async ValueTask ProcessBatchAsync(ImmutableArray<IDocumentSnapshot> item
                 return;
             }
 
-            var codeDocument = await document.GetGeneratedOutputAsync().ConfigureAwait(false);
+            var codeDocument = await document.GetGeneratedOutputAsync(token).ConfigureAwait(false);
 
             foreach (var listener in _listeners)
             {

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Diagnostics/RazorTranslateDiagnosticsService.cs

@@ -6,6 +6,7 @@
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Language.Syntax;
@@ -45,18 +46,27 @@ internal class RazorTranslateDiagnosticsService(IDocumentMappingService document
     ]).ToFrozenSet();
 
     /// <summary>
-    /// Translates code diagnostics from one representation into another.
+    ///  Translates code diagnostics from one representation into another.
     /// </summary>
-    /// <param name="diagnosticKind">The `RazorLanguageKind` of the `Diagnostic` objects included in `diagnostics`.</param>
-    /// <param name="diagnostics">An array of `Diagnostic` objects to translate.</param>
-    /// <param name="documentSnapshot">The `DocumentContext` for the code document associated with the diagnostics.</param>
+    /// <param name="diagnosticKind">
+    ///  The <see cref="RazorLanguageKind"/> of the <see cref="Diagnostic"/> objects
+    ///  included in <paramref name="diagnostics"/>.
+    /// </param>
+    /// <param name="diagnostics">
+    ///  An array of <see cref="Diagnostic"/> objects to translate.
+    /// </param>
+    /// <param name="documentSnapshot">
+    ///  The <see cref="IDocumentSnapshot"/> for the code document associated with the diagnostics.
+    /// </param>
+    /// <param name="cancellationToken">A token that can be checked to cancel work.</param>
     /// <returns>An array of translated diagnostics</returns>
     internal async Task<LspDiagnostic[]> TranslateAsync(
         RazorLanguageKind diagnosticKind,
         LspDiagnostic[] diagnostics,
-        IDocumentSnapshot documentSnapshot)
+        IDocumentSnapshot documentSnapshot,
+        CancellationToken cancellationToken)
     {
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (codeDocument.IsUnsupported() != false)
         {
             _logger.LogInformation($"Unsupported code document.");

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/FormattingContext.cs

@@ -7,6 +7,7 @@
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Language.Syntax;
@@ -224,11 +225,13 @@ public bool TryGetFormattingSpan(int absoluteIndex, [NotNullWhen(true)] out Form
         return false;
     }
 
-    public async Task<FormattingContext> WithTextAsync(SourceText changedText)
+    public async Task<FormattingContext> WithTextAsync(SourceText changedText, CancellationToken cancellationToken)
     {
         var changedSnapshot = OriginalSnapshot.WithText(changedText);
 
-        var codeDocument = await _codeDocumentProvider.GetCodeDocumentAsync(changedSnapshot).ConfigureAwait(false);
+        var codeDocument = await _codeDocumentProvider
+            .GetCodeDocumentAsync(changedSnapshot, cancellationToken)
+            .ConfigureAwait(false);
 
         DEBUG_ValidateComponents(CodeDocument, codeDocument);
 

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/IFormattingCodeDocumentProvider.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.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -9,5 +10,5 @@ namespace Microsoft.CodeAnalysis.Razor.Formatting;
 
 internal interface IFormattingCodeDocumentProvider
 {
-    Task<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot);
+    ValueTask<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/Passes/CSharpFormattingPass.cs

@@ -37,7 +37,7 @@ protected async override Task<ImmutableArray<TextChange>> ExecuteCoreAsync(Forma
         if (changes.Length > 0)
         {
             changedText = changedText.WithChanges(changes);
-            changedContext = await context.WithTextAsync(changedText).ConfigureAwait(false);
+            changedContext = await context.WithTextAsync(changedText, cancellationToken).ConfigureAwait(false);
         }
 
         cancellationToken.ThrowIfCancellationRequested();
@@ -48,7 +48,7 @@ protected async override Task<ImmutableArray<TextChange>> ExecuteCoreAsync(Forma
         if (csharpChanges.Length > 0)
         {
             changedText = changedText.WithChanges(csharpChanges);
-            changedContext = await changedContext.WithTextAsync(changedText).ConfigureAwait(false);
+            changedContext = await changedContext.WithTextAsync(changedText, cancellationToken).ConfigureAwait(false);
 
             _logger.LogTestOnly($"After FormatCSharpAsync:\r\n{changedText}");
         }