Merge pull request #20 from fretje/parallel

Add "UseTaskWhenAll" option

Author: KuraiAndras

CHANGELOG.md

@@ -8,6 +8,7 @@
 - Added documentation comments to `ICourier`, `CourierInjector` and `CourierOptions`
 - Use Invoke iso DynamicInvoke for better performance
 - Fix bug where RemoveWeak would sometimes remove too many handlers
+- Add `UseTaskWhenAll` option for parallel notification handling
 
 # 6.0.0
 - Update MediatR to version 12

MediatR.Courier.Tests/CourierTests.cs

@@ -119,6 +119,68 @@ async Task NotificationAction(TestNotification _, CancellationToken cancellation
         Assert.Equal(1, receivedMessageCount);
     }
 
+    [Fact]
+    public async Task Handlers_RunInParallel_WhenUseTaskWhenAllIsTrue()
+    {
+        var options = new CourierOptions { UseTaskWhenAll = true };
+        var mediatRCourier = new MediatRCourier(options);
+
+        var delays = new[] { 100, 200, 300 };
+        var started = new bool[delays.Length];
+        var completed = new bool[delays.Length];
+
+        for (int i = 0; i < delays.Length; i++)
+        {
+            int idx = i;
+            mediatRCourier.Subscribe<TestNotification>(async (_, ct) =>
+            {
+                started[idx] = true;
+                await Task.Delay(delays[idx], ct);
+                completed[idx] = true;
+            });
+        }
+
+        var sw = System.Diagnostics.Stopwatch.StartNew();
+        await mediatRCourier.Handle(new TestNotification(), CancellationToken.None);
+        sw.Stop();
+
+        Assert.All(started, Assert.True);
+        Assert.All(completed, Assert.True);
+        // Should complete in just over the max delay (parallel)
+        Assert.InRange(sw.ElapsedMilliseconds, delays.Max(), delays.Max() + 150);
+    }
+
+    [Fact]
+    public async Task Handlers_RunSequentially_WhenUseTaskWhenAllIsFalse()
+    {
+        var options = new CourierOptions { CaptureThreadContext = false, UseTaskWhenAll = false };
+        var mediatRCourier = new MediatRCourier(options);
+
+        var delays = new[] { 100, 200, 300 };
+        var started = new bool[delays.Length];
+        var completed = new bool[delays.Length];
+
+        for (int i = 0; i < delays.Length; i++)
+        {
+            int idx = i;
+            mediatRCourier.Subscribe<TestNotification>(async (_, ct) =>
+            {
+                started[idx] = true;
+                await Task.Delay(delays[idx], ct);
+                completed[idx] = true;
+            });
+        }
+
+        var sw = System.Diagnostics.Stopwatch.StartNew();
+        await mediatRCourier.Handle(new TestNotification(), CancellationToken.None);
+        sw.Stop();
+
+        Assert.All(started, s => Assert.True(s));
+        Assert.All(completed, c => Assert.True(c));
+        // Should complete in just over the sum of delays (sequential)
+        Assert.InRange(sw.ElapsedMilliseconds, delays.Sum(), delays.Sum() + 150);
+    }
+
     private sealed class AsyncTestData : IEnumerable<object[]>
     {
         public IEnumerator<object[]> GetEnumerator()

MediatR.Courier/CourierOptions.cs

@@ -11,4 +11,12 @@ public sealed class CourierOptions
     /// Defaults to false.
     /// </summary>
     public bool CaptureThreadContext { get; set; }
+
+    /// <summary>
+    /// Gets or sets a value indicating whether to use Task.WhenAll to await handler tasks concurrently.
+    /// When set to true, all notification handlers are collected and awaited using Task.WhenAll.
+    /// When set to false, notification handlers are awaited sequentially.
+    /// Defaults to false.
+    /// </summary>
+    public bool UseTaskWhenAll { get; set; }
 }

MediatR.Courier/MediatRCourier.cs

@@ -22,6 +22,7 @@ async Task HandleLocal(INotification n, CancellationToken c)
             if (!_weakActions.TryGetValue(notificationType, out var weakSubscribers)) weakSubscribers = new();
 
             var remainingSubscribers = new ConcurrentBag<(WeakReference<object> target, MethodInfo methodInfo, bool needsToken)>();
+            var tasks = new List<Task>();
 
             foreach (var (target, methodInfo, needsToken) in weakSubscribers)
             {
@@ -32,7 +33,17 @@ async Task HandleLocal(INotification n, CancellationToken c)
                         : new object[] { n };
 
                     var result = methodInfo.Invoke(handler, parameters);
-                    if (result is Task task) await task.ConfigureAwait(_options.CaptureThreadContext);
+                    if (result is Task task)
+                    {
+                        if (_options.UseTaskWhenAll)
+                        {
+                            tasks.Add(task);
+                        }
+                        else
+                        {
+                            await task.ConfigureAwait(_options.CaptureThreadContext);
+                        }
+                    }
 
                     remainingSubscribers.Add((target, methodInfo, needsToken));
                 }
@@ -48,7 +59,22 @@ async Task HandleLocal(INotification n, CancellationToken c)
                     : new object[] { n };
 
                 var result = action.Method.Invoke(action.Target, parameters);
-                if (result is Task task) await task.ConfigureAwait(_options.CaptureThreadContext);
+                if (result is Task task)
+                {
+                    if (_options.UseTaskWhenAll)
+                    {
+                        tasks.Add(task);
+                    }
+                    else
+                    {
+                        await task.ConfigureAwait(_options.CaptureThreadContext);
+                    }
+                }
+            }
+
+            if (_options.UseTaskWhenAll && tasks.Count > 0)
+            {
+                await Task.WhenAll(tasks).ConfigureAwait(_options.CaptureThreadContext);
             }
         }