Reduce unnecessary disposing JTF dependencies & task continuations

src/Microsoft.VisualStudio.Threading/JoinableTask.cs

@@ -590,13 +590,34 @@ private Task QueueNeedProcessEvent
         /// before the async operation has completed.
         /// </param>
         /// <returns>A task that completes after the asynchronous operation completes and the join is reverted.</returns>
-        public async Task JoinAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public Task JoinAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
             cancellationToken.ThrowIfCancellationRequested();
 
-            using (this.AmbientJobJoinsThis())
+            if (!cancellationToken.CanBeCanceled)
             {
-                await this.Task.WithCancellation(AwaitShouldCaptureSyncContext, cancellationToken).ConfigureAwait(AwaitShouldCaptureSyncContext);
+                // if a completed or failed JoinableTask will remove itself from parent dependency chains, so we don't repeat it which requires the sync lock.
+                _ = this.AmbientJobJoinsThis();
+                return this.Task;
+            }
+            else
+            {
+                return JoinSlowAsync(this, cancellationToken);
+            }
+
+            static async Task JoinSlowAsync(JoinableTask me, CancellationToken cancellationToken)
+            {
+                JoinRelease dependency = me.AmbientJobJoinsThis();
+
+                try
+                {
+                    await me.Task.WithCancellation(continueOnCapturedContext: AwaitShouldCaptureSyncContext, cancellationToken).ConfigureAwait(AwaitShouldCaptureSyncContext);
+                }
+                catch
+                {
+                    dependency.Dispose();
+                    throw;
+                }
             }
         }
 
@@ -1055,6 +1076,20 @@ internal int GetPendingEventCountForSynchronousTask(JoinableTask synchronousTask
             return this.queueNeedProcessEvent;
         }
 
+        private protected JoinRelease AmbientJobJoinsThis()
+        {
+            if (!this.IsCompleted)
+            {
+                JoinableTask? ambientJob = this.JoinableTaskContext.AmbientTask;
+                if (ambientJob is object && ambientJob != this)
+                {
+                    return JoinableTaskDependencyGraph.AddDependency(ambientJob, this);
+                }
+            }
+
+            return default(JoinRelease);
+        }
+
         private static bool TryDequeueSelfOrDependencies(IJoinableTaskDependent currentNode, bool onMainThread, HashSet<IJoinableTaskDependent> visited, [NotNullWhen(true)] out SingleExecuteProtector? work)
         {
             Requires.NotNull(currentNode, nameof(currentNode));
@@ -1191,19 +1226,5 @@ private void AddStateFlags(JoinableTaskFlags flags)
                 }
             }
         }
-
-        private JoinRelease AmbientJobJoinsThis()
-        {
-            if (!this.IsCompleted)
-            {
-                JoinableTask? ambientJob = this.JoinableTaskContext.AmbientTask;
-                if (ambientJob is object && ambientJob != this)
-                {
-                    return JoinableTaskDependencyGraph.AddDependency(ambientJob, this);
-                }
-            }
-
-            return default(JoinRelease);
-        }
     }
 }

src/Microsoft.VisualStudio.Threading/JoinableTask`1.cs

@@ -56,12 +56,29 @@ internal JoinableTask(JoinableTaskFactory owner, bool synchronouslyBlocking, Joi
                 return this.Task;
             }
 
-            return JoinSlowAsync(cancellationToken);
+            if (!cancellationToken.CanBeCanceled)
+            {
+                // if a completed or failed JoinableTask will remove itself from parent dependency chains, so we don't repeat it which requires the sync lock.
+                _ = this.AmbientJobJoinsThis();
+                return this.Task;
+            }
 
-            async Task<T> JoinSlowAsync(CancellationToken cancellationToken)
+            return JoinSlowAsync(this, cancellationToken);
+
+            static async Task<T> JoinSlowAsync(JoinableTask<T> me, CancellationToken cancellationToken)
             {
-                await base.JoinAsync(cancellationToken).ConfigureAwait(AwaitShouldCaptureSyncContext);
-                return await this.Task.ConfigureAwait(AwaitShouldCaptureSyncContext);
+                JoinableTaskCollection.JoinRelease dependency = me.AmbientJobJoinsThis();
+
+                try
+                {
+                    await me.Task.WithCancellation(continueOnCapturedContext: AwaitShouldCaptureSyncContext, cancellationToken).ConfigureAwait(AwaitShouldCaptureSyncContext);
+                    return await me.Task.ConfigureAwait(AwaitShouldCaptureSyncContext);
+                }
+                catch
+                {
+                    dependency.Dispose();
+                    throw;
+                }
             }
         }