Coalesce JoinableTaskFactory dispatcher posts

docfx/docs/threading_rules.md

@@ -299,9 +299,50 @@ The following describes how to replace the mechanism for getting to the
 UI thread in a host-independent way:
 
 You can set your own priority by creating your own derived type of
-`JoinableTaskFactory` and overriding the `PostToUnderlyingSynchronizationContext`
-method. This method is responsible both for initial switches to the UI
-thread as well as resuming on the UI thread after a yielding await.
+`JoinableTaskFactory`.
+
+The base implementation coalesces pending callbacks so that only one driver
+message at a time is queued to the underlying synchronization context. A
+derived type that does not override `PostToUnderlyingSynchronizationContext`
+inherits this behavior.
+
+For backward compatibility, an override of
+`PostToUnderlyingSynchronizationContext` remains fully authoritative and does
+not automatically coalesce. Existing derived types may suppress a post,
+redirect it, or apply semantics that the base class cannot safely assume. Such
+types therefore retain their original behavior.
+
+A derived type may explicitly opt into coalescing by routing
+`PostToUnderlyingSynchronizationContext` through
+`PostToUnderlyingSynchronizationContextWithCoalescing`, and overriding
+`PostToUnderlyingSynchronizationContextCore` with the actual dispatcher
+operation. Because `JoinableTaskFactory` is defined in another assembly, C#
+requires these `protected internal` base members to be declared `protected`
+when overridden:
+
+```csharp
+protected override void PostToUnderlyingSynchronizationContext(
+    SendOrPostCallback callback,
+    object state)
+{
+    this.PostToUnderlyingSynchronizationContextWithCoalescing(callback, state);
+}
+
+protected override void PostToUnderlyingSynchronizationContextCore(
+    SendOrPostCallback callback,
+    object state)
+{
+    this.UnderlyingSynchronizationContext!.Post(callback, state);
+}
+```
+
+`PostToUnderlyingSynchronizationContext` is responsible both for initial
+switches to the UI thread and for resuming on the UI thread after a yielding
+await. When coalescing is enabled, the core method may be called once for a
+sequence of pending callbacks and should only perform the underlying post; it
+should not call the coalescing helper. Replace the synchronization-context post
+shown above with the custom dispatcher or priority operation required by the
+derived factory.
 
 Note that the `JoinableTaskFactory` class has no default constructor, so when
 implementing your own `JoinableTaskFactory`-derived type you will need to add

src/Microsoft.VisualStudio.Threading/DispatcherExtensions.cs

@@ -76,6 +76,12 @@ internal DispatcherJoinableTaskFactory(JoinableTaskFactory innerFactory, Dispatc
 
         /// <inheritdoc />
         protected internal override void PostToUnderlyingSynchronizationContext(SendOrPostCallback callback, object state)
+        {
+            this.PostToUnderlyingSynchronizationContextWithCoalescing(callback, state);
+        }
+
+        /// <inheritdoc />
+        protected internal override void PostToUnderlyingSynchronizationContextCore(SendOrPostCallback callback, object state)
         {
             this.dispatcher.BeginInvoke(this.priority, callback, state);
         }

src/Microsoft.VisualStudio.Threading/JoinableTaskFactory.cs

@@ -23,18 +23,35 @@ namespace Microsoft.VisualStudio.Threading;
 /// </remarks>
 public partial class JoinableTaskFactory
 {
+    private static readonly ContextCallback ExecutePendingUnderlyingSynchronizationContextCallbackDelegate = state =>
+    {
+        var callback = ((SendOrPostCallback Callback, object State))state!;
+        callback.Callback(callback.State);
+    };
+
+    private static readonly SendOrPostCallback ExecuteOnePendingUnderlyingSynchronizationContextCallbackDelegate = state => ((UnderlyingSynchronizationContextCallback)state!).Execute();
+
+    [ThreadStatic]
+    private static List<JoinableTaskFactory>? synchronouslyPostingFactories;
+
     /// <summary>
     /// The <see cref="JoinableTaskContext"/> that owns this instance.
     /// </summary>
     private readonly JoinableTaskContext owner;
 
+    private readonly object pendingUnderlyingSynchronizationContextCallbacksLock = new();
+
     private readonly SynchronizationContext? mainThreadJobSyncContext;
 
     /// <summary>
     /// The collection to add all created tasks to. May be <see langword="null" />.
     /// </summary>
     private readonly JoinableTaskCollection? jobCollection;
 
+    private Queue<(SendOrPostCallback Callback, object State, ExecutionContext? ExecutionContext)>? pendingUnderlyingSynchronizationContextCallbacks;
+
+    private bool underlyingSynchronizationContextCallbackPending;
+
     /// <summary>
     /// Backing field for the <see cref="HangDetectionTimeout"/> property.
     /// </summary>
@@ -465,11 +482,35 @@ internal void Post(SendOrPostCallback callback, object? state, bool mainThreadAf
     /// </summary>
     /// <param name="callback">The callback to invoke.</param>
     /// <param name="state">State to pass to the callback.</param>
+    /// <remarks>
+    /// The base implementation coalesces pending callbacks. An override replaces that behavior entirely,
+    /// preserving the semantics of derived types written before coalescing was introduced. A derived type
+    /// may opt into coalescing by calling <see cref="PostToUnderlyingSynchronizationContextWithCoalescing"/>
+    /// from this override and overriding <see cref="PostToUnderlyingSynchronizationContextCore"/> to perform
+    /// the actual post.
+    /// </remarks>
     protected internal virtual void PostToUnderlyingSynchronizationContext(SendOrPostCallback callback, object state)
     {
         Requires.NotNull(callback, nameof(callback));
         Assumes.NotNull(this.UnderlyingSynchronizationContext);
 
+        this.PostToUnderlyingSynchronizationContextWithCoalescing(callback, state);
+    }
+
+    /// <summary>
+    /// Posts a message directly to the underlying synchronization context.
+    /// </summary>
+    /// <param name="callback">The callback to invoke.</param>
+    /// <param name="state">State to pass to the callback.</param>
+    /// <remarks>
+    /// Derived types that opt into coalescing should override this method, rather than
+    /// <see cref="PostToUnderlyingSynchronizationContext"/>, with their custom dispatcher operation.
+    /// </remarks>
+    protected internal virtual void PostToUnderlyingSynchronizationContextCore(SendOrPostCallback callback, object state)
+    {
+        Requires.NotNull(callback, nameof(callback));
+        Assumes.NotNull(this.UnderlyingSynchronizationContext);
+
         this.UnderlyingSynchronizationContext.Post(callback, state);
     }
 
@@ -634,6 +675,55 @@ protected void Add(JoinableTask joinable)
         }
     }
 
+    /// <summary>
+    /// Posts a message to the underlying synchronization context while coalescing pending messages.
+    /// </summary>
+    /// <param name="callback">The callback to invoke.</param>
+    /// <param name="state">
+    /// State to pass to the callback. Implementing <see cref="IPendingExecutionRequestState"/> allows
+    /// the callback to be removed from the private queue when it has already executed by another means.
+    /// </param>
+    /// <remarks>
+    /// This method is intended for derived types that override <see cref="PostToUnderlyingSynchronizationContext"/>
+    /// and explicitly opt into coalescing. Such types should also override
+    /// <see cref="PostToUnderlyingSynchronizationContextCore"/> to perform the actual dispatcher post.
+    /// </remarks>
+    protected void PostToUnderlyingSynchronizationContextWithCoalescing(SendOrPostCallback callback, object state)
+    {
+        Requires.NotNull(callback, nameof(callback));
+
+        ExecutionContext? executionContext = ExecutionContext.Capture();
+        bool postCallback = false;
+        lock (this.pendingUnderlyingSynchronizationContextCallbacksLock)
+        {
+            this.pendingUnderlyingSynchronizationContextCallbacks ??= new Queue<(SendOrPostCallback, object, ExecutionContext?)>();
+            this.RemoveCompletedPendingUnderlyingSynchronizationContextCallbacks();
+            this.pendingUnderlyingSynchronizationContextCallbacks.Enqueue((callback, state, executionContext));
+            if (!this.underlyingSynchronizationContextCallbackPending)
+            {
+                this.underlyingSynchronizationContextCallbackPending = true;
+                postCallback = true;
+            }
+        }
+
+        if (postCallback)
+        {
+            this.PostPendingUnderlyingSynchronizationContextCallback(propagateException: true);
+        }
+    }
+
+    /// <summary>
+    /// Checks whether this thread is currently inside an underlying synchronous post for the specified factory.
+    /// </summary>
+    /// <remarks>
+    /// The full chain is tracked so nested posts across factories (for example A to B to A) recognize an
+    /// ancestor factory and drain it iteratively instead of recursively posting another driver message.
+    /// </remarks>
+    private static bool IsSynchronouslyPosting(JoinableTaskFactory factory)
+    {
+        return synchronouslyPostingFactories?.Contains(factory) is true;
+    }
+
     /// <summary>
     /// Throws an exception if an active AsyncReaderWriterLock
     /// upgradeable read or write lock is held by the caller.
@@ -659,6 +749,150 @@ private static void VerifyNoNonConcurrentSyncContext()
         }
     }
 
+    /// <summary>
+    /// Executes one callback from the private queue and, when more work is already queued,
+    /// posts its successor to the underlying synchronization context before invoking the callback.
+    /// </summary>
+    /// <remarks>
+    /// Posting the successor first ensures that code invoked by the callback can enter a nested
+    /// message loop and find the next message already available. Synchronization contexts that
+    /// execute <see cref="SynchronizationContext.Post(SendOrPostCallback, object?)"/> inline are
+    /// drained iteratively instead to avoid recursive stack growth.
+    /// </remarks>
+    private void ExecuteOnePendingUnderlyingSynchronizationContextCallback()
+    {
+        bool continueSynchronously;
+        do
+        {
+            continueSynchronously = false;
+            (SendOrPostCallback Callback, object State, ExecutionContext? ExecutionContext)? callback = null;
+            bool postSuccessor = false;
+            bool completeSynchronousDrainAfterCallback = false;
+            try
+            {
+                lock (this.pendingUnderlyingSynchronizationContextCallbacksLock)
+                {
+                    this.RemoveCompletedPendingUnderlyingSynchronizationContextCallbacks();
+                    if (this.pendingUnderlyingSynchronizationContextCallbacks?.Count > 0)
+                    {
+                        callback = this.pendingUnderlyingSynchronizationContextCallbacks.Dequeue();
+                    }
+
+                    this.RemoveCompletedPendingUnderlyingSynchronizationContextCallbacks();
+                    if (IsSynchronouslyPosting(this))
+                    {
+                        completeSynchronousDrainAfterCallback = true;
+                        continueSynchronously = this.pendingUnderlyingSynchronizationContextCallbacks?.Count > 0;
+                    }
+                    else if (this.pendingUnderlyingSynchronizationContextCallbacks?.Count > 0)
+                    {
+                        postSuccessor = true;
+                    }
+                    else
+                    {
+                        this.pendingUnderlyingSynchronizationContextCallbacks = null;
+                        this.underlyingSynchronizationContextCallbackPending = false;
+                    }
+                }
+
+                if (postSuccessor)
+                {
+                    this.PostPendingUnderlyingSynchronizationContextCallback(propagateException: false);
+                }
+
+                if (callback is { } work)
+                {
+                    if (work.ExecutionContext is object)
+                    {
+                        ExecutionContext.Run(work.ExecutionContext, ExecutePendingUnderlyingSynchronizationContextCallbackDelegate, (work.Callback, work.State));
+                    }
+                    else
+                    {
+                        work.Callback(work.State);
+                    }
+                }
+            }
+            finally
+            {
+                if (completeSynchronousDrainAfterCallback)
+                {
+                    lock (this.pendingUnderlyingSynchronizationContextCallbacksLock)
+                    {
+                        this.RemoveCompletedPendingUnderlyingSynchronizationContextCallbacks();
+                        if (this.pendingUnderlyingSynchronizationContextCallbacks?.Count > 0)
+                        {
+                            continueSynchronously = true;
+                        }
+                        else
+                        {
+                            this.pendingUnderlyingSynchronizationContextCallbacks = null;
+                            this.underlyingSynchronizationContextCallbackPending = false;
+                        }
+                    }
+                }
+            }
+        }
+        while (continueSynchronously);
+    }
+
+    private void PostPendingUnderlyingSynchronizationContextCallback(bool propagateException)
+    {
+        try
+        {
+            List<JoinableTaskFactory> synchronousPostingChain = synchronouslyPostingFactories ??= new();
+            synchronousPostingChain.Add(this);
+            bool restoreFlow = !ExecutionContext.IsFlowSuppressed();
+            if (restoreFlow)
+            {
+                ExecutionContext.SuppressFlow();
+            }
+
+            try
+            {
+                this.PostToUnderlyingSynchronizationContextCore(ExecuteOnePendingUnderlyingSynchronizationContextCallbackDelegate, new UnderlyingSynchronizationContextCallback(this));
+            }
+            finally
+            {
+                if (restoreFlow)
+                {
+                    ExecutionContext.RestoreFlow();
+                }
+
+                synchronousPostingChain.RemoveAt(synchronousPostingChain.Count - 1);
+            }
+        }
+        catch
+        {
+            lock (this.pendingUnderlyingSynchronizationContextCallbacksLock)
+            {
+                // Every callback remains in its owning JoinableTask's execution queue, so abandon this
+                // secondary route when no underlying message was established.
+                this.pendingUnderlyingSynchronizationContextCallbacks = null;
+                this.underlyingSynchronizationContextCallbackPending = false;
+            }
+
+            if (propagateException)
+            {
+                throw;
+            }
+        }
+    }
+
+    private void RemoveCompletedPendingUnderlyingSynchronizationContextCallbacks()
+    {
+        Assumes.True(Monitor.IsEntered(this.pendingUnderlyingSynchronizationContextCallbacksLock));
+
+        // Only inspect the head to keep enqueue and drain operations O(1). Completed entries behind live work
+        // are removed as they reach the head, while coalescing still limits the underlying context to one driver.
+#pragma warning disable VSOnly // IPendingExecutionRequestState is intended for evaluation purposes only.
+        while (this.pendingUnderlyingSynchronizationContextCallbacks?.Count > 0
+            && this.pendingUnderlyingSynchronizationContextCallbacks.Peek().State is IPendingExecutionRequestState { IsCompleted: true })
+#pragma warning restore VSOnly
+        {
+            this.pendingUnderlyingSynchronizationContextCallbacks.Dequeue();
+        }
+    }
+
     /// <summary>
     /// Wraps the invocation of an async method such that it may
     /// execute asynchronously, but may potentially be
@@ -1411,4 +1645,19 @@ private void OnExecuting()
             }
         }
     }
+
+    private sealed class UnderlyingSynchronizationContextCallback
+    {
+        private JoinableTaskFactory? factory;
+
+        internal UnderlyingSynchronizationContextCallback(JoinableTaskFactory factory)
+        {
+            this.factory = factory;
+        }
+
+        internal void Execute()
+        {
+            Interlocked.Exchange(ref this.factory, null)?.ExecuteOnePendingUnderlyingSynchronizationContextCallback();
+        }
+    }
 }