Fix RendererSyncContext.Post() (#10451)

* Fix RendererSyncContext.Post()

Fixes: #9683 - SignalR connection breaks on large DOM

The root cause here is a misbehaving sync context. It's not legal for a
Post() implementation to run a callback synchronously. We want that
behavior for most of the functionality Blazor calls directly, but Post()
should always go async or else various threading primitives are broken.

* Fix incorrect tests

These tests have the assumption that setting the result of a TCS will
execution continuations synchronously. This was a bug in our
SyncContext, and these tests needed updating to be more resiliant.

* Remove a delegate allocation

src/Components/Components/src/Rendering/RendererSynchronizationContext.cs

@@ -39,12 +39,13 @@ namespace Microsoft.AspNetCore.Components.Rendering
 
         public Task Invoke(Action action)
         {
-            var completion = new TaskCompletionSource<object>();
+            var completion = new RendererSynchronizationTaskCompletionSource<Action, object>(action);
-            Post(_ =>
+            ExecuteSynchronouslyIfPossible((state) =>
             {
+                var completion = (RendererSynchronizationTaskCompletionSource<Action, object>)state;
                 try
                 {
-                    action();
+                    completion.Callback();
                     completion.SetResult(null);
                 }
                 catch (OperationCanceledException)
@@ -55,19 +56,20 @@ namespace Microsoft.AspNetCore.Components.Rendering
                 {
                     completion.SetException(exception);
                 }
-            }, null);
+            }, completion);
 
             return completion.Task;
         }
 
         public Task InvokeAsync(Func<Task> asyncAction)
         {
-            var completion = new TaskCompletionSource<object>();
+            var completion = new RendererSynchronizationTaskCompletionSource<Func<Task>, object>(asyncAction);
-            Post(async (_) =>
+            ExecuteSynchronouslyIfPossible(async (state) =>
             {
+                var completion = (RendererSynchronizationTaskCompletionSource<Func<Task>, object>)state;
                 try
                 {
-                    await asyncAction();
+                    await completion.Callback();
                     completion.SetResult(null);
                 }
                 catch (OperationCanceledException)
@@ -78,19 +80,20 @@ namespace Microsoft.AspNetCore.Components.Rendering
                 {
                     completion.SetException(exception);
                 }
-            }, null);
+            }, completion);
 
             return completion.Task;
         }
 
         public Task<TResult> Invoke<TResult>(Func<TResult> function)
         {
-            var completion = new TaskCompletionSource<TResult>();
+            var completion = new RendererSynchronizationTaskCompletionSource<Func<TResult>, TResult>(function);
-            Post(_ =>
+            ExecuteSynchronouslyIfPossible((state) =>
             {
+                var completion = (RendererSynchronizationTaskCompletionSource<Func<TResult>, TResult>)state;
                 try
                 {
-                    var result = function();
+                    var result = completion.Callback();
                     completion.SetResult(result);
                 }
                 catch (OperationCanceledException)
@@ -101,19 +104,20 @@ namespace Microsoft.AspNetCore.Components.Rendering
                 {
                     completion.SetException(exception);
                 }
-            }, null);
+            }, completion);
 
             return completion.Task;
         }
 
         public Task<TResult> InvokeAsync<TResult>(Func<Task<TResult>> asyncFunction)
         {
-            var completion = new TaskCompletionSource<TResult>();
+            var completion = new RendererSynchronizationTaskCompletionSource<Func<Task<TResult>>, TResult>(asyncFunction);
-            Post(async (_) =>
+            ExecuteSynchronouslyIfPossible(async (state) =>
             {
+                var completion = (RendererSynchronizationTaskCompletionSource<Func<Task<TResult>>, TResult>)state;
                 try
                 {
-                    var result = await asyncFunction();
+                    var result = await completion.Callback();
                     completion.SetResult(result);
                 }
                 catch (OperationCanceledException)
@@ -124,30 +128,20 @@ namespace Microsoft.AspNetCore.Components.Rendering
                 {
                     completion.SetException(exception);
                 }
-            }, null);
+            }, completion);
 
             return completion.Task;
         }
 
         // asynchronously runs the callback
+        //
+        // NOTE: this must always run async. It's not legal here to execute the work item synchronously.
         public override void Post(SendOrPostCallback d, object state)
         {
-            TaskCompletionSource<object> completion;
             lock (_state.Lock)
             {
-                if (!_state.Task.IsCompleted)
+                _state.Task = Enqueue(_state.Task, d, state, forceAsync: true);
-                {
-                    _state.Task = Enqueue(_state.Task, d, state);
-                    return;
-                }
-
-                // We can execute this synchronously because nothing is currently running
-                // or queued.
-                completion = new TaskCompletionSource<object>();
-                _state.Task = completion.Task;
             }
-
-            ExecuteSynchronously(completion, d, state);
         }
 
         // synchronously runs the callback
@@ -177,10 +171,33 @@ namespace Microsoft.AspNetCore.Components.Rendering
             return new RendererSynchronizationContext(_state);
         }
 
-        private Task Enqueue(Task antecedant, SendOrPostCallback d, object state)
+        // Similar to Post, but it can runs the work item synchronously if the context is not busy.
+        //
+        // This is the main code path used by components, we want to be able to run async work but only dispatch
+        // if necessary.
+        private void ExecuteSynchronouslyIfPossible(SendOrPostCallback d, object state)
         {
-            // If we get here is means that a callback is being queued while something is currently executing
+            TaskCompletionSource<object> completion;
-            // in this context. Let's instead add it to the queue and yield.
+            lock (_state.Lock)
+            {
+                if (!_state.Task.IsCompleted)
+                {
+                    _state.Task = Enqueue(_state.Task, d, state);
+                    return;
+                }
+
+                // We can execute this synchronously because nothing is currently running
+                // or queued.
+                completion = new TaskCompletionSource<object>();
+                _state.Task = completion.Task;
+            }
+
+            ExecuteSynchronously(completion, d, state);
+        }
+
+        private Task Enqueue(Task antecedant, SendOrPostCallback d, object state, bool forceAsync = false)
+        {
+            // If we get here is means that a callback is being explicitly queued. Let's instead add it to the queue and yield.
             //
             // We use our own queue here to maintain the execution order of the callbacks scheduled here. Also
             // we need a queue rather than just scheduling an item in the thread pool - those items would immediately
@@ -194,13 +211,14 @@ namespace Microsoft.AspNetCore.Components.Rendering
                 executionContext = ExecutionContext.Capture();
             }
 
+            var flags = forceAsync ? TaskContinuationOptions.RunContinuationsAsynchronously : TaskContinuationOptions.None;
             return antecedant.ContinueWith(BackgroundWorkThunk, new WorkItem()
             {
                 SynchronizationContext = this,
                 ExecutionContext = executionContext,
                 Callback = d,
                 State = state,
-            }, CancellationToken.None, TaskContinuationOptions.None, TaskScheduler.Current);
+            }, CancellationToken.None, flags, TaskScheduler.Current);
         }
 
         private void ExecuteSynchronously(
@@ -280,5 +298,15 @@ namespace Microsoft.AspNetCore.Components.Rendering
             public SendOrPostCallback Callback;
             public object State;
         }
+
+        private class RendererSynchronizationTaskCompletionSource<TCallback, TResult> : TaskCompletionSource<TResult>
+        {
+            public RendererSynchronizationTaskCompletionSource(TCallback callback)
+            {
+                Callback = callback;
+            }
+
+            public TCallback Callback { get; }
+        }
     }
 }

src/Components/Components/test/Auth/AuthorizeViewTest.cs

@@ -2,9 +2,11 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Diagnostics;
 using System.Linq;
 using System.Security.Claims;
 using System.Security.Principal;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Components.RenderTree;
 using Microsoft.AspNetCore.Components.Test.Helpers;
@@ -14,6 +16,10 @@ namespace Microsoft.AspNetCore.Components
 {
     public class AuthorizeViewTest
     {
+        // Nothing should exceed the timeout in a successful run of the the tests, this is just here to catch
+        // failures.
+        private static readonly TimeSpan Timeout = Debugger.IsAttached ? System.Threading.Timeout.InfiniteTimeSpan : TimeSpan.FromSeconds(10);
+
         [Fact]
         public void RendersNothingIfNotAuthorized()
         {
@@ -180,7 +186,11 @@ namespace Microsoft.AspNetCore.Components
         public void RendersNothingUntilAuthorizationCompleted()
         {
             // Arrange
-            var renderer = new TestRenderer();
+            var @event = new ManualResetEventSlim();
+            var renderer = new TestRenderer()
+            {
+                OnUpdateDisplayComplete = () => { @event.Set(); },
+            };
             var rootComponent = WrapInAuthorizeView(
                 notAuthorizedContent: builder => builder.AddContent(0, "You are not authorized"));
             var authTcs = new TaskCompletionSource<AuthenticationState>();
@@ -195,7 +205,12 @@ namespace Microsoft.AspNetCore.Components
             Assert.Empty(diff1.Edits);
 
             // Act/Assert 2: Auth process completes asynchronously
+            @event.Reset();
             authTcs.SetResult(new AuthenticationState(new ClaimsPrincipal()));
+
+            // We need to wait here because the continuations of SetResult will be scheduled to run asynchronously.
+            @event.Wait(Timeout);
+
             Assert.Equal(2, renderer.Batches.Count);
             var batch2 = renderer.Batches[1];
             var diff2 = batch2.DiffsByComponentId[authorizeViewComponentId].Single();
@@ -212,7 +227,11 @@ namespace Microsoft.AspNetCore.Components
         public void RendersAuthorizingContentUntilAuthorizationCompleted()
         {
             // Arrange
-            var renderer = new TestRenderer();
+            var @event = new ManualResetEventSlim();
+            var renderer = new TestRenderer()
+            {
+                OnUpdateDisplayComplete = () => { @event.Set(); },
+            };
             var rootComponent = WrapInAuthorizeView(
                 authorizingContent: builder => builder.AddContent(0, "Auth pending..."),
                 authorizedContent: context => builder => builder.AddContent(0, $"Hello, {context.User.Identity.Name}!"));
@@ -234,7 +253,12 @@ namespace Microsoft.AspNetCore.Components
             });
 
             // Act/Assert 2: Auth process completes asynchronously
+            @event.Reset();
             authTcs.SetResult(CreateAuthenticationState("Monsieur").Result);
+
+            // We need to wait here because the continuations of SetResult will be scheduled to run asynchronously.
+            @event.Wait(Timeout);
+
             Assert.Equal(2, renderer.Batches.Count);
             var batch2 = renderer.Batches[1];
             var diff2 = batch2.DiffsByComponentId[authorizeViewComponentId].Single();

src/Components/Components/test/ComponentBaseTest.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Components.RenderTree;
@@ -12,6 +13,10 @@ namespace Microsoft.AspNetCore.Components.Test
 {
     public class ComponentBaseTest
     {
+        // Nothing should exceed the timeout in a successful run of the the tests, this is just here to catch
+        // failures.
+        private static readonly TimeSpan Timeout = Debugger.IsAttached ? System.Threading.Timeout.InfiniteTimeSpan : TimeSpan.FromSeconds(10);
+
         [Fact]
         public void RunsOnInitWhenRendered()
         {
@@ -178,7 +183,12 @@ namespace Microsoft.AspNetCore.Components.Test
         public async Task RendersAfterParametersSetAndInitAsyncTasksAreCompleted()
         {
             // Arrange
-            var renderer = new TestRenderer();
+            var @event = new ManualResetEventSlim();
+
+            var renderer = new TestRenderer()
+            {
+                OnUpdateDisplayComplete = () => { @event.Set();  },
+            };
             var component = new TestComponent();
 
             component.Counter = 1;
@@ -197,10 +207,16 @@ namespace Microsoft.AspNetCore.Components.Test
             // A rendering should have happened after the synchronous execution of Init
             Assert.Single(renderer.Batches);
 
+            @event.Reset();
+
             // Completes task started by OnInitAsync
             component.Counter = 2;
             initTask.SetResult(true);
 
+            // We need to wait here, because the continuation from SetResult needs to be scheduled.
+            @event.Wait(Timeout);
+            @event.Reset();
+
             // Component should be rendered once, after set parameters
             Assert.Equal(2, renderer.Batches.Count);
 
@@ -209,6 +225,7 @@ namespace Microsoft.AspNetCore.Components.Test
             parametersSetTask.SetResult(false);
 
             await renderTask;
+            Assert.True(@event.IsSet);
 
             // Component should be rendered again
             // after the async part of onparameterssetasync completes

src/Components/Components/test/RendererTest.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Runtime.ExceptionServices;
 using System.Threading;
@@ -19,6 +20,10 @@ namespace Microsoft.AspNetCore.Components.Test
 {
     public class RendererTest
     {
+        // Nothing should exceed the timeout in a successful run of the the tests, this is just here to catch
+        // failures.
+        private static readonly TimeSpan Timeout = Debugger.IsAttached ? System.Threading.Timeout.InfiniteTimeSpan : TimeSpan.FromSeconds(10);
+
         private const string EventActionsName = nameof(NestedAsyncComponent.EventActions);
         private const string WhatToRenderName = nameof(NestedAsyncComponent.WhatToRender);
         private const string LogName = nameof(NestedAsyncComponent.Log);
@@ -2458,13 +2463,17 @@ namespace Microsoft.AspNetCore.Components.Test
         public void CallsAfterRenderAfterTheUIHasFinishedUpdatingAsynchronously()
         {
             // Arrange
+            var @event = new ManualResetEventSlim();
             var tcs = new TaskCompletionSource<object>();
             var afterRenderTcs = new TaskCompletionSource<object>();
             var onAfterRenderCallCountLog = new List<int>();
-            var component = new AsyncAfterRenderComponent(afterRenderTcs.Task);
+            var component = new AsyncAfterRenderComponent(afterRenderTcs.Task)
+            {
+                OnAfterRenderComplete = () => @event.Set(),
+            };
             var renderer = new AsyncUpdateTestRenderer()
             {
-                OnUpdateDisplayAsync = _ => tcs.Task
+                OnUpdateDisplayAsync = _ => tcs.Task,
             };
             renderer.AssignRootComponentId(component);
 
@@ -2473,6 +2482,9 @@ namespace Microsoft.AspNetCore.Components.Test
             tcs.SetResult(null);
             afterRenderTcs.SetResult(null);
 
+            // We need to wait here because the completions from SetResult will be scheduled.
+            @event.Wait(Timeout);
+
             // Assert
             Assert.True(component.Called);
         }
@@ -2481,9 +2493,13 @@ namespace Microsoft.AspNetCore.Components.Test
         public void CallsAfterRenderAfterTheUIHasFinishedUpdatingSynchronously()
         {
             // Arrange
+            var @event = new ManualResetEventSlim();
             var afterRenderTcs = new TaskCompletionSource<object>();
             var onAfterRenderCallCountLog = new List<int>();
-            var component = new AsyncAfterRenderComponent(afterRenderTcs.Task);
+            var component = new AsyncAfterRenderComponent(afterRenderTcs.Task)
+            {
+                OnAfterRenderComplete = () => @event.Set(),
+            };
             var renderer = new AsyncUpdateTestRenderer()
             {
                 OnUpdateDisplayAsync = _ => Task.CompletedTask
@@ -2494,6 +2510,9 @@ namespace Microsoft.AspNetCore.Components.Test
             component.TriggerRender();
             afterRenderTcs.SetResult(null);
 
+            // We need to wait here because the completions from SetResult will be scheduled.
+            @event.Wait(Timeout);
+
             // Assert
             Assert.True(component.Called);
         }
@@ -2798,7 +2817,12 @@ namespace Microsoft.AspNetCore.Components.Test
             // code paths are special cased for the first render because of prerendering.
 
             // Arrange
-            var renderer = new TestRenderer { ShouldHandleExceptions = true };
+            var @event = new ManualResetEventSlim();
+            var renderer = new TestRenderer()
+            {
+                ShouldHandleExceptions = true,
+                OnExceptionHandled = () => { @event.Set(); },
+            };
             var taskToAwait = Task.CompletedTask;
             var component = new TestComponent(builder =>
             {
@@ -2815,8 +2839,13 @@ namespace Microsoft.AspNetCore.Components.Test
 
             // Act
             var exception = new InvalidOperationException();
+
+            @event.Reset();
             asyncExceptionTcs.SetException(exception);
 
+            // We need to wait here because the continuations of SetException will be scheduled to run asynchronously.
+            @event.Wait(Timeout);
+
             // Assert
             Assert.Same(exception, Assert.Single(renderer.HandledExceptions).GetBaseException());
         }
@@ -3829,10 +3858,14 @@ namespace Microsoft.AspNetCore.Components.Test
 
             public bool Called { get; private set; }
 
+            public Action OnAfterRenderComplete { get; set; }
+
             public async Task OnAfterRenderAsync()
             {
                 await _task;
                 Called = true;
+
+                OnAfterRenderComplete?.Invoke();
             }
 
             protected override void BuildRenderTree(RenderTreeBuilder builder)

src/Components/Components/test/Rendering/RendererSynchronizationContextTest.cs

@@ -1,8 +1,9 @@
+// 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;
 using System.Diagnostics;
 using System.Globalization;
-using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
@@ -16,34 +17,51 @@ namespace Microsoft.AspNetCore.Components.Rendering
         public TimeSpan Timeout = Debugger.IsAttached ? System.Threading.Timeout.InfiniteTimeSpan : TimeSpan.FromSeconds(10);
 
         [Fact]
-        public void Post_CanRunSynchronously_WhenNotBusy()
+        public void Post_RunsAsynchronously_WhenNotBusy()
         {
             // Arrange
             var context = new RendererSynchronizationContext();
             var thread = Thread.CurrentThread;
             Thread capturedThread = null;
 
+            var e = new ManualResetEventSlim();
+            
             // Act
             context.Post((_) =>
             {
                 capturedThread = Thread.CurrentThread;
+
+                e.Set();
             }, null);
 
             // Assert
-            Assert.Same(thread, capturedThread);
+            Assert.True(e.Wait(Timeout), "timeout");
+            Assert.NotSame(thread, capturedThread);
         }
 
         [Fact]
-        public void Post_CanRunSynchronously_WhenNotBusy_Exception()
+        public void Post_RunsAynchronously_WhenNotBusy_Exception()
         {
             // Arrange
             var context = new RendererSynchronizationContext();
 
-            // Act & Assert
+            Exception exception = null;
-            Assert.Throws<InvalidTimeZoneException>(() => context.Post((_) =>
+            context.UnhandledException += (sender, e) =>
+            {
+                exception = (InvalidTimeZoneException)e.ExceptionObject;
+            };
+
+            // Act
+            context.Post((_) =>
             {
                 throw new InvalidTimeZoneException();
-            }, null));
+            }, null);
+
+            // Assert
+            //
+            // Use another item to 'push through' the throwing one
+            context.Send((_) => { }, null);
+            Assert.NotNull(exception);
         }
 
         [Fact]
@@ -573,7 +591,7 @@ namespace Microsoft.AspNetCore.Components.Rendering
             Thread capturedThread = null;
 
             // Act
-            var task = context.Invoke(() =>
+            var task = context.InvokeAsync(() =>
             {
                 capturedThread = Thread.CurrentThread;
                 return Task.CompletedTask;

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

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Text.Encodings.Web;
 using System.Threading;
 using System.Threading.Tasks;
@@ -19,6 +20,10 @@ namespace Microsoft.AspNetCore.Components.Browser.Rendering
 {
     public class RemoteRendererTest : HtmlRendererTestBase
     {
+        // Nothing should exceed the timeout in a successful run of the the tests, this is just here to catch
+        // failures.
+        private static readonly TimeSpan Timeout = Debugger.IsAttached ? System.Threading.Timeout.InfiniteTimeSpan : TimeSpan.FromSeconds(10);
+
         protected override HtmlRenderer GetHtmlRenderer(IServiceProvider serviceProvider)
         {
             return GetRemoteRenderer(serviceProvider, new CircuitClientProxy());
@@ -50,6 +55,7 @@ namespace Microsoft.AspNetCore.Components.Browser.Rendering
         public async Task ProcessBufferedRenderBatches_WritesRenders()
         {
             // Arrange
+            var @event = new ManualResetEventSlim();
             var serviceProvider = new ServiceCollection().BuildServiceProvider();
             var renderIds = new List<long>();
 
@@ -78,8 +84,13 @@ namespace Microsoft.AspNetCore.Components.Browser.Rendering
             var componentId = renderer.AssignRootComponentId(component);
             component.TriggerRender();
             renderer.OnRenderCompleted(2, null);
+
+            @event.Reset();
             firstBatchTCS.SetResult(null);
 
+            // Waiting is required here because the continuations of SetResult will not execute synchronously.
+            @event.Wait(Timeout);
+
             circuitClient.SetDisconnected();
             component.TriggerRender();
             component.TriggerRender();
@@ -261,7 +272,7 @@ namespace Microsoft.AspNetCore.Components.Browser.Rendering
                 NullLogger.Instance);
         }
 
-        private class TestComponent : IComponent
+        private class TestComponent : IComponent, IHandleAfterRender
         {
             private RenderHandle _renderHandle;
             private RenderFragment _renderFragment = (builder) =>
@@ -280,11 +291,19 @@ namespace Microsoft.AspNetCore.Components.Browser.Rendering
                 _renderFragment = renderFragment;
             }
 
+            public Action OnAfterRenderComplete { get; set; }
+
             public void Configure(RenderHandle renderHandle)
             {
                 _renderHandle = renderHandle;
             }
 
+            public Task OnAfterRenderAsync()
+            {
+                OnAfterRenderComplete?.Invoke();
+                return Task.CompletedTask;
+            }
+
             public Task SetParametersAsync(ParameterCollection parameters)
             {
                 TriggerRender();

src/Components/Shared/test/TestRenderer.cs

@@ -25,8 +25,12 @@ namespace Microsoft.AspNetCore.Components.Test.Helpers
         {
         }
 
+        public Action OnExceptionHandled { get; set; }
+
         public Action<RenderBatch> OnUpdateDisplay { get; set; }
 
+        public Action OnUpdateDisplayComplete { get; set; }
+
         public List<CapturedBatch> Batches { get; }
             = new List<CapturedBatch>();
 
@@ -81,6 +85,7 @@ namespace Microsoft.AspNetCore.Components.Test.Helpers
             }
 
             HandledExceptions.Add(exception);
+            OnExceptionHandled?.Invoke();
         }
 
         protected override Task UpdateDisplayAsync(in RenderBatch renderBatch)
@@ -102,6 +107,8 @@ namespace Microsoft.AspNetCore.Components.Test.Helpers
 
             // This renderer updates the UI synchronously, like the WebAssembly one.
             // To test async UI updates, subclass TestRenderer and override UpdateDisplayAsync.
+
+            OnUpdateDisplayComplete?.Invoke();
             return Task.CompletedTask;
         }
     }