Cleanup the code we have to run copilot code in parallel (#78324)

Author: CyrusNajmabadi

src/Features/ExternalAccess/Copilot/Internal/Completion/CSharpContextProviderService.cs

@@ -6,80 +6,34 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Composition;
-using System.Linq;
-using System.Runtime.CompilerServices;
 using System.Threading;
-using System.Threading.Tasks;
-using Microsoft.CodeAnalysis;
-using Microsoft.CodeAnalysis.ErrorReporting;
 using Microsoft.CodeAnalysis.ExternalAccess.Copilot.Completion;
 using Microsoft.CodeAnalysis.Host.Mef;
-using Microsoft.VisualStudio.Threading;
+using Microsoft.CodeAnalysis.Shared.Utilities;
 
 namespace Microsoft.CodeAnalysis.ExternalAccess.Copilot.Internal.Completion;
 
-[Shared]
-[Export(typeof(ICSharpCopilotContextProviderService))]
-internal sealed class CSharpContextProviderService : ICSharpCopilotContextProviderService
+[Export(typeof(ICSharpCopilotContextProviderService)), Shared]
+[method: ImportingConstructor]
+[method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+internal sealed class CSharpContextProviderService([ImportMany] IEnumerable<IContextProvider> providers)
+    : ICSharpCopilotContextProviderService
 {
-    // Exposed for testing
-    public ImmutableArray<IContextProvider> Providers { get; }
-
-    [ImportingConstructor]
-    [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
-    public CSharpContextProviderService([ImportMany] IEnumerable<IContextProvider> providers)
-    {
-        Providers = providers.ToImmutableArray();
-    }
-
-    public async IAsyncEnumerable<IContextItem> GetContextItemsAsync(Document document, int position, IReadOnlyDictionary<string, object> activeExperiments, [EnumeratorCancellation] CancellationToken cancellationToken)
-    {
-        var queue = new AsyncQueue<IContextItem>();
-        var tasks = this.Providers.Select(provider => Task.Run(async () =>
-        {
-            try
-            {
-                await provider.ProvideContextItemsAsync(document, position, activeExperiments, ProvideItemsAsync, cancellationToken).ConfigureAwait(false);
-            }
-            catch (Exception exception) when (FatalError.ReportAndCatchUnlessCanceled(exception, ErrorSeverity.General))
-            {
-            }
-        },
-        cancellationToken));
-
-        // Let all providers run in parallel in the background, so we can steam results as they come in.
-        // Complete the queue when all providers are done.
-        _ = Task.WhenAll(tasks)
-            .ContinueWith((_, __) => queue.Complete(),
-                          null,
-                          cancellationToken,
-                          TaskContinuationOptions.ExecuteSynchronously,
-                          TaskScheduler.Default);
-
-        while (true)
-        {
-            IContextItem item;
-            try
-            {
-                item = await queue.DequeueAsync(cancellationToken).ConfigureAwait(false);
-            }
-            catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
-            {
-                // Dequeue is cancelled because the queue is empty and completed, we can break out of the loop.
-                break;
-            }
-
-            yield return item;
-        }
-
-        ValueTask ProvideItemsAsync(ImmutableArray<IContextItem> items, CancellationToken cancellationToken)
-        {
-            foreach (var item in items)
-            {
-                queue.Enqueue(item);
-            }
-
-            return default;
-        }
-    }
+    private readonly ImmutableArray<IContextProvider> _providers = [.. providers];
+
+    public IAsyncEnumerable<IContextItem> GetContextItemsAsync(Document document, int position, IReadOnlyDictionary<string, object> activeExperiments, CancellationToken cancellationToken)
+        => ProducerConsumer<IContextItem>.RunParallelStreamAsync(
+            _providers,
+            static async (provider, callback, args, cancellationToken) =>
+                await provider.ProvideContextItemsAsync(
+                    args.document, args.position, args.activeExperiments,
+                    (items, cancellationToken) =>
+                    {
+                        foreach (var item in items)
+                            callback(item);
+
+                        return default;
+                    }, cancellationToken).ConfigureAwait(false),
+            args: (document, position, activeExperiments),
+            cancellationToken);
 }

src/Features/ExternalAccess/Copilot/Microsoft.CodeAnalysis.ExternalAccess.Copilot.csproj

@@ -30,10 +30,6 @@
     <ProjectReference Include="..\..\Core\Portable\Microsoft.CodeAnalysis.Features.csproj" />
   </ItemGroup>
 
-  <ItemGroup>
-    <PackageReference Include="Microsoft.VisualStudio.Threading" />
-  </ItemGroup>
-
   <ItemGroup>
     <PublicAPI Include="PublicAPI.Shipped.txt" />
     <PublicAPI Include="PublicAPI.Unshipped.txt" />

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/LanguageServer/Handler/CopilotCompletion/CopilotCompletionResolveContextHandler.cs

@@ -2,34 +2,29 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using System.Collections.Immutable;
 using System.Composition;
 using Microsoft.CodeAnalysis.ExternalAccess.Copilot.Completion;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Text;
 using Roslyn.LanguageServer.Protocol;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.Handler.Copilot;
 
-[Shared]
 [Method(MethodName)]
-[ExportCSharpVisualBasicStatelessLspService(typeof(CopilotCompletionResolveContextHandler), WellKnownLspServerKinds.Any)]
-internal sealed class CopilotCompletionResolveContextHandler : ILspServiceDocumentRequestHandler<ContextResolveParam, IContextItem[]>
+[ExportCSharpVisualBasicStatelessLspService(typeof(CopilotCompletionResolveContextHandler), WellKnownLspServerKinds.Any), Shared]
+[method: ImportingConstructor]
+[method: Obsolete("This exported object must be obtained through the MEF export provider.", error: true)]
+internal sealed class CopilotCompletionResolveContextHandler(ICSharpCopilotContextProviderService contextProviderService)
+    : ILspServiceDocumentRequestHandler<ContextResolveParam, IContextItem[]>
 {
     // "@2" prefix to differentiate it from the implementation previously located in devkit extension.
     private const string MethodName = "roslyn/resolveContext@2";
 
-    [ImportingConstructor]
-    [Obsolete("This exported object must be obtained through the MEF export provider.", error: true)]
-    public CopilotCompletionResolveContextHandler(ICSharpCopilotContextProviderService contextProviderService)
-    {
-        ContextProviderService = contextProviderService;
-    }
-
     public bool MutatesSolutionState => false;
 
     public bool RequiresLSPSolution => true;
 
-    public ICSharpCopilotContextProviderService ContextProviderService { get; }
+    public ICSharpCopilotContextProviderService ContextProviderService { get; } = contextProviderService;
 
     public TextDocumentIdentifier GetTextDocumentIdentifier(ContextResolveParam request)
         => request.DocumentContext.TextDocument;
@@ -41,13 +36,11 @@ public async Task<IContextItem[]> HandleRequestAsync(ContextResolveParam param,
 
         var text = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);
         var position = text.Lines.GetPosition(linePosition);
-        var builder = ImmutableArray.CreateBuilder<IContextItem>();
         var activeExperiments = param.GetUnpackedActiveExperiments();
 
+        using var _ = ArrayBuilder<IContextItem>.GetInstance(out var builder);
         await foreach (var item in ContextProviderService.GetContextItemsAsync(document, position, activeExperiments, cancellationToken).ConfigureAwait(false))
-        {
             builder.Add(item);
-        }
 
         return builder.ToArray();
     }

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Utilities/ProducerConsumer.cs

@@ -281,6 +281,28 @@ public static async IAsyncEnumerable<TItem> RunAsync<TArgs>(
             yield return item;
     }
 
+    /// <summary>
+    /// Equivalent to <see cref="RunParallelAsync{TSource, TArgs}(IEnumerable{TSource}, Func{TSource, Action{TItem}, TArgs, CancellationToken, Task}, TArgs, CancellationToken)"/>,
+    /// but returns value as an <see cref="IAsyncEnumerable{TItem}"/>.  Versus an <see cref="ImmutableArray{TItem}"/>.  
+    /// This is useful for cases where the caller wants to stream over the results as they are produced, rather than
+    /// waiting on the full set to be produced before processing them.
+    /// </summary>
+    public static IAsyncEnumerable<TItem> RunParallelStreamAsync<TSource, TArgs>(
+        IEnumerable<TSource> source,
+        Func<TSource, Action<TItem>, TArgs, CancellationToken, Task> produceItems,
+        TArgs args,
+        CancellationToken cancellationToken)
+    {
+        return RunAsync(
+            static (callback, args, cancellationToken) =>
+                RoslynParallel.ForEachAsync(
+                    args.source, cancellationToken,
+                    async (source, cancellationToken) => await args.produceItems(
+                        source, callback, args.args, cancellationToken).ConfigureAwait(false)),
+            args: (source, produceItems, args),
+            cancellationToken);
+    }
+
     /// <summary>
     /// Helper utility for the pattern of a pair of a production routine and consumption routine using a channel to
     /// coordinate data transfer.  The provided <paramref name="options"/> are used to create a <see