Coalesce JoinableTaskFactory dispatcher posts

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,31 @@ namespace Microsoft.VisualStudio.Threading;
 /// </remarks>
 public partial class JoinableTaskFactory
 {
+    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)>? pendingUnderlyingSynchronizationContextCallbacks;
+
+    private bool underlyingSynchronizationContextCallbackPending;
+
+    private TaskCompletionSource<object?>? underlyingSynchronizationContextPostCompletion;
+
     /// <summary>
     /// Backing field for the <see cref="HangDetectionTimeout"/> property.
     /// </summary>
@@ -470,6 +483,19 @@ protected internal virtual void PostToUnderlyingSynchronizationContext(SendOrPos
         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>
+    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 +660,51 @@ 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>
+    protected void PostToUnderlyingSynchronizationContextWithCoalescing(SendOrPostCallback callback, object state)
+    {
+        Requires.NotNull(callback, nameof(callback));
+
+        TaskCompletionSource<object?>? postCompletion = null;
+        Task? waitForPost = null;
+        lock (this.pendingUnderlyingSynchronizationContextCallbacksLock)
+        {
+            this.pendingUnderlyingSynchronizationContextCallbacks ??= new Queue<(SendOrPostCallback, object)>();
+            this.RemoveCompletedPendingUnderlyingSynchronizationContextCallbacks();
+            this.pendingUnderlyingSynchronizationContextCallbacks.Enqueue((callback, state));
+            if (!this.underlyingSynchronizationContextCallbackPending)
+            {
+                this.underlyingSynchronizationContextCallbackPending = true;
+                postCompletion = this.underlyingSynchronizationContextPostCompletion = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
+            }
+            else
+            {
+                waitForPost = this.underlyingSynchronizationContextPostCompletion?.Task;
+            }
+        }
+
+        if (postCompletion is object)
+        {
+            this.PostPendingUnderlyingSynchronizationContextCallback(postCompletion);
+        }
+        else if (!IsSynchronouslyPosting(this))
+        {
+            waitForPost?.GetAwaiter().GetResult();
+        }
+    }
+
+    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 +730,146 @@ 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)? callback = null;
+            TaskCompletionSource<object?>? postCompletion = null;
+            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)
+                    {
+                        postCompletion = this.underlyingSynchronizationContextPostCompletion = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
+                    }
+                    else
+                    {
+                        this.pendingUnderlyingSynchronizationContextCallbacks = null;
+                        this.underlyingSynchronizationContextCallbackPending = false;
+                    }
+                }
+
+                if (postCompletion is object)
+                {
+                    this.PostPendingUnderlyingSynchronizationContextCallback(postCompletion, propagateException: false);
+                }
+
+                if (callback is { } work)
+                {
+                    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(TaskCompletionSource<object?> postCompletion, bool propagateException = true)
+    {
+        try
+        {
+            List<JoinableTaskFactory> synchronousPostingChain = synchronouslyPostingFactories ??= new();
+            synchronousPostingChain.Add(this);
+            try
+            {
+                this.PostToUnderlyingSynchronizationContextCore(ExecuteOnePendingUnderlyingSynchronizationContextCallbackDelegate, new UnderlyingSynchronizationContextCallback(this));
+            }
+            finally
+            {
+                synchronousPostingChain.RemoveAt(synchronousPostingChain.Count - 1);
+            }
+
+            lock (this.pendingUnderlyingSynchronizationContextCallbacksLock)
+            {
+                if (this.underlyingSynchronizationContextPostCompletion == postCompletion)
+                {
+                    this.underlyingSynchronizationContextPostCompletion = null;
+                }
+            }
+
+            postCompletion.SetResult(null);
+        }
+        catch (Exception ex)
+        {
+            lock (this.pendingUnderlyingSynchronizationContextCallbacksLock)
+            {
+                if (this.underlyingSynchronizationContextPostCompletion == postCompletion)
+                {
+                    // 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;
+                    this.underlyingSynchronizationContextPostCompletion = null;
+                }
+            }
+
+            postCompletion.SetException(ex);
+            _ = postCompletion.Task.Exception;
+            if (propagateException)
+            {
+                throw;
+            }
+        }
+    }
+
+    private void RemoveCompletedPendingUnderlyingSynchronizationContextCallbacks()
+    {
+        Assumes.True(Monitor.IsEntered(this.pendingUnderlyingSynchronizationContextCallbacksLock));
+
+#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 +1622,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();
+        }
+    }
 }