Harden RendererSynchronizationContext restoration

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

@@ -229,7 +229,7 @@ private void SendIfQuiescedOrElsePost<TState>(Action<TState> callback, TState st
 
     /// <summary>
     /// Sets the current synchronization context to this instance, invokes the <paramref name="callback"/>,
-    /// resets the synchronization context, and sets marks the builder as completed.
+    /// marks the builder as completed, and resets the synchronization context.
     /// </summary>
     private void InvokeWithThisAsCurrentSyncCtxThenSetResult<TState>(
         AsyncTaskMethodBuilder completion,
@@ -244,8 +244,10 @@ private void InvokeWithThisAsCurrentSyncCtxThenSetResult<TState>(
         }
         finally
         {
-            SetSynchronizationContext(original);
+            // Complete the queue marker while this context is still current so that queued
+            // continuations are not inlined onto the caller's thread.
             completion.SetResult();
+            SetSynchronizationContext(original);
         }
     }
 

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

@@ -694,7 +694,6 @@ public async Task InvokeAsync_FuncTaskT_CanRunAsynchronously_WhenBusy()
     {
         // Arrange
         var context = new RendererSynchronizationContext();
-        var thread = Thread.CurrentThread;
 
         var e1 = new ManualResetEventSlim();
         var e2 = new ManualResetEventSlim();
@@ -719,12 +718,12 @@ public async Task InvokeAsync_FuncTaskT_CanRunAsynchronously_WhenBusy()
         });
 
         // Assert
-        Assert.False(e2.IsSet);
+        Assert.False(e3.IsSet);
         e2.Set(); // Unblock the first item
         await task1;
 
-        Assert.True(e3.Wait(Timeout), "timeout");
-        Assert.NotSame(thread, await task2);
+        await task2;
+        Assert.True(e3.IsSet);
     }
 
     [Fact]
@@ -790,4 +789,41 @@ public async Task InvokeAsync_SyncWorkInAsyncTaskIsCompletedFirst()
         // Assert
         Assert.Equal(expected, actual);
     }
+
+    [Fact]
+    public void InvokeAsync_FuncTask_RestoresContextWhenQueuedWorkDoesNotFlowExecutionContext()
+    {
+        var context = new RendererSynchronizationContext();
+        var contextEntered = new ManualResetEventSlim();
+        var releaseContext = new ManualResetEventSlim();
+        var queuedWorkDone = new ManualResetEventSlim();
+        var callerDone = new ManualResetEventSlim();
+        SynchronizationContext original = null;
+        SynchronizationContext actual = null;
+
+        ThreadPool.UnsafeQueueUserWorkItem(_ =>
+        {
+            original = SynchronizationContext.Current;
+            _ = context.InvokeAsync(async () =>
+            {
+                contextEntered.Set();
+                Assert.True(releaseContext.Wait(Timeout), "timeout");
+                await Task.CompletedTask;
+            });
+
+            Assert.True(queuedWorkDone.Wait(Timeout), "timeout");
+            actual = SynchronizationContext.Current;
+            callerDone.Set();
+        }, null);
+
+        Assert.True(contextEntered.Wait(Timeout), "timeout");
+        using (ExecutionContext.SuppressFlow())
+        {
+            _ = context.InvokeAsync(queuedWorkDone.Set);
+        }
+
+        releaseContext.Set();
+        Assert.True(callerDone.Wait(Timeout), "timeout");
+        Assert.Same(original, actual);
+    }
 }

src/Components/test/E2ETest/ServerExecutionTests/ServerComponentRenderingTest.cs

@@ -32,4 +32,17 @@ public void ThrowsIfRenderIsRequestedOutsideSyncContext()
             $"{typeof(InvalidOperationException).FullName}: The current thread is not associated with the Dispatcher. Use InvokeAsync() to switch execution to the Dispatcher when triggering rendering or component state.",
             () => result.Text);
     }
+
+    [Fact]
+    public void RestoresContextWhenQueuedWorkDoesNotFlowExecutionContext()
+    {
+        var appElement = Browser.MountTestComponent<DispatchingComponent>();
+        var result = appElement.FindElement(By.Id("suppressed-execution-context-result"));
+
+        appElement.FindElement(By.Id("run-with-suppressed-execution-context")).Click();
+
+        Browser.Equal(
+            "Context leaked: False; Dispatcher overlapped: False",
+            () => result.Text);
+    }
 }

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

@@ -14,9 +14,17 @@
 <button id="run-with-dispatch" @onclick=RunWithDispatch>Run with dispatch</button>
 <button id="run-with-double-dispatch" @onclick=RunWithDoubleDispatch>Run with double dispatch</button>
 <button id="run-async-with-dispatch" @onclick=RunAsyncWorkWithDispatch>Run async work with dispatch</button>
+<button id="run-with-suppressed-execution-context" @onclick=RunWithSuppressedExecutionContext>
+    Run with suppressed execution context
+</button>
+
+<p id="suppressed-execution-context-result">@suppressedExecutionContextResult</p>
 
 @code {
+    static readonly TimeSpan Timeout = TimeSpan.FromSeconds(10);
+
     string result;
+    string suppressedExecutionContextResult;
 
     async Task RunWithoutDispatch()
     {
@@ -67,6 +75,95 @@
         result += " Fifth";
     }
 
+    void RunWithSuppressedExecutionContext()
+    {
+        suppressedExecutionContextResult = "Running";
+
+        ThreadPool.UnsafeQueueUserWorkItem(_ =>
+        {
+            try
+            {
+                InvokeAsync(() => { }).GetAwaiter().GetResult();
+                var (contextLeaked, dispatcherOverlapped) = RunSuppressedExecutionContextScenario();
+                suppressedExecutionContextResult =
+                    $"Context leaked: {contextLeaked}; Dispatcher overlapped: {dispatcherOverlapped}";
+            }
+            catch (Exception exception)
+            {
+                suppressedExecutionContextResult = exception.ToString();
+            }
+            finally
+            {
+                _ = InvokeAsync(StateHasChanged);
+            }
+        }, null);
+    }
+
+    (bool ContextLeaked, bool DispatcherOverlapped) RunSuppressedExecutionContextScenario()
+    {
+        using var rendererEntered = new ManualResetEventSlim();
+        using var releaseRenderer = new ManualResetEventSlim();
+        using var leakChecked = new ManualResetEventSlim();
+        using var dispatcherOccupied = new ManualResetEventSlim();
+        using var releaseDispatcher = new ManualResetEventSlim();
+        using var holderDone = new ManualResetEventSlim();
+        var contextLeaked = false;
+        var dispatcherOverlapped = false;
+
+        ThreadPool.UnsafeQueueUserWorkItem(_ =>
+        {
+            var originalContext = SynchronizationContext.Current;
+            InvokeAsync(async () =>
+            {
+                rendererEntered.Set();
+                Wait(releaseRenderer, nameof(releaseRenderer));
+                await Task.CompletedTask;
+            }).GetAwaiter().GetResult();
+
+            contextLeaked = !ReferenceEquals(SynchronizationContext.Current, originalContext);
+            leakChecked.Set();
+
+            Wait(dispatcherOccupied, nameof(dispatcherOccupied));
+            var componentUpdate = InvokeAsync(() =>
+            {
+                dispatcherOverlapped = dispatcherOccupied.IsSet && !releaseDispatcher.IsSet;
+                StateHasChanged();
+            });
+
+            releaseDispatcher.Set();
+            componentUpdate.GetAwaiter().GetResult();
+            holderDone.Set();
+        }, null);
+
+        Wait(rendererEntered, nameof(rendererEntered));
+        Task queuedNotification;
+        using (ExecutionContext.SuppressFlow())
+        {
+            queuedNotification = InvokeAsync(() => { });
+        }
+
+        releaseRenderer.Set();
+        Wait(leakChecked, nameof(leakChecked));
+
+        InvokeAsync(() =>
+        {
+            dispatcherOccupied.Set();
+            Wait(releaseDispatcher, nameof(releaseDispatcher));
+        }).GetAwaiter().GetResult();
+
+        Wait(holderDone, nameof(holderDone));
+        queuedNotification.GetAwaiter().GetResult();
+        return (contextLeaked, dispatcherOverlapped);
+    }
+
+    static void Wait(ManualResetEventSlim signal, string name)
+    {
+        if (!signal.Wait(Timeout))
+        {
+            throw new TimeoutException($"Timed out waiting for {name}.");
+        }
+    }
+
     void AttemptToRender()
     {
         try