Make a lot more things cancellable

- Cancellation tokens are now propagated to all code lens providers.
- Provider code that was heavy in synchronous code has some cancellation
  checks and `Task.Yield` calls to help keep thread pool threads more
  flexible.

Author: SeeminglyScience

src/PowerShellEditorServices/Services/CodeLens/ICodeLensProvider.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
@@ -25,8 +26,9 @@ internal interface ICodeLensProvider
         /// <param name="scriptFile">
         /// The document for which CodeLenses should be provided.
         /// </param>
+        /// <param name="cancellationToken"></param>
         /// <returns>An array of CodeLenses.</returns>
-        CodeLens[] ProvideCodeLenses(ScriptFile scriptFile);
+        CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken);
 
         /// <summary>
         /// Resolves a CodeLens that was created without a Command.
@@ -37,11 +39,13 @@ internal interface ICodeLensProvider
         /// <param name="scriptFile">
         /// The ScriptFile to resolve it in (sometimes unused).
         /// </param>
+        /// <param name="cancellationToken"></param>
         /// <returns>
         /// A Task which returns the resolved CodeLens when completed.
         /// </returns>
         Task<CodeLens> ResolveCodeLens(
             CodeLens codeLens,
-            ScriptFile scriptFile);
+            ScriptFile scriptFile,
+            CancellationToken cancellationToken);
     }
 }

src/PowerShellEditorServices/Services/CodeLens/PesterCodeLensProvider.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 using System;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Services.Symbols;
@@ -99,8 +100,9 @@ private static CodeLens[] GetPesterLens(PesterSymbolReference pesterSymbol, Scri
         /// Get all Pester CodeLenses for a given script file.
         /// </summary>
         /// <param name="scriptFile">The script file to get Pester CodeLenses for.</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>All Pester CodeLenses for the given script file.</returns>
-        public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
+        public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
         {
             // Don't return anything if codelens setting is disabled
             if (!_configurationService.CurrentSettings.Pester.CodeLens)
@@ -116,6 +118,7 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
                     continue;
                 }
 
+                cancellationToken.ThrowIfCancellationRequested();
                 if (_configurationService.CurrentSettings.Pester.UseLegacyCodeLens
                         && pesterSymbol.Command != PesterCommandType.Describe)
                 {
@@ -133,8 +136,9 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
         /// </summary>
         /// <param name="codeLens">The code lens to resolve.</param>
         /// <param name="scriptFile">The script file.</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>The given CodeLens, wrapped in a task.</returns>
-        public Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile) =>
+        public Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile, CancellationToken cancellationToken) =>
             // This provider has no specific behavior for
             // resolving CodeLenses.
             Task.FromResult(codeLens);

src/PowerShellEditorServices/Services/CodeLens/ReferencesCodeLensProvider.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Services.Symbols;
@@ -53,12 +54,14 @@ public ReferencesCodeLensProvider(WorkspaceService workspaceService, SymbolsServ
         /// Get all reference code lenses for a given script file.
         /// </summary>
         /// <param name="scriptFile">The PowerShell script file to get code lenses for.</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>An array of CodeLenses describing all functions in the given script file.</returns>
-        public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
+        public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
         {
             List<CodeLens> acc = new();
             foreach (SymbolReference sym in _symbolProvider.ProvideDocumentSymbols(scriptFile))
             {
+                cancellationToken.ThrowIfCancellationRequested();
                 if (sym.SymbolType == SymbolType.Function)
                 {
                     acc.Add(new CodeLens
@@ -68,7 +71,7 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
                             Uri = scriptFile.DocumentUri,
                             ProviderId = nameof(ReferencesCodeLensProvider)
                         }, LspSerializer.Instance.JsonSerializer),
-                        Range = sym.ScriptRegion.ToRange()
+                        Range = sym.ScriptRegion.ToRange(),
                     });
                 }
             }
@@ -81,8 +84,12 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
         /// </summary>
         /// <param name="codeLens">The old code lens to get updated references for.</param>
         /// <param name="scriptFile"></param>
+        /// <param name="cancellationToken"></param>
         /// <returns>A new code lens object describing the same data as the old one but with updated references.</returns>
-        public async Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile)
+        public async Task<CodeLens> ResolveCodeLens(
+            CodeLens codeLens,
+            ScriptFile scriptFile,
+            CancellationToken cancellationToken)
         {
             ScriptFile[] references = _workspaceService.ExpandScriptReferences(
                 scriptFile);
@@ -93,9 +100,10 @@ public async Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile script
                 codeLens.Range.Start.Character + 1);
 
             List<SymbolReference> referencesResult = await _symbolsService.FindReferencesOfSymbol(
-                foundSymbol,
-                references,
-                _workspaceService).ConfigureAwait(false);
+                    foundSymbol,
+                    references,
+                    _workspaceService,
+                    cancellationToken).ConfigureAwait(false);
 
             Location[] referenceLocations;
             if (referencesResult == null)
@@ -107,6 +115,10 @@ public async Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile script
                 List<Location> acc = new();
                 foreach (SymbolReference foundReference in referencesResult)
                 {
+                    // This async method is pretty dense with synchronous code
+                    // so it's helpful to add some yields.
+                    await Task.Yield();
+                    cancellationToken.ThrowIfCancellationRequested();
                     if (IsReferenceDefinition(foundSymbol, foundReference))
                     {
                         continue;
@@ -140,9 +152,9 @@ public async Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile script
                     Title = GetReferenceCountHeader(referenceLocations.Length),
                     Arguments = JArray.FromObject(new object[]
                     {
-                        scriptFile.DocumentUri,
-                        codeLens.Range.Start,
-                        referenceLocations
+                    scriptFile.DocumentUri,
+                    codeLens.Range.Start,
+                    referenceLocations
                     },
                     LspSerializer.Instance.JsonSerializer)
                 }

src/PowerShellEditorServices/Services/Symbols/SymbolsService.cs

@@ -11,6 +11,7 @@
 using System.Management.Automation.Language;
 using System.Runtime.InteropServices;
 using System.Text.RegularExpressions;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using Microsoft.PowerShell.EditorServices.CodeLenses;
@@ -160,18 +161,25 @@ public static SymbolReference FindSymbolAtLocation(
         /// <param name="foundSymbol">The symbol to find all references for</param>
         /// <param name="referencedFiles">An array of scriptFiles too search for references in</param>
         /// <param name="workspace">The workspace that will be searched for symbols</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>FindReferencesResult</returns>
         public async Task<List<SymbolReference>> FindReferencesOfSymbol(
             SymbolReference foundSymbol,
             ScriptFile[] referencedFiles,
-            WorkspaceService workspace)
+            WorkspaceService workspace,
+            CancellationToken cancellationToken = default)
         {
             if (foundSymbol == null)
             {
                 return null;
             }
 
-            (Dictionary<string, List<string>> cmdletToAliases, Dictionary<string, string> aliasToCmdlets) = await CommandHelpers.GetAliasesAsync(_executionService).ConfigureAwait(false);
+            CommandHelpers.AliasMap aliases = await CommandHelpers.GetAliasesAsync(
+                _executionService,
+                cancellationToken).ConfigureAwait(false);
+
+            Dictionary<string, List<string>> cmdletToAliases = aliases.CmdletToAliases;
+            Dictionary<string, string> aliasToCmdlets = aliases.AliasToCmdlets;
 
             // We want to look for references first in referenced files, hence we use ordered dictionary
             // TODO: File system case-sensitivity is based on filesystem not OS, but OS is a much cheaper heuristic
@@ -188,6 +196,10 @@ public async Task<List<SymbolReference>> FindReferencesOfSymbol(
             {
                 if (!fileMap.Contains(filePath))
                 {
+                    // This async method is pretty dense with synchronous code
+                    // so it's helpful to add some yields.
+                    await Task.Yield();
+                    cancellationToken.ThrowIfCancellationRequested();
                     if (!workspace.TryGetFile(filePath, out ScriptFile scriptFile))
                     {
                         // If we can't access the file for some reason, just ignore it
@@ -223,6 +235,9 @@ public async Task<List<SymbolReference>> FindReferencesOfSymbol(
                     reference.FilePath = file.FilePath;
                     symbolReferences.Add(reference);
                 }
+
+                await Task.Yield();
+                cancellationToken.ThrowIfCancellationRequested();
             }
 
             return symbolReferences;

src/PowerShellEditorServices/Services/TextDocument/Handlers/CodeLensHandlers.cs

@@ -41,11 +41,11 @@ public PsesCodeLensHandlers(ILoggerFactory factory, SymbolsService symbolsServic
         public override Task<CodeLensContainer> Handle(CodeLensParams request, CancellationToken cancellationToken)
         {
             ScriptFile scriptFile = _workspaceService.GetFile(request.TextDocument.Uri);
-            CodeLens[] codeLensResults = ProvideCodeLenses(scriptFile);
+            CodeLens[] codeLensResults = ProvideCodeLenses(scriptFile, cancellationToken);
             return Task.FromResult(new CodeLensContainer(codeLensResults));
         }
 
-        public override Task<CodeLens> Handle(CodeLens request, CancellationToken cancellationToken)
+        public override async Task<CodeLens> Handle(CodeLens request, CancellationToken cancellationToken)
         {
             // TODO: Catch deserializtion exception on bad object
             CodeLensData codeLensData = request.Data.ToObject<CodeLensData>();
@@ -55,18 +55,19 @@ public override Task<CodeLens> Handle(CodeLens request, CancellationToken cancel
                 .FirstOrDefault(provider => provider.ProviderId.Equals(codeLensData.ProviderId, StringComparison.Ordinal));
 
             ScriptFile scriptFile = _workspaceService.GetFile(codeLensData.Uri);
-
-            return originalProvider.ResolveCodeLens(request, scriptFile);
+            return await originalProvider.ResolveCodeLens(request, scriptFile, cancellationToken)
+                .ConfigureAwait(false);
         }
 
         /// <summary>
         /// Get all the CodeLenses for a given script file.
         /// </summary>
         /// <param name="scriptFile">The PowerShell script file to get CodeLenses for.</param>
+        /// <param name="cancellationToken"></param>
         /// <returns>All generated CodeLenses for the given script file.</returns>
-        private CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
+        private CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
         {
-            return InvokeProviders(provider => provider.ProvideCodeLenses(scriptFile))
+            return InvokeProviders(provider => provider.ProvideCodeLenses(scriptFile, cancellationToken))
                 .SelectMany(codeLens => codeLens)
                 .ToArray();
         }

src/PowerShellEditorServices/Services/TextDocument/Handlers/ReferencesHandler.cs

@@ -45,7 +45,8 @@ public override async Task<LocationContainer> Handle(ReferenceParams request, Ca
                 await _symbolsService.FindReferencesOfSymbol(
                     foundSymbol,
                     _workspaceService.ExpandScriptReferences(scriptFile),
-                    _workspaceService).ConfigureAwait(false);
+                    _workspaceService,
+                    cancellationToken).ConfigureAwait(false);
 
             List<Location> locations = new();
 

src/PowerShellEditorServices/Services/Workspace/Handlers/WorkspaceSymbolsHandler.cs

@@ -32,7 +32,7 @@ public PsesWorkspaceSymbolsHandler(ILoggerFactory loggerFactory, SymbolsService
 
         protected override WorkspaceSymbolRegistrationOptions CreateRegistrationOptions(WorkspaceSymbolCapability capability, ClientCapabilities clientCapabilities) => new() { };
 
-        public override Task<Container<SymbolInformation>> Handle(WorkspaceSymbolParams request, CancellationToken cancellationToken)
+        public override async Task<Container<SymbolInformation>> Handle(WorkspaceSymbolParams request, CancellationToken cancellationToken)
         {
             List<SymbolInformation> symbols = new();
 
@@ -47,6 +47,10 @@ public override Task<Container<SymbolInformation>> Handle(WorkspaceSymbolParams
 
                 foreach (SymbolReference foundOccurrence in foundSymbols)
                 {
+                    // This async method is pretty dense with synchronous code
+                    // so it's helpful to add some yields.
+                    await Task.Yield();
+                    cancellationToken.ThrowIfCancellationRequested();
                     if (!IsQueryMatch(request.Query, foundOccurrence.SymbolName))
                     {
                         continue;
@@ -68,8 +72,7 @@ public override Task<Container<SymbolInformation>> Handle(WorkspaceSymbolParams
                 }
             }
             _logger.LogWarning("Logging in a handler works now.");
-
-            return Task.FromResult(new Container<SymbolInformation>(symbols));
+            return new Container<SymbolInformation>(symbols);
         }
 
         #region private Methods