Optimize CodeLens providers

Removing the timer and unnecessary allocations.

Author: andyleejordan

src/PowerShellEditorServices/Services/CodeLens/ICodeLensProvider.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
@@ -27,8 +28,8 @@ internal interface ICodeLensProvider
         /// The document for which CodeLenses should be provided.
         /// </param>
         /// <param name="cancellationToken"></param>
-        /// <returns>An array of CodeLenses.</returns>
-        CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken);
+        /// <returns>An IEnumerable of CodeLenses.</returns>
+        IEnumerable<CodeLens> ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken);
 
         /// <summary>
         /// Resolves a CodeLens that was created without a Command.

src/PowerShellEditorServices/Services/CodeLens/PesterCodeLensProvider.cs

@@ -1,6 +1,5 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
-using System;
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
@@ -99,18 +98,20 @@ private static CodeLens[] GetPesterLens(PesterSymbolReference pesterSymbol, Scri
         /// <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, CancellationToken cancellationToken)
+        public IEnumerable<CodeLens> ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
         {
             // Don't return anything if codelens setting is disabled
             if (!_configurationService.CurrentSettings.Pester.CodeLens)
             {
-                return Array.Empty<CodeLens>();
+                yield break;
             }
 
-            List<CodeLens> lenses = new();
             foreach (SymbolReference symbol in _symbolProvider.ProvideDocumentSymbols(scriptFile))
             {
-                cancellationToken.ThrowIfCancellationRequested();
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    yield break;
+                }
 
                 if (symbol is not PesterSymbolReference pesterSymbol)
                 {
@@ -129,10 +130,11 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken can
                     continue;
                 }
 
-                lenses.AddRange(GetPesterLens(pesterSymbol, scriptFile));
+                foreach (CodeLens codeLens in GetPesterLens(pesterSymbol, scriptFile))
+                {
+                    yield return codeLens;
+                }
             }
-
-            return lenses.ToArray();
         }
 
         /// <summary>

src/PowerShellEditorServices/Services/CodeLens/ReferencesCodeLensProvider.cs

@@ -53,33 +53,34 @@ public ReferencesCodeLensProvider(WorkspaceService workspaceService, SymbolsServ
         /// </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, classes and enums in the given script file.</returns>
-        public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
+        /// <returns>An IEnumerable of CodeLenses describing all functions, classes and enums in the given script file.</returns>
+        public IEnumerable<CodeLens> ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
         {
-            List<CodeLens> acc = new();
             foreach (SymbolReference symbol in _symbolProvider.ProvideDocumentSymbols(scriptFile))
             {
-                cancellationToken.ThrowIfCancellationRequested();
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    yield break;
+                }
+
                 // TODO: Can we support more here?
                 if (symbol.IsDeclaration &&
                     symbol.Type is
                     SymbolType.Function or
                     SymbolType.Class or
                     SymbolType.Enum)
                 {
-                    acc.Add(new CodeLens
+                    yield return new CodeLens
                     {
                         Data = JToken.FromObject(new
                         {
                             Uri = scriptFile.DocumentUri,
                             ProviderId = nameof(ReferencesCodeLensProvider)
                         }, LspSerializer.Instance.JsonSerializer),
                         Range = symbol.NameRegion.ToRange(),
-                    });
+                    };
                 }
             }
-
-            return acc.ToArray();
         }
 
         /// <summary>

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

@@ -3,13 +3,11 @@
 
 using System;
 using System.Collections.Generic;
-using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
 using Microsoft.PowerShell.EditorServices.CodeLenses;
-using Microsoft.PowerShell.EditorServices.Logging;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
 using Microsoft.PowerShell.EditorServices.Utility;
@@ -21,6 +19,7 @@ namespace Microsoft.PowerShell.EditorServices.Handlers
 {
     internal class PsesCodeLensHandlers : CodeLensHandlerBase
     {
+        private static readonly CodeLensContainer s_emptyCodeLensContainer = new();
         private readonly ILogger _logger;
         private readonly SymbolsService _symbolsService;
         private readonly WorkspaceService _workspaceService;
@@ -40,12 +39,17 @@ public PsesCodeLensHandlers(ILoggerFactory factory, SymbolsService symbolsServic
 
         public override Task<CodeLensContainer> Handle(CodeLensParams request, CancellationToken cancellationToken)
         {
+            _logger.LogDebug($"Handling code lens request for {request.TextDocument.Uri}");
+
             ScriptFile scriptFile = _workspaceService.GetFile(request.TextDocument.Uri);
-            CodeLens[] codeLensResults = ProvideCodeLenses(scriptFile, cancellationToken);
-            return Task.FromResult(new CodeLensContainer(codeLensResults));
+            IEnumerable<CodeLens> codeLensResults = ProvideCodeLenses(scriptFile, cancellationToken);
+
+            return !codeLensResults.Any()
+                ? Task.FromResult(s_emptyCodeLensContainer)
+                : Task.FromResult(new CodeLensContainer(codeLensResults));
         }
 
-        public override async Task<CodeLens> Handle(CodeLens request, CancellationToken cancellationToken)
+        public override Task<CodeLens> Handle(CodeLens request, CancellationToken cancellationToken)
         {
             // TODO: Catch deserialization exception on bad object
             CodeLensData codeLensData = request.Data.ToObject<CodeLensData>();
@@ -55,8 +59,7 @@ public override async Task<CodeLens> Handle(CodeLens request, CancellationToken
                 .FirstOrDefault(provider => provider.ProviderId.Equals(codeLensData.ProviderId, StringComparison.Ordinal));
 
             ScriptFile scriptFile = _workspaceService.GetFile(codeLensData.Uri);
-            return await originalProvider.ResolveCodeLens(request, scriptFile, cancellationToken)
-                .ConfigureAwait(false);
+            return originalProvider.ResolveCodeLens(request, scriptFile, cancellationToken);
         }
 
         /// <summary>
@@ -65,43 +68,20 @@ public override async Task<CodeLens> Handle(CodeLens request, CancellationToken
         /// <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, CancellationToken cancellationToken)
-        {
-            return InvokeProviders(provider => provider.ProvideCodeLenses(scriptFile, cancellationToken))
-                .SelectMany(codeLens => codeLens)
-                .ToArray();
-        }
-
-        /// <summary>
-        /// Invokes the given function synchronously against all
-        /// registered providers.
-        /// </summary>
-        /// <param name="invokeFunc">The function to be invoked.</param>
-        /// <returns>
-        /// An IEnumerable containing the results of all providers
-        /// that were invoked successfully.
-        /// </returns>
-        private IEnumerable<TResult> InvokeProviders<TResult>(Func<ICodeLensProvider, TResult> invokeFunc)
+        private IEnumerable<CodeLens> ProvideCodeLenses(ScriptFile scriptFile, CancellationToken cancellationToken)
         {
-            Stopwatch invokeTimer = new();
-            List<TResult> providerResults = new();
-
             foreach (ICodeLensProvider provider in _symbolsService.GetCodeLensProviders())
             {
-                try
-                {
-                    invokeTimer.Restart();
-                    providerResults.Add(invokeFunc(provider));
-                    invokeTimer.Stop();
-                    _logger.LogTrace($"Invocation of provider '{provider.GetType().Name}' completed in {invokeTimer.ElapsedMilliseconds}ms.");
-                }
-                catch (Exception e)
+                foreach (CodeLens codeLens in provider.ProvideCodeLenses(scriptFile, cancellationToken))
                 {
-                    _logger.LogException($"Exception caught while invoking provider {provider.GetType().Name}:", e);
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        yield break;
+                    }
+
+                    yield return codeLens;
                 }
             }
-
-            return providerResults;
         }
     }
 }