Simplify channel-processing code in SemanticSearch. (#78060)

Author: CyrusNajmabadi

src/Features/Core/Portable/SemanticSearch/Tools/ReferencingSyntaxFinder.cs

@@ -2,19 +2,17 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-#if NET6_0_OR_GREATER
+#if NET
 
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Threading;
-using System.Threading.Channels;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.FindSymbols;
-using Microsoft.CodeAnalysis.FindUsages;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.Utilities;
-using Microsoft.CodeAnalysis.Threading;
 
 namespace Microsoft.CodeAnalysis.SemanticSearch;
 
@@ -34,85 +32,39 @@ public IEnumerable<SyntaxNode> Find(ISymbol symbol)
 
     public async IAsyncEnumerable<SyntaxNode> FindAsync(ISymbol symbol)
     {
-        var channel = Channel.CreateUnbounded<ReferenceLocation>(new()
+        using var _ = PooledHashSet<SyntaxNode>.GetInstance(out var cachedRoots);
+
+        // Kick off the SymbolFinder.FindReferencesAsync call on the provided symbol/solution.  As it finds
+        // ReferenceLocations, it will push those into the 'callback' delegate passed into it. ProducerConsumer will
+        // then convert this to a simple IAsyncEnumerable<ReferenceLocation> that we can iterate over, converting those
+        // locations to SyntaxNodes in the corresponding C# or VB document.
+        await foreach (var item in ProducerConsumer<ReferenceLocation>.RunAsync(
+            FindReferencesAsync, args: (solution, symbol), cancellationToken))
         {
-            SingleReader = true,
-            SingleWriter = false,
-        });
-
-        using var _ = cancellationToken.Register(
-            static (obj, cancellationToken) => ((Channel<SourceReferenceItem>)obj!).Writer.TryComplete(new OperationCanceledException(cancellationToken)),
-            state: channel);
-
-        var progress = new Progress(channel);
-
-        var writeTask = ProduceItemsAndWriteToChannelAsync();
-
-        await foreach (var reference in channel.Reader.ReadAllAsync(cancellationToken))
-        {
-            // TODO: consider grouping by document to avoid repeated syntax root lookup
-
-            var root = await reference.Document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
+            var root = await item.Document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
             if (root == null)
-            {
                 continue;
-            }
-
-            yield return root.FindNode(reference.Location.SourceSpan, findInTrivia: true, getInnermostNodeForTie: true);
-        }
-
-        await writeTask.ConfigureAwait(false);
-
-        async Task ProduceItemsAndWriteToChannelAsync()
-        {
-            await Task.Yield().ConfigureAwait(false);
 
-            Exception? exception = null;
-            try
-            {
-                //await service.FindReferencesAsync(context, document, location.SourceSpan.Start, classificationOptions, cancellationToken).ConfigureAwait(false);
-                await SymbolFinder.FindReferencesAsync(
-                    symbol,
-                    solution,
-                    progress,
-                    documents: null,
-                    s_options,
-                    cancellationToken).ConfigureAwait(false);
-            }
-            catch (Exception ex) when ((exception = ex) == null)
-            {
-                throw ExceptionUtilities.Unreachable();
-            }
-            finally
-            {
-                // No matter what path we take (exceptional or non-exceptional), always complete the channel so the
-                // writing task knows it's done.
-                channel.Writer.TryComplete(exception);
-            }
+            // Hold onto the root so that if we find more references in the same document, we don't have to reparse it.
+            cachedRoots.Add(root);
+            yield return item.Location.FindNode(findInsideTrivia: true, getInnermostNodeForTie: true, cancellationToken);
         }
     }
 
-    private sealed class Progress(Channel<ReferenceLocation> channel) : IStreamingFindReferencesProgress
-    {
-        public ValueTask OnStartedAsync(CancellationToken cancellationToken)
-            => ValueTask.CompletedTask;
-
-        public ValueTask OnCompletedAsync(CancellationToken cancellationToken)
-            => ValueTask.CompletedTask;
+    private static Task FindReferencesAsync(Action<ReferenceLocation> callback, (Solution solution, ISymbol symbol) args, CancellationToken cancellationToken)
+        => SymbolFinder.FindReferencesAsync(
+            args.symbol, args.solution, new Progress(callback), documents: null, s_options, cancellationToken);
 
-        public ValueTask OnDefinitionFoundAsync(SymbolGroup group, CancellationToken cancellationToken)
-            => ValueTask.CompletedTask;
+    private sealed class Progress(Action<ReferenceLocation> callback) : IStreamingFindReferencesProgress
+    {
+        public ValueTask OnStartedAsync(CancellationToken cancellationToken) => ValueTask.CompletedTask;
+        public ValueTask OnCompletedAsync(CancellationToken cancellationToken) => ValueTask.CompletedTask;
+        public ValueTask OnDefinitionFoundAsync(SymbolGroup group, CancellationToken cancellationToken) => ValueTask.CompletedTask;
 
         public ValueTask OnReferencesFoundAsync(ImmutableArray<(SymbolGroup group, ISymbol symbol, ReferenceLocation location)> references, CancellationToken cancellationToken)
         {
             foreach (var (_, _, location) in references)
-            {
-                // It's ok to use TryWrite here.  TryWrite always succeeds unless the channel is completed. And the
-                // channel is only ever completed by us (after produceItems completes or throws an exception) or if the
-                // cancellationToken is triggered above in RunAsync. In that latter case, it's ok for writing to the
-                // channel to do nothing as we no longer need to write out those assets to the pipe.
-                _ = channel.Writer.TryWrite(location);
-            }
+                callback(location);
 
             return ValueTask.CompletedTask;
         }

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

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Channels;
 using System.Threading.Tasks;
@@ -258,6 +259,28 @@ await args.produceItems(source, callback, args.args, cancellationToken).Configur
             cancellationToken);
     }
 
+    /// <summary>
+    /// Version of <see cref="RunChannelAsync"/> when caller the prefers to just push all the results into a channel
+    /// that it receives in the return value to process asynchronously.
+    /// </summary>
+    public static async IAsyncEnumerable<TItem> RunAsync<TArgs>(
+        Func<Action<TItem>, TArgs, CancellationToken, Task> produceItems,
+        TArgs args,
+        [EnumeratorCancellation] CancellationToken cancellationToken)
+    {
+        var channelReader = await RunChannelAsync(
+            // We're the only reader (in the foreach loop below).  So we can use the single reader options.
+            ProducerConsumerOptions.SingleReaderOptions,
+            produceItems,
+            // Trivially grab the reader and return it.  We don't need to do any processing of the values, as the
+            // callers just wants them as is.
+            consumeItems: static (reader, _, _) => Task.FromResult(reader),
+            args, cancellationToken).ConfigureAwait(false);
+
+        await foreach (var item in channelReader.ReadAllAsync(cancellationToken))
+            yield return item;
+    }
+
     /// <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