Merge pull request #792 from microsoft/dev/lifengl/joinableTaskCleanup

JoinTillEmptyAsync hangs for ref-counted JTF collection

Author: lifengl

src/Microsoft.VisualStudio.Threading/JoinableTask.cs

@@ -995,7 +995,7 @@ internal void OnQueueCompleted()
 
                 foreach (IJoinableTaskDependent? collection in this.dependencyParents)
                 {
-                    JoinableTaskDependencyGraph.RemoveDependency(collection, this);
+                    JoinableTaskDependencyGraph.RemoveDependency(collection, this, forceCleanup: true);
                 }
 
                 if (this.mainThreadJobSyncContext is object)

src/Microsoft.VisualStudio.Threading/JoinableTaskDependencyGraph.cs

@@ -71,10 +71,11 @@ internal static JoinableTaskCollection.JoinRelease AddDependency(IJoinableTaskDe
         /// </summary>
         /// <param name="taskItem">The current joinableTask or collection.</param>
         /// <param name="child">The <see cref="IJoinableTaskDependent"/> to join as a child.</param>
-        internal static void RemoveDependency(IJoinableTaskDependent taskItem, IJoinableTaskDependent child)
+        /// <param name="forceCleanup">Ignore refCount, it is being used when the child task is completed.</param>
+        internal static void RemoveDependency(IJoinableTaskDependent taskItem, IJoinableTaskDependent child, bool forceCleanup = false)
         {
             Requires.NotNull(taskItem, nameof(taskItem));
-            JoinableTaskDependentData.RemoveDependency(taskItem, child);
+            JoinableTaskDependentData.RemoveDependency(taskItem, child, forceCleanup);
         }
 
         /// <summary>
@@ -400,7 +401,8 @@ internal static JoinableTaskCollection.JoinRelease AddDependency(IJoinableTaskDe
             /// </summary>
             /// <param name="parentTaskOrCollection">The current joinableTask or collection contains to remove a dependency.</param>
             /// <param name="joinChild">The <see cref="IJoinableTaskDependent"/> to join as a child.</param>
-            internal static void RemoveDependency(IJoinableTaskDependent parentTaskOrCollection, IJoinableTaskDependent joinChild)
+            /// <param name="forceCleanup">Ignore refCount, it is being used when the child task is completed.</param>
+            internal static void RemoveDependency(IJoinableTaskDependent parentTaskOrCollection, IJoinableTaskDependent joinChild, bool forceCleanup)
             {
                 Requires.NotNull(parentTaskOrCollection, nameof(parentTaskOrCollection));
                 Requires.NotNull(joinChild, nameof(joinChild));
@@ -412,7 +414,7 @@ internal static void RemoveDependency(IJoinableTaskDependent parentTaskOrCollect
                     {
                         if (data.childDependentNodes is object && data.childDependentNodes.TryGetValue(joinChild, out int refCount))
                         {
-                            if (refCount == 1)
+                            if (refCount == 1 || forceCleanup)
                             {
                                 joinChild.OnRemovedFromDependency(parentTaskOrCollection);
 

test/Microsoft.VisualStudio.Threading.Tests/JoinableTaskCollectionTests.cs

@@ -196,4 +196,21 @@ public void JoinDisposedTwice()
             return Task.CompletedTask;
         });
     }
+
+    [Fact]
+    public void JoinTillEmptyWorksWithRefCounting()
+    {
+        var finishTaskEvent = new AsyncManualResetEvent();
+        JoinableTask task = this.JoinableFactory.RunAsync(async delegate { await finishTaskEvent.WaitAsync().ConfigureAwait(false); });
+
+        var collection = new JoinableTaskCollection(this.context, refCountAddedJobs: true);
+
+        collection.Add(task);
+        collection.Add(task);
+
+        finishTaskEvent.Set();
+
+        Task waiter = collection.JoinTillEmptyAsync();
+        Assert.True(waiter.Wait(UnexpectedTimeout));
+    }
 }