Merge pull request #791 from microsoft/dev/lifengl/fixInconsistentDependencyLogic

An inconsistence whether dependencies of completed tasks should be included leads corrupted data structure

Author: lifengl

src/Microsoft.VisualStudio.Threading/JoinableTask.cs

@@ -940,6 +940,9 @@ internal void CompleteOnCurrentThread()
                             }
                             else if (tryAgainAfter is object)
                             {
+                                // prevent referencing tasks which may be GCed during the waiting cycle.
+                                visited?.Clear();
+
                                 ThreadingEventSource.Instance.WaitSynchronouslyStart();
                                 this.owner.WaitSynchronously(tryAgainAfter);
                                 ThreadingEventSource.Instance.WaitSynchronouslyStop();
@@ -1070,7 +1073,7 @@ private static bool TryDequeueSelfOrDependencies(IJoinableTaskDependent currentN
 
                 if (work is null)
                 {
-                    if (joinableTask?.IsFullyCompleted != true)
+                    if (joinableTask?.IsCompleteRequested != true)
                     {
                         foreach (IJoinableTaskDependent? item in JoinableTaskDependencyGraph.GetDirectDependentNodes(currentNode))
                         {

src/Microsoft.VisualStudio.Threading/JoinableTaskDependencyGraph.cs

@@ -242,7 +242,7 @@ internal static bool CleanUpPotentialUnreachableDependentItems(JoinableTask sync
             if (possibleUnreachableItems is object && possibleUnreachableItems.Count > 0)
             {
                 var reachableNodes = new HashSet<IJoinableTaskDependent>();
-                var syncTaskItem = (IJoinableTaskDependent)syncTask;
+                IJoinableTaskDependent syncTaskItem = syncTask;
 
                 JoinableTaskDependentData.ComputeSelfAndDescendentOrJoinedJobsAndRemainTasks(syncTaskItem, reachableNodes, possibleUnreachableItems);
 
@@ -455,7 +455,17 @@ internal static void AddSelfAndDescendentOrJoinedJobs(IJoinableTaskDependent tas
 
                 if (taskOrCollection is JoinableTask thisJoinableTask)
                 {
-                    if (thisJoinableTask.IsFullyCompleted || !joinables.Add(thisJoinableTask))
+                    if (thisJoinableTask.IsCompleteRequested)
+                    {
+                        if (!thisJoinableTask.IsFullyCompleted)
+                        {
+                            joinables.Add(thisJoinableTask);
+                        }
+
+                        return;
+                    }
+
+                    if (!joinables.Add(thisJoinableTask))
                     {
                         return;
                     }
@@ -485,7 +495,6 @@ internal static void OnSynchronousTaskStartToBlockWaiting(JoinableTask syncTask,
                 Requires.NotNull(syncTask, nameof(syncTask));
 
                 pendingRequestsCount = 0;
-                taskHasPendingRequests = null;
 
                 taskHasPendingRequests = AddDependingSynchronousTask(syncTask, syncTask, ref pendingRequestsCount);
             }
@@ -504,7 +513,11 @@ internal static void OnSynchronousTaskEndToBlockWaiting(JoinableTask syncTask)
                     lock (syncTask.Factory.Context.SyncContextLock)
                     {
                         // Remove itself from the tracking list, after the task is completed.
-                        RemoveDependingSynchronousTask(syncTask, syncTask, force: true);
+                        IJoinableTaskDependent syncTaskItem = syncTask;
+                        if (syncTaskItem.GetJoinableTaskDependentData().dependingSynchronousTaskTracking is object)
+                        {
+                            RemoveDependingSynchronousTask(syncTask, syncTask, force: true);
+                        }
 
                         if (syncTask.PotentialUnreachableDependents is object && syncTask.PotentialUnreachableDependents.Count > 0)
                         {
@@ -537,6 +550,11 @@ internal static void ComputeSelfAndDescendentOrJoinedJobsAndRemainTasks(IJoinabl
                             return;
                         }
 
+                        if ((taskOrCollection as JoinableTask)?.IsCompleteRequested == true)
+                        {
+                            return;
+                        }
+
                         WeakKeyDictionary<IJoinableTaskDependent, int>? dependencies = taskOrCollection.GetJoinableTaskDependentData().childDependentNodes;
                         if (dependencies is object)
                         {
@@ -846,11 +864,7 @@ private static void RemoveDependingSynchronousTask(IJoinableTaskDependent taskOr
                 Requires.NotNull(syncTask, nameof(syncTask));
                 Assumes.True(Monitor.IsEntered(taskOrCollection.JoinableTaskContext.SyncContextLock));
 
-                var syncTaskItem = (IJoinableTaskDependent)syncTask;
-                if (syncTaskItem.GetJoinableTaskDependentData().dependingSynchronousTaskTracking is object)
-                {
-                    RemoveDependingSynchronousTaskFrom(new IJoinableTaskDependent[] { taskOrCollection }, syncTask, force);
-                }
+                RemoveDependingSynchronousTaskFrom(new IJoinableTaskDependent[] { taskOrCollection }, syncTask, force);
             }
 
             /// <summary>

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

@@ -4116,6 +4116,44 @@ public void GetAwaiterCompletedTDoesNotLock()
         lockHolder.Wait();
     }
 
+    [Fact]
+    public void JoinableTaskCleanUpDependenciesForLoopDependencies()
+    {
+        this.SimulateUIThread(async delegate
+        {
+            var joinableTaskCollection = new JoinableTaskCollection(this.context, refCountAddedJobs: true);
+            JoinableTaskFactory taskFactory = this.context.CreateFactory(joinableTaskCollection);
+            ((DerivedJoinableTaskFactory)taskFactory).AssumeConcurrentUse = true;
+
+            bool isMainThreadBlockedResult = true;
+            var unblockTask2 = new AsyncManualResetEvent();
+            JoinableTask? childTask = null;
+
+            this.context.Factory.Run(async () =>
+            {
+                await TaskScheduler.Default.SwitchTo(alwaysYield: true);
+
+                childTask = await this.SpinOffMainThreadTaskForJoinableTaskDependenciesHandledAfterTaskCompletion(taskFactory, joinableTaskCollection, unblockTask2);
+            });
+
+            using (this.context.SuppressRelevance())
+            {
+                isMainThreadBlockedResult = await taskFactory.RunAsync(() =>
+                {
+                    return Task.FromResult(this.context.IsMainThreadBlocked());
+                });
+            }
+
+            using (this.context.SuppressRelevance())
+            {
+                unblockTask2.Set();
+                childTask!.Task.Wait();
+            }
+
+            Assert.False(isMainThreadBlockedResult);
+        });
+    }
+
     protected override JoinableTaskContext CreateJoinableTaskContext()
     {
         return new DerivedJoinableTaskContext();
@@ -4126,6 +4164,87 @@ private static async void SomeFireAndForgetMethod()
         await Task.Yield();
     }
 
+    private async Task<JoinableTask> SpinOffMainThreadTaskForJoinableTaskDependenciesHandledAfterTaskCompletion(
+        JoinableTaskFactory joinableTaskFactory,
+        JoinableTaskCollection joinableTaskCollection,
+        AsyncManualResetEvent unblockTask2)
+    {
+        JoinableTask? task2 = null;
+        Task? spinOffTask = null;
+
+        JoinableTask? joinableTask = null;
+
+        var mainThreadJoinedCollection = new JoinableTaskCollection(this.context);
+        await TaskScheduler.Default.SwitchTo(alwaysYield: true);
+
+        using (this.context.SuppressRelevance())
+        {
+            // this task creates a circular loop.
+            task2 = joinableTaskFactory.RunAsync(async () =>
+            {
+                joinableTaskCollection.Join();
+                await unblockTask2.WaitAsync();
+            });
+
+            var unblockTask1 = new AsyncManualResetEvent();
+            var spinOffIsReady = new AsyncManualResetEvent();
+            var spinOffIsDone = new AsyncManualResetEvent();
+
+            joinableTask = joinableTaskFactory.RunAsync(async () =>
+            {
+                joinableTaskCollection.Join();
+
+                spinOffTask = Task.Run(async () =>
+                {
+                    JoinableTaskFactory.MainThreadAwaiter awaiter = this.context.Factory.SwitchToMainThreadAsync().GetAwaiter();
+                    awaiter.OnCompleted(() =>
+                    {
+                        spinOffIsDone.Set();
+                    });
+
+                    spinOffIsReady.Set();
+
+                    await spinOffIsDone.WaitAsync();
+                });
+
+                await spinOffIsReady.WaitAsync();
+                await unblockTask1.WaitAsync();
+            });
+
+            // Increase refcount
+            mainThreadJoinedCollection.Add(joinableTask);
+            joinableTaskCollection.Add(joinableTask);
+
+            unblockTask1.Set();
+
+            await joinableTask.Task;
+        }
+
+        await this.context.Factory.SwitchToMainThreadAsync();
+
+        // Due to circular reference, this add/remove reference to lead incompleted state
+        // the joinableTaskCollection will retain refcount to the JTF.Run task.
+        using (joinableTaskCollection.Join())
+        {
+        }
+
+        using (mainThreadJoinedCollection.Join())
+        {
+            // This IsMainThreadBlocked triggers the incompleted state to be cleaned up.
+            // JTF.Run joins the completed task through the middle collection to expose the
+            // inconsistent between two logic, so the recomputation won't clean up the circular
+            // dependency loop correctly.
+            await this.context.Factory.RunAsync(() =>
+            {
+                return Task.FromResult(this.context.IsMainThreadBlocked());
+            });
+
+            await spinOffTask!;
+        }
+
+        return task2;
+    }
+
     private async Task SomeOperationThatMayBeOnMainThreadAsync()
     {
         await Task.Yield();