Reduce JTF dependency graph lock contention

src/Microsoft.VisualStudio.Threading/InternalUtilities.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Collections.Generic;
+using System.Threading;
 
 namespace Microsoft.VisualStudio.Threading;
 
@@ -10,6 +11,50 @@ namespace Microsoft.VisualStudio.Threading;
 /// </summary>
 internal static class InternalUtilities
 {
+    /// <summary>
+    /// Takes a list from a single-slot pool, allocating only when the slot is empty or its capacity is insufficient.
+    /// </summary>
+    /// <typeparam name="T">The type of list elements.</typeparam>
+    /// <param name="pool">The slot containing an available, empty list.</param>
+    /// <param name="capacity">The minimum required capacity.</param>
+    /// <returns>An exclusively owned, empty list with the requested capacity.</returns>
+    internal static List<T> RentList<T>(ref List<T>? pool, int capacity)
+    {
+        Requires.Range(capacity >= 0, nameof(capacity));
+
+        List<T>? list = Interlocked.Exchange(ref pool, null);
+        if (list is null)
+        {
+            return new List<T>(capacity);
+        }
+
+        if (list.Capacity < capacity)
+        {
+            list.Capacity = capacity;
+        }
+
+        return list;
+    }
+
+    /// <summary>
+    /// Clears a list and returns it to a single-slot pool if its capacity is within the retention limit.
+    /// </summary>
+    /// <typeparam name="T">The type of list elements.</typeparam>
+    /// <param name="pool">The slot to return the list to if empty.</param>
+    /// <param name="list">The exclusively owned list to return.</param>
+    /// <param name="maximumRetainedCapacity">The maximum capacity of a list to retain.</param>
+    internal static void ReturnList<T>(ref List<T>? pool, List<T> list, int maximumRetainedCapacity)
+    {
+        Requires.NotNull(list, nameof(list));
+        Requires.Range(maximumRetainedCapacity >= 0, nameof(maximumRetainedCapacity));
+
+        list.Clear();
+        if (list.Capacity <= maximumRetainedCapacity)
+        {
+            Interlocked.CompareExchange(ref pool, list, null);
+        }
+    }
+
     /// <summary>
     /// Removes an element from the middle of a queue without disrupting the other elements.
     /// </summary>

src/Microsoft.VisualStudio.Threading/JoinableTask.cs

@@ -275,6 +275,17 @@ internal bool IsFullyCompleted
     {
         get
         {
+            JoinableTaskFlags currentState = this.state;
+            if ((currentState & JoinableTaskFlags.CompleteFinalized) != 0)
+            {
+                return true;
+            }
+
+            if ((currentState & JoinableTaskFlags.CompleteRequested) == 0)
+            {
+                return false;
+            }
+
             using (this.JoinableTaskContext.NoMessagePumpSynchronizationContext.Apply())
             {
                 lock (this.JoinableTaskContext.SyncContextLock)

src/Microsoft.VisualStudio.Threading/JoinableTaskDependencyGraph.cs

@@ -16,8 +16,12 @@ namespace Microsoft.VisualStudio.Threading;
 /// </summary>
 internal static class JoinableTaskDependencyGraph
 {
+    private const int MaxRetainedNotificationListCapacity = 64;
+
     private static readonly HashSet<IJoinableTaskDependent> EmptySet = new HashSet<IJoinableTaskDependent>();
 
+    private static List<AsyncManualResetEvent>? notificationListPool;
+
     /// <summary>
     /// Gets a value indicating whether there is no child depenent item.
     /// This method is expected to be used with the JTF lock.
@@ -297,7 +301,7 @@ internal struct JoinableTaskDependentData
         /// <remarks>
         /// When the value in an entry is decremented to 0, the entry is removed from the map.
         /// </remarks>
-        private Dictionary<IJoinableTaskDependent, int> childDependentNodes;
+        private SmallDependencyMap<IJoinableTaskDependent> childDependentNodes;
 
         /// <summary>
         /// The head of a singly linked list of records to track which task may process events of this task.
@@ -348,58 +352,56 @@ internal static JoinableTaskCollection.JoinRelease AddDependency(IJoinableTaskDe
             using (parentTaskOrCollection.JoinableTaskContext.NoMessagePumpSynchronizationContext.Apply())
             {
                 List<AsyncManualResetEvent>? eventsNeedNotify = null;
-                lock (parentTaskOrCollection.JoinableTaskContext.SyncContextLock)
+                try
                 {
-                    var joinableTask = joinChild as JoinableTask;
-                    if (joinableTask?.IsFullyCompleted == true)
+                    lock (parentTaskOrCollection.JoinableTaskContext.SyncContextLock)
                     {
-                        return default(JoinableTaskCollection.JoinRelease);
-                    }
+                        var joinableTask = joinChild as JoinableTask;
+                        if (joinableTask?.IsFullyCompleted == true)
+                        {
+                            return default(JoinableTaskCollection.JoinRelease);
+                        }
 
-                    ref JoinableTaskDependentData data = ref parentTaskOrCollection.GetJoinableTaskDependentData();
-                    if (data.childDependentNodes is null)
-                    {
-                        data.childDependentNodes = new Dictionary<IJoinableTaskDependent, int>(capacity: 2);
-                    }
+                        ref JoinableTaskDependentData data = ref parentTaskOrCollection.GetJoinableTaskDependentData();
+                        if (data.childDependentNodes is null)
+                        {
+                            data.childDependentNodes = new SmallDependencyMap<IJoinableTaskDependent>();
+                        }
 
-                    if (data.childDependentNodes.TryGetValue(joinChild, out int refCount) && !parentTaskOrCollection.NeedRefCountChildDependencies)
-                    {
-                        return default(JoinableTaskCollection.JoinRelease);
+                        if (data.childDependentNodes.TryGetValue(joinChild, out int refCount) && !parentTaskOrCollection.NeedRefCountChildDependencies)
+                        {
+                            return default(JoinableTaskCollection.JoinRelease);
+                        }
+
+                        data.childDependentNodes[joinChild] = ++refCount;
+                        if (refCount == 1)
+                        {
+                            // This constitutes a significant change, so we should apply synchronous task tracking to the new child.
+                            joinChild.OnAddedToDependency(parentTaskOrCollection);
+                            AddDependingSynchronousTaskToChild(parentTaskOrCollection, joinChild, ref eventsNeedNotify);
+
+                            parentTaskOrCollection.OnDependencyAdded(joinChild);
+                        }
                     }
 
-                    data.childDependentNodes[joinChild] = ++refCount;
-                    if (refCount == 1)
+                    // We explicitly do this outside our lock.
+                    if (eventsNeedNotify is object)
                     {
-                        // This constitutes a significant change, so we should apply synchronous task tracking to the new child.
-                        joinChild.OnAddedToDependency(parentTaskOrCollection);
-                        IReadOnlyCollection<PendingNotification>? tasksNeedNotify = AddDependingSynchronousTaskToChild(parentTaskOrCollection, joinChild);
-                        if (tasksNeedNotify.Count > 0)
+                        foreach (AsyncManualResetEvent? queueEvent in eventsNeedNotify)
                         {
-                            eventsNeedNotify = new List<AsyncManualResetEvent>(tasksNeedNotify.Count);
-                            foreach (PendingNotification taskToNotify in tasksNeedNotify)
-                            {
-                                AsyncManualResetEvent? notifyEvent = taskToNotify.SynchronousTask.RegisterPendingEventsForSynchrousTask(taskToNotify.TaskHasPendingMessages, taskToNotify.NewPendingMessagesCount);
-                                if (notifyEvent is object)
-                                {
-                                    eventsNeedNotify.Add(notifyEvent);
-                                }
-                            }
+                            queueEvent.PulseAll();
                         }
-
-                        parentTaskOrCollection.OnDependencyAdded(joinChild);
                     }
-                }
 
-                // We explicitly do this outside our lock.
-                if (eventsNeedNotify is object)
+                    return new JoinableTaskCollection.JoinRelease(parentTaskOrCollection, joinChild);
+                }
+                finally
                 {
-                    foreach (AsyncManualResetEvent? queueEvent in eventsNeedNotify)
+                    if (eventsNeedNotify is object)
                     {
-                        queueEvent.PulseAll();
+                        InternalUtilities.ReturnList(ref notificationListPool, eventsNeedNotify, MaxRetainedNotificationListCapacity);
                     }
                 }
-
-                return new JoinableTaskCollection.JoinRelease(parentTaskOrCollection, joinChild);
             }
         }
 
@@ -478,7 +480,7 @@ internal static void AddSelfAndDescendentOrJoinedJobs(IJoinableTaskDependent tas
                 }
             }
 
-            Dictionary<IJoinableTaskDependent, int>? childDependentNodes = taskOrCollection.GetJoinableTaskDependentData().childDependentNodes;
+            SmallDependencyMap<IJoinableTaskDependent>? childDependentNodes = taskOrCollection.GetJoinableTaskDependentData().childDependentNodes;
             if (childDependentNodes is object)
             {
                 foreach (KeyValuePair<IJoinableTaskDependent, int> item in childDependentNodes)
@@ -562,7 +564,7 @@ internal static void ComputeSelfAndDescendentOrJoinedJobsAndRemainTasks(IJoinabl
                         return;
                     }
 
-                    Dictionary<IJoinableTaskDependent, int>? dependencies = taskOrCollection.GetJoinableTaskDependentData().childDependentNodes;
+                    SmallDependencyMap<IJoinableTaskDependent>? dependencies = taskOrCollection.GetJoinableTaskDependentData().childDependentNodes;
                     if (dependencies is object)
                     {
                         foreach (KeyValuePair<IJoinableTaskDependent, int> item in dependencies)
@@ -690,7 +692,7 @@ internal void OnTaskCompleted(IJoinableTaskDependent thisDependentNode)
 
                 if (this.childDependentNodes is object)
                 {
-                    Dictionary<IJoinableTaskDependent, int>.KeyCollection? childrenTasks = this.childDependentNodes.Keys;
+                    IReadOnlyCollection<IJoinableTaskDependent> childrenTasks = this.childDependentNodes.Keys;
                     while (existingTaskTracking is object)
                     {
                         RemoveDependingSynchronousTaskFrom(childrenTasks, existingTaskTracking.SynchronousTask, force: existingTaskTracking.SynchronousTask == thisDependentNode);
@@ -763,8 +765,8 @@ internal IReadOnlyCollection<JoinableTask> GetDependingSynchronousTasks(bool for
         /// </summary>
         /// <param name="dependentNode">The current joinableTask or collection owns the data.</param>
         /// <param name="child">The new child task.</param>
-        /// <returns>Pairs of synchronous tasks we need notify and the event source triggering it, plus the number of pending events.</returns>
-        private static IReadOnlyCollection<PendingNotification> AddDependingSynchronousTaskToChild(IJoinableTaskDependent dependentNode, IJoinableTaskDependent child)
+        /// <param name="eventsNeedNotify">Receives the events to pulse after releasing the JTF lock.</param>
+        private static void AddDependingSynchronousTaskToChild(IJoinableTaskDependent dependentNode, IJoinableTaskDependent child, ref List<AsyncManualResetEvent>? eventsNeedNotify)
         {
             Requires.NotNull(dependentNode, nameof(dependentNode));
             Requires.NotNull(child, nameof(child));
@@ -774,24 +776,26 @@ private static IReadOnlyCollection<PendingNotification> AddDependingSynchronousT
             int count = data.CountOfDependingSynchronousTasks();
             if (count == 0)
             {
-                return Array.Empty<PendingNotification>();
+                return;
             }
 
-            var tasksNeedNotify = new List<PendingNotification>(count);
             DependentSynchronousTask? existingTaskTracking = data.dependingSynchronousTaskTracking;
             while (existingTaskTracking is object)
             {
                 int totalEventNumber = 0;
                 JoinableTask? eventTriggeringTask = AddDependingSynchronousTask(child, existingTaskTracking.SynchronousTask, ref totalEventNumber);
                 if (eventTriggeringTask is object)
                 {
-                    tasksNeedNotify.Add(new PendingNotification(existingTaskTracking.SynchronousTask, eventTriggeringTask, totalEventNumber));
+                    AsyncManualResetEvent? notifyEvent = existingTaskTracking.SynchronousTask.RegisterPendingEventsForSynchrousTask(eventTriggeringTask, totalEventNumber);
+                    if (notifyEvent is object)
+                    {
+                        eventsNeedNotify ??= InternalUtilities.RentList(ref notificationListPool, count);
+                        eventsNeedNotify.Add(notifyEvent);
+                    }
                 }
 
                 existingTaskTracking = existingTaskTracking.Next;
             }
-
-            return tasksNeedNotify;
         }
 
         /// <summary>
@@ -1051,38 +1055,6 @@ private void RemoveDependingSynchronousTaskFromChild(IJoinableTaskDependent chil
             }
         }
 
-        /// <summary>
-        /// The record of a pending notification we need send to the synchronous task that we have some new messages to process.
-        /// </summary>
-        private readonly struct PendingNotification
-        {
-            internal PendingNotification(JoinableTask synchronousTask, JoinableTask taskHasPendingMessages, int newPendingMessagesCount)
-            {
-                Requires.NotNull(synchronousTask, nameof(synchronousTask));
-                Requires.NotNull(taskHasPendingMessages, nameof(taskHasPendingMessages));
-
-                this.SynchronousTask = synchronousTask;
-                this.TaskHasPendingMessages = taskHasPendingMessages;
-                this.NewPendingMessagesCount = newPendingMessagesCount;
-            }
-
-            /// <summary>
-            /// Gets the synchronous task which need process new messages.
-            /// </summary>
-            internal JoinableTask SynchronousTask { get; }
-
-            /// <summary>
-            /// Gets one JoinableTask which may have pending messages. We may have multiple new JoinableTasks which contains pending messages.
-            /// This is just one of them.  It gives the synchronous task a way to start quickly without searching all messages.
-            /// </summary>
-            internal JoinableTask TaskHasPendingMessages { get; }
-
-            /// <summary>
-            /// Gets the total number of new pending messages.  The real number could be less than that, but should not be more than that.
-            /// </summary>
-            internal int NewPendingMessagesCount { get; }
-        }
-
         /// <summary>
         /// A single linked list to maintain synchronous JoinableTask depends on the current task,
         ///  which may process the queue of the current task.