Follow up 'Ability to monitor Blazor Server circuit activity'

src/Components/Server/src/Circuits/CircuitHandler.cs

@@ -72,4 +72,11 @@ public abstract class CircuitHandler
     /// <param name="cancellationToken">The <see cref="CancellationToken"/>.</param>
     /// <returns><see cref="Task"/> that represents the asynchronous execution operation.</returns>
     public virtual Task OnCircuitClosedAsync(Circuit circuit, CancellationToken cancellationToken) => Task.CompletedTask;
+
+    /// <summary>
+    /// Creates a handler that gets invoked when inbound activity on the circuit causes an asynchronous task to be dispatched on the server.
+    /// </summary>
+    /// <param name="next">The next handler to invoke.</param>
+    /// <returns>A handler function that returns a <see cref="Task"/> that completes when the activity has finished.</returns>
+    public virtual Func<CircuitInboundActivityContext, Task> CreateInboundActivityHandler(Func<CircuitInboundActivityContext, Task> next) => next;
 }

src/Components/Server/src/Circuits/CircuitHost.cs

@@ -606,24 +606,21 @@ internal async Task<TResult> HandleInboundActivityAsync<TResult>(Func<Task<TResu
 
     private static Func<Func<Task>, Task> BuildInboundActivityDispatcher(IReadOnlyList<CircuitHandler> circuitHandlers, Circuit circuit)
     {
-        Func<CircuitInboundActivityContext, Task>? result = null;
-
-        for (var i = circuitHandlers.Count - 1; i >= 0; i--)
+        if (circuitHandlers.Count == 0)
         {
-            if (circuitHandlers[i] is IHandleCircuitActivity inboundActivityHandler)
-            {
-                var next = result ?? (static (context) => context.Handler());
-                result = (context) => inboundActivityHandler.HandleInboundActivityAsync(context, next);
-            }
+            // If there are no registered handlers, there is no need to allocate a context on each call.
+            return static handler => handler();
         }
 
-        if (result is null)
+        var result = static (CircuitInboundActivityContext context) => context.Handler();
+
+        for (var i = circuitHandlers.Count - 1; i >= 0; i--)
         {
-            // If there are no registered handlers, there is no need to allocate a context on each call.
-            return static (handler) => handler();
+            var next = result;
+            result = circuitHandlers[i].CreateInboundActivityHandler(next);
         }
 
-        return (handler) => result(new(handler, circuit));
+        return handler => result(new(handler, circuit));
     }
 
     private void AssertInitialized()

src/Components/Server/src/PublicAPI.Unshipped.txt

@@ -1,5 +1,4 @@
 #nullable enable
 Microsoft.AspNetCore.Components.Server.Circuits.CircuitInboundActivityContext
 Microsoft.AspNetCore.Components.Server.Circuits.CircuitInboundActivityContext.Circuit.get -> Microsoft.AspNetCore.Components.Server.Circuits.Circuit!
-Microsoft.AspNetCore.Components.Server.Circuits.IHandleCircuitActivity
-Microsoft.AspNetCore.Components.Server.Circuits.IHandleCircuitActivity.HandleInboundActivityAsync(Microsoft.AspNetCore.Components.Server.Circuits.CircuitInboundActivityContext! context, System.Func<Microsoft.AspNetCore.Components.Server.Circuits.CircuitInboundActivityContext!, System.Threading.Tasks.Task!>! next) -> System.Threading.Tasks.Task!
+virtual Microsoft.AspNetCore.Components.Server.Circuits.CircuitHandler.CreateInboundActivityHandler(System.Func<Microsoft.AspNetCore.Components.Server.Circuits.CircuitInboundActivityContext!, System.Threading.Tasks.Task!>! next) -> System.Func<Microsoft.AspNetCore.Components.Server.Circuits.CircuitInboundActivityContext!, System.Threading.Tasks.Task!>!

src/Components/Server/test/Circuits/CircuitHostTest.cs

@@ -157,6 +157,8 @@ public async Task InitializeAsync_InvokesHandlers()
         var handler2 = new Mock<CircuitHandler>(MockBehavior.Strict);
         var sequence = new MockSequence();
 
+        SetupMockInboundActivityHandlers(sequence, handler1, handler2);
+
         handler1
             .InSequence(sequence)
             .Setup(h => h.OnCircuitOpenedAsync(It.IsAny<Circuit>(), cancellationToken))
@@ -242,6 +244,8 @@ public async Task InitializeAsync_ReportsOwnAsyncExceptions()
         var tcs = new TaskCompletionSource();
         var reportedErrors = new List<UnhandledExceptionEventArgs>();
 
+        SetupMockInboundActivityHandler(handler);
+
         handler
             .Setup(h => h.OnCircuitOpenedAsync(It.IsAny<Circuit>(), It.IsAny<CancellationToken>()))
             .Returns(tcs.Task)
@@ -284,6 +288,8 @@ public async Task DisposeAsync_InvokesCircuitHandler()
         var handler2 = new Mock<CircuitHandler>(MockBehavior.Strict);
         var sequence = new MockSequence();
 
+        SetupMockInboundActivityHandlers(sequence, handler1, handler2);
+
         handler1
             .InSequence(sequence)
             .Setup(h => h.OnConnectionDownAsync(It.IsAny<Circuit>(), cancellationToken))
@@ -327,29 +333,31 @@ public async Task HandleInboundActivityAsync_InvokesCircuitActivityHandlers()
         var handler3 = new Mock<CircuitHandler>(MockBehavior.Strict);
         var sequence = new MockSequence();
 
-        // We deliberately avoid making handler2 an inbound activity handler
-        var activityHandler1 = handler1.As<IHandleCircuitActivity>();
-        var activityHandler3 = handler3.As<IHandleCircuitActivity>();
-
         var asyncLocal1 = new AsyncLocal<bool>();
         var asyncLocal3 = new AsyncLocal<bool>();
 
-        activityHandler1
+        handler3
             .InSequence(sequence)
-            .Setup(h => h.HandleInboundActivityAsync(It.IsAny<CircuitInboundActivityContext>(), It.IsAny<Func<CircuitInboundActivityContext, Task>>()))
-            .Returns(async (CircuitInboundActivityContext context, Func<CircuitInboundActivityContext, Task> next) =>
+            .Setup(h => h.CreateInboundActivityHandler(It.IsAny<Func<CircuitInboundActivityContext, Task>>()))
+            .Returns((Func<CircuitInboundActivityContext, Task> next) => async (CircuitInboundActivityContext context) =>
             {
-                asyncLocal1.Value = true;
+                asyncLocal3.Value = true;
                 await next(context);
             })
             .Verifiable();
 
-        activityHandler3
+        handler2
+            .InSequence(sequence)
+            .Setup(h => h.CreateInboundActivityHandler(It.IsAny<Func<CircuitInboundActivityContext, Task>>()))
+            .Returns((Func<CircuitInboundActivityContext, Task> next) => next)
+            .Verifiable();
+
+        handler1
             .InSequence(sequence)
-            .Setup(h => h.HandleInboundActivityAsync(It.IsAny<CircuitInboundActivityContext>(), It.IsAny<Func<CircuitInboundActivityContext, Task>>()))
-            .Returns(async (CircuitInboundActivityContext context, Func<CircuitInboundActivityContext, Task> next) =>
+            .Setup(h => h.CreateInboundActivityHandler(It.IsAny<Func<CircuitInboundActivityContext, Task>>()))
+            .Returns((Func<CircuitInboundActivityContext, Task> next) => async (CircuitInboundActivityContext context) =>
             {
-                asyncLocal3.Value = true;
+                asyncLocal1.Value = true;
                 await next(context);
             })
             .Verifiable();
@@ -367,8 +375,9 @@ await circuitHost.HandleInboundActivityAsync(() =>
         });
 
         // Assert
-        activityHandler1.VerifyAll();
-        activityHandler3.VerifyAll();
+        handler1.VerifyAll();
+        handler2.VerifyAll();
+        handler3.VerifyAll();
 
         Assert.False(asyncLocal1.Value);
         Assert.False(asyncLocal3.Value);
@@ -404,6 +413,26 @@ private static TestRemoteRenderer GetRemoteRenderer()
             Mock.Of<IClientProxy>());
     }
 
+    private static void SetupMockInboundActivityHandlers(MockSequence sequence, params Mock<CircuitHandler>[] circuitHandlers)
+    {
+        for (var i = circuitHandlers.Length - 1; i >= 0; i--)
+        {
+            circuitHandlers[i]
+                .InSequence(sequence)
+                .Setup(h => h.CreateInboundActivityHandler(It.IsAny<Func<CircuitInboundActivityContext, Task>>()))
+                .Returns((Func<CircuitInboundActivityContext, Task> next) => next)
+                .Verifiable();
+        }
+    }
+
+    private static void SetupMockInboundActivityHandler(Mock<CircuitHandler> circuitHandler)
+    {
+        circuitHandler
+            .Setup(h => h.CreateInboundActivityHandler(It.IsAny<Func<CircuitInboundActivityContext, Task>>()))
+            .Returns((Func<CircuitInboundActivityContext, Task> next) => next)
+            .Verifiable();
+    }
+
     private class TestRemoteRenderer : RemoteRenderer
     {
         public TestRemoteRenderer(IServiceProvider serviceProvider, IClientProxy client)

src/Components/test/testassets/TestServer/TestCircuitContextAccessor.cs

@@ -5,16 +5,18 @@
 
 namespace TestServer;
 
-public class TestCircuitContextAccessor : CircuitHandler, IHandleCircuitActivity
+public class TestCircuitContextAccessor : CircuitHandler
 {
     private readonly AsyncLocal<bool> _hasCircuitContext = new();
 
     public bool HasCircuitContext => _hasCircuitContext.Value;
 
-    public async Task HandleInboundActivityAsync(CircuitInboundActivityContext context, Func<CircuitInboundActivityContext, Task> next)
-    {
-        _hasCircuitContext.Value = true;
-        await next(context);
-        _hasCircuitContext.Value = false;
-    }
+    public override Func<CircuitInboundActivityContext, Task> CreateInboundActivityHandler(
+        Func<CircuitInboundActivityContext, Task> next)
+        => async (context) =>
+        {
+            _hasCircuitContext.Value = true;
+            await next(context);
+            _hasCircuitContext.Value = false;
+        };
 }