Fix event during onafterrender (#32017)

Co-authored-by: Pranav K <[email protected]>

src/Components/WebAssembly/WebAssembly/src/Hosting/WebAssemblyCallQueue.cs

@@ -0,0 +1,70 @@
+// 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;
+using System.Collections.Generic;
+
+namespace Microsoft.AspNetCore.Components.WebAssembly.Hosting
+{
+    // A mechanism for queuing JS-to-.NET calls so they aren't nested on the call stack and hence
+    // have the same ordering behaviors as in Blazor Server. This eliminates serveral inconsistency
+    // problems and bugs that otherwise require special-case solutions in other parts of the code.
+    //
+    // The reason for not using an actual SynchronizationContext for this is that, historically,
+    // Blazor WebAssembly has not enforced any rule around having to dispatch to a sync context.
+    // Adding such a rule now would be too breaking, given how component libraries may be reliant
+    // on being able to render at any time without InvokeAsync. If we add true multithreading in the
+    // future, we should start enforcing dispatch if (and only if) multithreading is enabled.
+    // For now, this minimal work queue is an internal detail of how the framework dispatches
+    // incoming JS->.NET calls and makes sure they get deferred if a renderbatch is in process.
+
+    internal static class WebAssemblyCallQueue
+    {
+        private static bool _isCallInProgress;
+        private static readonly Queue<Action> _pendingWork = new();
+
+        public static bool IsInProgress => _isCallInProgress;
+        public static bool HasUnstartedWork => _pendingWork.Count > 0;
+
+        /// <summary>
+        /// Runs the supplied callback when possible. If the call queue is empty, the callback is executed
+        /// synchronously. If some call is already executing within the queue, the callback is added to the
+        /// back of the queue and will be executed in turn.
+        /// </summary>
+        /// <typeparam name="T">The type of a state parameter for the callback</typeparam>
+        /// <param name="state">A state parameter for the callback. If the callback is able to execute synchronously, this allows us to avoid any allocations for the closure.</param>
+        /// <param name="callback">The callback to run.</param>
+        /// <remarks>
+        /// In most cases this should only be used for callbacks that will not throw, because
+        /// [1] Unhandled exceptions will be fatal to the application, as the work queue will no longer process
+        ///     further items (just like unhandled hub exceptions in Blazor Server)
+        /// [2] The exception will be thrown at the point of the top-level caller, which is not necessarily the
+        ///     code that scheduled the callback, so you may not be able to observe it.
+        ///
+        /// We could change this to return a Task and do the necessary try/catch things to direct exceptions back
+        /// to the code that scheduled the callback, but it's not required for current use cases and would require
+        /// at least an extra allocation and layer of try/catch per call, plus more work to schedule continuations
+        /// at the call site.
+        /// </remarks>
+        public static void Schedule<T>(T state, Action<T> callback)
+        {
+            if (_isCallInProgress)
+            {
+                _pendingWork.Enqueue(() => callback(state));
+            }
+            else
+            {
+                _isCallInProgress = true;
+                callback(state);
+
+                // Now run any queued work items
+                while (_pendingWork.TryDequeue(out var nextWorkItem))
+                {
+                    nextWorkItem();
+                }
+
+                _isCallInProgress = false;
+            }
+        }
+    }
+}

src/Components/WebAssembly/WebAssembly/src/Hosting/WebAssemblyHost.cs

@@ -158,13 +158,28 @@ namespace Microsoft.AspNetCore.Components.WebAssembly.Hosting
                 var loggerFactory = Services.GetRequiredService<ILoggerFactory>();
                 _renderer = new WebAssemblyRenderer(Services, loggerFactory);
 
-                var rootComponents = _rootComponents;
-                for (var i = 0; i < rootComponents.Length; i++)
+                var initializationTcs = new TaskCompletionSource();
+                WebAssemblyCallQueue.Schedule((_rootComponents, _renderer, initializationTcs), static async state =>
                 {
-                    var rootComponent = rootComponents[i];
-                    await _renderer.AddComponentAsync(rootComponent.ComponentType, rootComponent.Selector, rootComponent.Parameters);
-                }
-
+                    var (rootComponents, renderer, initializationTcs) = state;
+
+                    try
+                    {
+                        for (var i = 0; i < rootComponents.Length; i++)
+                        {
+                            var rootComponent = rootComponents[i];
+                            await renderer.AddComponentAsync(rootComponent.ComponentType, rootComponent.Selector, rootComponent.Parameters);
+                        }
+
+                        initializationTcs.SetResult();
+                    }
+                    catch (Exception ex)
+                    {
+                        initializationTcs.SetException(ex);
+                    }
+                });
+
+                await initializationTcs.Task;
                 store.ExistingState.Clear();
 
                 await tcs.Task;

src/Components/WebAssembly/WebAssembly/src/Rendering/WebAssemblyRenderer.cs

@@ -2,10 +2,10 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
-using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Components.RenderTree;
+using Microsoft.AspNetCore.Components.WebAssembly.Hosting;
 using Microsoft.AspNetCore.Components.WebAssembly.Services;
 using Microsoft.Extensions.Logging;
 using static Microsoft.AspNetCore.Internal.LinkerFlags;
@@ -20,9 +20,6 @@ namespace Microsoft.AspNetCore.Components.WebAssembly.Rendering
     {
         private readonly ILogger _logger;
         private readonly int _webAssemblyRendererId;
-        private readonly QueueWithLast<IncomingEventInfo> deferredIncomingEvents = new();
-
-        private bool isDispatchingEvent;
 
         /// <summary>
         /// Constructs an instance of <see cref="WebAssemblyRenderer"/>.
@@ -95,6 +92,32 @@ namespace Microsoft.AspNetCore.Components.WebAssembly.Rendering
             RendererRegistry.TryRemove(_webAssemblyRendererId);
         }
 
+        /// <inheritdoc />
+        protected override void ProcessPendingRender()
+        {
+            // For historical reasons, Blazor WebAssembly doesn't enforce that you use InvokeAsync
+            // to dispatch calls that originated from outside the system. Changing that now would be
+            // too breaking, at least until we can make it a prerequisite for multithreading.
+            // So, we don't have a way to guarantee that calls to here are already on our work queue.
+            //
+            // We do need rendering to happen on the work queue so that incoming events can be deferred
+            // until we've finished this rendering process (and other similar cases where we want
+            // execution order to be consistent with Blazor Server, which queues all JS->.NET calls).
+            //
+            // So, if we find that we're here and are not yet on the work queue, get onto it. Either
+            // way, rendering must continue synchronously here and is not deferred until later.
+            if (WebAssemblyCallQueue.IsInProgress)
+            {
+                base.ProcessPendingRender();
+            }
+            else
+            {
+                WebAssemblyCallQueue.Schedule(this, static @this => @this.CallBaseProcessPendingRender());
+            }
+        }
+
+        private void CallBaseProcessPendingRender() => base.ProcessPendingRender();
+
         /// <inheritdoc />
         protected override Task UpdateDisplayAsync(in RenderBatch batch)
         {
@@ -103,22 +126,21 @@ namespace Microsoft.AspNetCore.Components.WebAssembly.Rendering
                 _webAssemblyRendererId,
                 batch);
 
-            if (deferredIncomingEvents.Count == 0)
+            if (WebAssemblyCallQueue.HasUnstartedWork)
             {
-                // In the vast majority of cases, since the call to update the UI is synchronous,
-                // we just return a pre-completed task from here.
-                return Task.CompletedTask;
+                // Because further incoming calls from JS to .NET are already queued (e.g., event notifications),
+                // we have to delay the renderbatch acknowledgement until it gets to the front of that queue.
+                // This is for consistency with Blazor Server which queues all JS-to-.NET calls relative to each
+                // other, and because various bits of cleanup logic rely on this ordering.
+                var tcs = new TaskCompletionSource();
+                WebAssemblyCallQueue.Schedule(tcs, static tcs => tcs.SetResult());
+                return tcs.Task;
             }
             else
             {
-                // However, in the rare case where JS sent us any event notifications that we had to
-                // defer until later, we behave as if the renderbatch isn't acknowledged until we have at
-                // least dispatched those event calls. This is to make the WebAssembly behavior more
-                // consistent with the Server behavior, which receives batch acknowledgements asynchronously
-                // and they are queued up with any other calls from JS such as event calls. If we didn't
-                // do this, then the order of execution could be inconsistent with Server, and in fact
-                // leads to a specific bug: https://github.com/dotnet/aspnetcore/issues/26838
-                return deferredIncomingEvents.Last.StartHandlerCompletionSource.Task;
+                // Nothing else is pending, so we can treat the renderbatch as acknowledged synchronously.
+                // This lets upstream code skip an expensive code path and avoids some allocations.
+                return Task.CompletedTask;
             }
         }
 
@@ -138,90 +160,6 @@ namespace Microsoft.AspNetCore.Components.WebAssembly.Rendering
             }
         }
 
-        /// <inheritdoc />
-        public override Task DispatchEventAsync(ulong eventHandlerId, EventFieldInfo? eventFieldInfo, EventArgs eventArgs)
-        {
-            // Be sure we only run one event handler at once. Although they couldn't run
-            // simultaneously anyway (there's only one thread), they could run nested on
-            // the stack if somehow one event handler triggers another event synchronously.
-            // We need event handlers not to overlap because (a) that's consistent with
-            // server-side Blazor which uses a sync context, and (b) the rendering logic
-            // relies completely on the idea that within a given scope it's only building
-            // or processing one batch at a time.
-            //
-            // The only currently known case where this makes a difference is in the E2E
-            // tests in ReorderingFocusComponent, where we hit what seems like a Chrome bug
-            // where mutating the DOM cause an element's "change" to fire while its "input"
-            // handler is still running (i.e., nested on the stack) -- this doesn't happen
-            // in Firefox. Possibly a future version of Chrome may fix this, but even then,
-            // it's conceivable that DOM mutation events could trigger this too.
-
-            if (isDispatchingEvent)
-            {
-                var info = new IncomingEventInfo(eventHandlerId, eventFieldInfo, eventArgs);
-                deferredIncomingEvents.Enqueue(info);
-                return info.FinishHandlerCompletionSource.Task;
-            }
-            else
-            {
-                try
-                {
-                    isDispatchingEvent = true;
-                    return base.DispatchEventAsync(eventHandlerId, eventFieldInfo, eventArgs);
-                }
-                finally
-                {
-                    isDispatchingEvent = false;
-
-                    if (deferredIncomingEvents.Count > 0)
-                    {
-                        // Fire-and-forget because the task we return from this method should only reflect the
-                        // completion of its own event dispatch, not that of any others that happen to be queued.
-                        // Also, ProcessNextDeferredEventAsync deals with its own async errors.
-                        _ = ProcessNextDeferredEventAsync();
-                    }
-                }
-            }
-        }
-
-        private async Task ProcessNextDeferredEventAsync()
-        {
-            var info = deferredIncomingEvents.Dequeue();
-
-            try
-            {
-                var handlerTask = DispatchEventAsync(info.EventHandlerId, info.EventFieldInfo, info.EventArgs);
-                info.StartHandlerCompletionSource.SetResult();
-                await handlerTask;
-                info.FinishHandlerCompletionSource.SetResult();
-            }
-            catch (Exception ex)
-            {
-                // Even if the handler threw synchronously, we at least started processing, so always complete successfully
-                info.StartHandlerCompletionSource.TrySetResult();
-
-                info.FinishHandlerCompletionSource.SetException(ex);
-            }
-        }
-
-        readonly struct IncomingEventInfo
-        {
-            public readonly ulong EventHandlerId;
-            public readonly EventFieldInfo? EventFieldInfo;
-            public readonly EventArgs EventArgs;
-            public readonly TaskCompletionSource StartHandlerCompletionSource;
-            public readonly TaskCompletionSource FinishHandlerCompletionSource;
-
-            public IncomingEventInfo(ulong eventHandlerId, EventFieldInfo? eventFieldInfo, EventArgs eventArgs)
-            {
-                EventHandlerId = eventHandlerId;
-                EventFieldInfo = eventFieldInfo;
-                EventArgs = eventArgs;
-                StartHandlerCompletionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-                FinishHandlerCompletionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
-            }
-        }
-
         private static class Log
         {
             private static readonly Action<ILogger, string, Exception> _unhandledExceptionRenderingComponent;
@@ -247,30 +185,5 @@ namespace Microsoft.AspNetCore.Components.WebAssembly.Rendering
                     exception);
             }
         }
-
-        private class QueueWithLast<T>
-        {
-            private readonly Queue<T> _items = new();
-
-            public int Count => _items.Count;
-
-            public T? Last { get; private set; }
-
-            public T Dequeue()
-            {
-                if (_items.Count == 1)
-                {
-                    Last = default;
-                }
-
-                return _items.Dequeue();
-            }
-
-            public void Enqueue(T item)
-            {
-                Last = item;
-                _items.Enqueue(item);
-            }
-        }
     }
 }

src/Components/WebAssembly/WebAssembly/src/Services/DefaultWebAssemblyJSRuntime.cs

@@ -4,6 +4,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
 using System.Text.Json;
+using Microsoft.AspNetCore.Components.WebAssembly.Hosting;
 using Microsoft.JSInterop.Infrastructure;
 using Microsoft.JSInterop.WebAssembly;
 
@@ -35,7 +36,14 @@ namespace Microsoft.AspNetCore.Components.WebAssembly.Services
 
         // Invoked via Mono's JS interop mechanism (invoke_method)
         public static void EndInvokeJS(string argsJson)
-            => DotNetDispatcher.EndInvokeJS(Instance, argsJson);
+        {
+            WebAssemblyCallQueue.Schedule(argsJson, static argsJson =>
+            {
+                // This is not expected to throw, as it takes care of converting any unhandled user code
+                // exceptions into a failure on the Task that was returned when calling InvokeAsync.
+                DotNetDispatcher.EndInvokeJS(Instance, argsJson);
+            });
+        }
 
         // Invoked via Mono's JS interop mechanism (invoke_method)
         public static void BeginInvokeDotNet(string callId, string assemblyNameOrDotNetObjectId, string methodIdentifier, string argsJson)
@@ -57,7 +65,12 @@ namespace Microsoft.AspNetCore.Components.WebAssembly.Services
             }
 
             var callInfo = new DotNetInvocationInfo(assemblyName, methodIdentifier, dotNetObjectId, callId);
-            DotNetDispatcher.BeginInvokeDotNet(Instance, callInfo, argsJson);
+            WebAssemblyCallQueue.Schedule((callInfo, argsJson), static state =>
+            {
+                // This is not expected to throw, as it takes care of converting any unhandled user code
+                // exceptions into a failure on the JS Promise object.
+                DotNetDispatcher.BeginInvokeDotNet(Instance, state.callInfo, state.argsJson);
+            });
         }
     }
 }

src/Components/test/E2ETest/Tests/ComponentRenderingTest.cs

@@ -444,6 +444,28 @@ namespace Microsoft.AspNetCore.Components.E2ETest.Tests
             long getPageYOffset() => (long)((IJavaScriptExecutor)Browser).ExecuteScript("return window.pageYOffset");
         }
 
+        [Theory]
+        [InlineData("focus-button-onafterrender-invoke")]
+        [InlineData("focus-button-onafterrender-await")]
+        public void CanFocusDuringOnAfterRenderAsyncWithFocusInEvent(string triggerButton)
+        {
+            // Represents https://github.com/dotnet/aspnetcore/issues/30070, plus a more complicated
+            // variant where the initial rendering doesn't start from a JS interop call and hence
+            // isn't automatically part of the WebAssemblyCallQueue.
+
+            var appElement = Browser.MountTestComponent<ElementFocusComponent>();
+            var didReceiveFocusLabel = appElement.FindElement(By.Id("focus-event-received"));
+            Browser.Equal("False", () => didReceiveFocusLabel.Text);
+
+            appElement.FindElement(By.Id(triggerButton)).Click();
+            Browser.Equal("True", () => didReceiveFocusLabel.Text);
+            Browser.Equal("focus-input-onafterrender", () => Browser.SwitchTo().ActiveElement().GetAttribute("id"));
+
+            // As well as actually focusing and triggering the onfocusin event, we should not be seeing any errors
+            var log = Browser.Manage().Logs.GetLog(LogType.Browser);
+            Assert.DoesNotContain(log, entry => entry.Level == LogLevel.Severe);
+        }
+
         [Fact]
         public void CanCaptureReferencesToDynamicallyAddedElements()
         {

src/Components/test/testassets/BasicTestApp/ElementFocusComponent.razor

@@ -1,16 +1,52 @@
-@using Microsoft.JSInterop
+@if (performFocusDuringOnAfterRender)
+{
+    <input id="focus-input-onafterrender" @ref="inputForOnAfterRenderReference" @onfocusin="@(() => { didReceiveFocusIn = true; })" />
+}
 
 <button id="focus-button" @onclick="@(() => FocusInput(false))">Click to focus!</button>
 <hr />
 <button id="focus-button-prevented" @onclick="@(() => FocusInput(true))">Click to focus with preventScroll!</button>
 <hr />
+<button id="focus-button-onafterrender-invoke" @onclick="@FocusDuringOnAfterRenderViaInvoke">Focus during OnAfterRender via Invoke</button>
+<hr />
+<button id="focus-button-onafterrender-await" @onclick="@FocusDuringOnAfterRenderViaAwait">Focus during OnAfterRender via await</button>
+<hr />
+<p>Received focus in: <span id="focus-event-received">@didReceiveFocusIn</span></p>
+<hr />
 <input id="focus-input" @ref="inputReference" style="margin-top: 10000px" />
 
 @code {
     private ElementReference inputReference;
+    private ElementReference inputForOnAfterRenderReference;
+    private bool performFocusDuringOnAfterRender;
+    private bool didReceiveFocusIn;
 
     private async Task FocusInput(bool preventScroll)
     {
         await inputReference.FocusAsync(preventScroll);
     }
+
+    private void FocusDuringOnAfterRenderViaInvoke()
+    {
+        Task.Run(async () =>
+        {
+            await Task.Yield();
+            performFocusDuringOnAfterRender = true;
+            _ = InvokeAsync(StateHasChanged);
+        });
+    }
+
+    private async Task FocusDuringOnAfterRenderViaAwait()
+    {
+        await Task.Yield();
+        performFocusDuringOnAfterRender = true;
+    }
+
+    protected override async Task OnAfterRenderAsync(bool firstRender)
+    {
+        if (performFocusDuringOnAfterRender)
+        {
+            await inputForOnAfterRenderReference.FocusAsync();
+        }
+    }
 }