Use async foreach to consume IAsyncEnumerable

- Flow CancellationToken to IAsyncEnumerable.GetAsyncEnumerator

Author: halter73

src/SignalR/server/Core/src/Internal/AsyncEnumerableAdapters.cs

@@ -0,0 +1,135 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Threading;
+using System.Threading.Channels;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.SignalR.Internal
+{
+    // True-internal because this is a weird and tricky class to use :)
+    internal static class AsyncEnumerableAdapters
+    {
+        public static IAsyncEnumerable<object> MakeCancelableAsyncEnumerable<T>(IAsyncEnumerable<T> asyncEnumerable, CancellationToken cancellationToken = default)
+        {
+            return new CancelableAsyncEnumerable<T>(asyncEnumerable, cancellationToken);
+        }
+
+        public static IAsyncEnumerable<object> GetAsyncEnumerableFromChannel<T>(ChannelReader<T> channel, CancellationToken cancellationToken = default)
+        {
+            return new ChannelAsyncEnumerable<T>(channel, cancellationToken);
+        }
+
+        /// <summary>Converts an IAsyncEnumerable of T to an IAsyncEnumerable of object.</summary>
+        private class CancelableAsyncEnumerable<T> : IAsyncEnumerable<object>
+        {
+            private readonly IAsyncEnumerable<T> _asyncEnumerable;
+            private readonly CancellationToken _cancellationToken;
+
+            public CancelableAsyncEnumerable(IAsyncEnumerable<T> asyncEnumerable, CancellationToken cancellationToken)
+            {
+                _asyncEnumerable = asyncEnumerable;
+                _cancellationToken = cancellationToken;
+            }
+
+            public IAsyncEnumerator<object> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+            {
+                Debug.Assert(cancellationToken == default);
+                return new CancelableAsyncEnumerator(_asyncEnumerable.GetAsyncEnumerator(_cancellationToken));
+            }
+
+            private class CancelableAsyncEnumerator : IAsyncEnumerator<object>
+            {
+                private IAsyncEnumerator<T> _asyncEnumerator;
+
+                public CancelableAsyncEnumerator(IAsyncEnumerator<T> asyncEnumerator)
+                {
+                    _asyncEnumerator = asyncEnumerator;
+                }
+
+                public object Current => _asyncEnumerator.Current;
+
+                public ValueTask<bool> MoveNextAsync()
+                {
+                    return _asyncEnumerator.MoveNextAsync();
+                }
+
+                public ValueTask DisposeAsync()
+                {
+                    return _asyncEnumerator.DisposeAsync();
+                }
+            }
+        }
+
+        /// <summary>Provides an IAsyncEnumerable of object for the data in a channel.</summary>
+        private class ChannelAsyncEnumerable<T> : IAsyncEnumerable<object>
+        {
+            private readonly ChannelReader<T> _channel;
+            private readonly CancellationToken _cancellationToken;
+
+            public ChannelAsyncEnumerable(ChannelReader<T> channel, CancellationToken cancellationToken)
+            {
+                _channel = channel;
+                _cancellationToken = cancellationToken;
+            }
+
+            public IAsyncEnumerator<object> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+            {
+                Debug.Assert(cancellationToken == default);
+                return new ChannelAsyncEnumerator(_channel, _cancellationToken);
+            }
+
+            private class ChannelAsyncEnumerator : IAsyncEnumerator<object>
+            {
+                /// <summary>The channel being enumerated.</summary>
+                private readonly ChannelReader<T> _channel;
+                /// <summary>Cancellation token used to cancel the enumeration.</summary>
+                private readonly CancellationToken _cancellationToken;
+                /// <summary>The current element of the enumeration.</summary>
+                private T _current;
+
+                public ChannelAsyncEnumerator(ChannelReader<T> channel, CancellationToken cancellationToken)
+                {
+                    _channel = channel;
+                    _cancellationToken = cancellationToken;
+                }
+
+                public object Current => _current;
+
+                public ValueTask<bool> MoveNextAsync()
+                {
+                    var result = _channel.ReadAsync(_cancellationToken);
+
+                    if (result.IsCompletedSuccessfully)
+                    {
+                        _current = result.Result;
+                        return new ValueTask<bool>(true);
+                    }
+
+                    return new ValueTask<bool>(MoveNextAsyncAwaited(result));
+                }
+
+                private async Task<bool> MoveNextAsyncAwaited(ValueTask<T> channelReadTask)
+                {
+                    try
+                    {
+                        _current = await channelReadTask;
+                    }
+                    catch (ChannelClosedException ex) when (ex.InnerException == null)
+                    {
+                        return false;
+                    }
+
+                    return true;
+                }
+
+                public ValueTask DisposeAsync()
+                {
+                    return default;
+                }
+            }
+        }
+    }
+}

src/SignalR/server/Core/src/Internal/AsyncEnumeratorAdapters.cs

@@ -263,16 +263,20 @@ await SendInvocationError(hubMethodInvocationMessage.InvocationId, connection,
                     {
                         var result = await ExecuteHubMethod(methodExecutor, hub, arguments);
 
-                        if (!TryGetStreamingEnumerator(connection, hubMethodInvocationMessage.InvocationId, descriptor, result, out var enumerator, ref cts))
+                        if (result == null)
                         {
                             Log.InvalidReturnValueFromStreamingMethod(_logger, methodExecutor.MethodInfo.Name);
                             await SendInvocationError(hubMethodInvocationMessage.InvocationId, connection,
-                                $"The value returned by the streaming method '{methodExecutor.MethodInfo.Name}' is not a ChannelReader<>.");
+                                $"The value returned by the streaming method '{methodExecutor.MethodInfo.Name}' is not a ChannelReader<> or IAsyncEnumerable<>.");
                             return;
                         }
 
+                        cts = cts ?? CancellationTokenSource.CreateLinkedTokenSource(connection.ConnectionAborted);
+                        connection.ActiveRequestCancellationSources.TryAdd(hubMethodInvocationMessage.InvocationId, cts);
+                        var enumerable = descriptor.FromReturnedStream(result, cts.Token);
+
                         Log.StreamingResult(_logger, hubMethodInvocationMessage.InvocationId, methodExecutor);
-                        _ = StreamResultsAsync(hubMethodInvocationMessage.InvocationId, connection, enumerator, scope, hubActivator, hub, cts);
+                        _ = StreamResultsAsync(hubMethodInvocationMessage.InvocationId, connection, enumerable, scope, hubActivator, hub, cts);
                     }
 
                     else if (string.IsNullOrEmpty(hubMethodInvocationMessage.InvocationId))
@@ -358,7 +362,7 @@ await SendInvocationError(hubMethodInvocationMessage.InvocationId, connection,
             }
         }
 
-        private async Task StreamResultsAsync(string invocationId, HubConnectionContext connection, IAsyncEnumerator<object> enumerator, IServiceScope scope,
+        private async Task StreamResultsAsync(string invocationId, HubConnectionContext connection, IAsyncEnumerable<object> enumerable, IServiceScope scope,
             IHubActivator<THub> hubActivator, THub hub, CancellationTokenSource streamCts)
         {
             string error = null;
@@ -367,10 +371,11 @@ private async Task StreamResultsAsync(string invocationId, HubConnectionContext
             {
                 try
                 {
-                    while (await enumerator.MoveNextAsync())
+
+                    await foreach (var streamItem in enumerable)
                     {
                         // Send the stream item
-                        await connection.WriteAsync(new StreamItemMessage(invocationId, enumerator.Current));
+                        await connection.WriteAsync(new StreamItemMessage(invocationId, streamItem));
                     }
                 }
                 catch (ChannelClosedException ex)
@@ -389,8 +394,6 @@ private async Task StreamResultsAsync(string invocationId, HubConnectionContext
                 }
                 finally
                 {
-                    await enumerator.DisposeAsync();
-
                     hubActivator.Release(hub);
 
                     // Dispose the linked CTS for the stream.
@@ -470,10 +473,10 @@ private static async Task<bool> IsHubMethodAuthorizedSlow(IServiceProvider provi
             return authorizationResult.Succeeded;
         }
 
-        private async Task<bool> ValidateInvocationMode(HubMethodDescriptor hubMethodDescriptor, bool isStreamedInvocation,
+        private async Task<bool> ValidateInvocationMode(HubMethodDescriptor hubMethodDescriptor, bool isStreamResponse,
             HubMethodInvocationMessage hubMethodInvocationMessage, HubConnectionContext connection)
         {
-            if (hubMethodDescriptor.IsStreamable && !isStreamedInvocation)
+            if (hubMethodDescriptor.IsStreamResponse && !isStreamResponse)
             {
                 // Non-null/empty InvocationId? Blocking
                 if (!string.IsNullOrEmpty(hubMethodInvocationMessage.InvocationId))
@@ -486,7 +489,7 @@ await connection.WriteAsync(CompletionMessage.WithError(hubMethodInvocationMessa
                 return false;
             }
 
-            if (!hubMethodDescriptor.IsStreamable && isStreamedInvocation)
+            if (!hubMethodDescriptor.IsStreamResponse && isStreamResponse)
             {
                 Log.NonStreamingMethodCalledWithStream(_logger, hubMethodInvocationMessage);
                 await connection.WriteAsync(CompletionMessage.WithError(hubMethodInvocationMessage.InvocationId,
@@ -498,26 +501,6 @@ await connection.WriteAsync(CompletionMessage.WithError(hubMethodInvocationMessa
             return true;
         }
 
-        private bool TryGetStreamingEnumerator(HubConnectionContext connection, string invocationId, HubMethodDescriptor hubMethodDescriptor, object result, out IAsyncEnumerator<object> enumerator, ref CancellationTokenSource streamCts)
-        {
-            if (result != null)
-            {
-                if (hubMethodDescriptor.IsStreamable)
-                {
-                    if (streamCts == null)
-                    {
-                        streamCts = CancellationTokenSource.CreateLinkedTokenSource(connection.ConnectionAborted);
-                    }
-                    connection.ActiveRequestCancellationSources.TryAdd(invocationId, streamCts);
-                    enumerator = hubMethodDescriptor.FromReturnedStream(result, streamCts.Token);
-                    return true;
-                }
-            }
-
-            enumerator = null;
-            return false;
-        }
-
         private void DiscoverHubMethods()
         {
             var hubType = typeof(THub);