Added HandleMany

Toolkit.Foundation/Mediator.cs

@@ -1,5 +1,6 @@
 using Microsoft.Extensions.DependencyInjection;
 using System.Reflection;
+using System.Runtime.CompilerServices;
 
 namespace Toolkit.Foundation;
 
@@ -12,56 +13,37 @@ public class Mediator(IHandlerProvider handlerProvider,
         CancellationToken cancellationToken = default)
         where TMessage : notnull
     {
-        Type messageType = message.GetType();
+        List<object?> handlers = GetHandlers<TMessage, TResponse>(message, key);
-        Type handlerWrapperType = typeof(HandlerWrapper<,>).MakeGenericType(messageType, 
-            typeof(TResponse));
 
-        Dictionary<Type, List<object?>> handlers = [];
+        foreach (object? handler in handlers)
-        foreach (object? handler in key is not null ? provider.GetKeyedServices(handlerWrapperType, key) :
-            provider.GetServices(handlerWrapperType))
         {
-            if (handler?.GetType() is Type serviceType)
+            MethodInfo? handleMethod = handler?.GetType().GetMethod("Handle", [message.GetType(), typeof(CancellationToken)]);
+            if (handleMethod != null)
             {
-                if (!handlers.TryGetValue(serviceType, out List<object?>? handlerList))
+                return await (Task<TResponse?>)handleMethod.Invoke(handler, new object[] { message, cancellationToken })!;
-                {
-                    handlerList = [];
-                    handlers.Add(serviceType, handlerList);
-                }
-
-                handlerList.Add(handler);
-            }
-        }
-
-        foreach (object? handler in handlerProvider.Get(messageType, key))
-        {
-            if (handler is not null)
-            {
-                Type handlerType = handler.GetType();
-                if (!handlers.TryGetValue(handlerType, out List<object?>? handlerList))
-                {
-                    handlerList = [];
-                    handlers.Add(handlerType, handlerList);
-                }
-
-                handlerList.Add(handler);
-            }
-        }
-
-        foreach (KeyValuePair<Type, List<object?>> handlerEntry in handlers)
-        {
-            foreach (object? handler in handlerEntry.Value)
-            {
-                if (handler?.GetType().GetMethod("Handle", [messageType, typeof(CancellationToken)]) is MethodInfo handleMethod)
-                {
-                    return await (Task<TResponse?>)handleMethod.Invoke(handler,
-                        new object[] { message, cancellationToken })!;
-                }
             }
         }
 
         return default;
     }
 
+    public async IAsyncEnumerable<TResponse?> HandleMany<TMessage, TResponse>(TMessage message,
+        object? key = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        where TMessage : notnull
+    {
+        List<object?> handlers = GetHandlers<TMessage, TResponse>(message, key);
+
+        foreach (object? handler in handlers)
+        {
+            MethodInfo? handleMethod = handler?.GetType().GetMethod("Handle", [message.GetType(), typeof(CancellationToken)]);
+            if (handleMethod != null)
+            {
+                yield return await (Task<TResponse?>)handleMethod.Invoke(handler, new object[] { message, cancellationToken })!;
+            }
+        }
+    }
+
     public async Task<object?> Handle(object message,
         object? key = null,
         CancellationToken cancellationToken = default)
@@ -72,31 +54,44 @@ public class Mediator(IHandlerProvider handlerProvider,
             requestType.GetGenericArguments().Length == 1)
         {
             Type responseType = requestType.GetGenericArguments()[0];
-            Type handlerWrapperType = typeof(HandlerWrapper<,>).MakeGenericType(message.GetType(),
+            Type handlerWrapperType = typeof(HandlerWrapper<,>).MakeGenericType(message.GetType(), responseType);
-                responseType);
 
+            List<object?> handlers = GetHandlers(message, handlerWrapperType, key);
+
+            foreach (object? handler in handlers)
+            {
+                MethodInfo? handleMethod = handler?.GetType().GetMethod("Handle", [messageType, typeof(CancellationToken)]);
+                if (handleMethod != null)
+                {
+                    return await (Task<object?>)handleMethod.Invoke(handler, new object[] { message, cancellationToken })!;
+                }
+            }
+        }
+
+        return default;
+    }
+
+    private List<object?> GetHandlers<TMessage, TResponse>(TMessage message, object? key)
+        where TMessage : notnull
+    {
+        Type messageType = message.GetType();
+        Type handlerWrapperType = typeof(HandlerWrapper<,>).MakeGenericType(messageType, typeof(TResponse));
+
+        return GetHandlers(message, handlerWrapperType, key);
+    }
+
+    private List<object?> GetHandlers(object message, Type handlerWrapperType, object? key)
+    {
+        Type messageType = message.GetType();
         Dictionary<Type, List<object?>> handlers = [];
 
-            foreach (object? handler in handlerProvider.Get(messageType))
+        void AddHandlers(IEnumerable<object?> newHandlers)
         {
-                if (handler is not null)
+            foreach (object? handler in newHandlers)
             {
-                    Type handlerType = handler.GetType();
+                if (handler == null) continue;
-                    if (!handlers.TryGetValue(handlerType, out List<object?>? handlerList))
-                    {
-                        handlerList = [];
-                        handlers.Add(handlerType, handlerList);
-                    }
 
-                    handlerList.Add(handler);
+                Type serviceType = handler.GetType();
-                }
-            }
-
-            foreach (object? handler in key is not null ? provider.GetKeyedServices(handlerWrapperType, key) :
-                provider.GetServices(handlerWrapperType))
-            {
-                if (handler?.GetType() is Type serviceType)
-                {
                 if (!handlers.TryGetValue(serviceType, out List<object?>? handlerList))
                 {
                     handlerList = [];
@@ -107,19 +102,14 @@ public class Mediator(IHandlerProvider handlerProvider,
             }
         }
 
-            foreach (KeyValuePair<Type, List<object?>> handlerEntry in handlers)
+        IEnumerable<object?> keyedServices = key != null ? provider.GetKeyedServices(handlerWrapperType, key) : 
-            {
+            provider.GetServices(handlerWrapperType);
-                foreach (object? handler in handlerEntry.Value)
+        AddHandlers(keyedServices);
-                {
+
-                    if (handler?.GetType().GetMethod("Handle", [messageType, typeof(CancellationToken)]) is MethodInfo handleMethod)
+        IEnumerable<object?> additionalHandlers = handlerProvider.Get(messageType, key);
-                    {
+        AddHandlers(additionalHandlers);
-                        return await(Task<object?>)handleMethod.Invoke(handler,
+
-                            new object[] { message, cancellationToken })!;
+        return handlers.SelectMany(entry => entry.Value).ToList();
-                    }
-                }
-            }
     }
 
-        return default;
-    }
 }