Add a JTF-derived type that cancels a DisposalToken and drains the JoinableTaskCollection when disposed

src/Microsoft.VisualStudio.Threading/DelegatingJoinableTaskFactory.cs

@@ -1,10 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 

src/Microsoft.VisualStudio.Threading/DisposableJoinableTaskFactory.cs

@@ -0,0 +1,89 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.VisualStudio.Threading;
+
+/// <summary>
+/// A variant on <see cref="JoinableTaskFactory"/> that tracks pending tasks and blocks disposal until all tasks have completed.
+/// A cancellation token is provided so pending tasks can cooperatively cancel when disposal is requested.
+/// </summary>
+/// <remarks>
+/// <para>
+/// Cancellation of pending tasks is cooperative.
+/// If a pending task does not observe <see cref="DisposalToken" />, then disposal may take longer to complete,
+/// or even never complete if a pending task never completes.
+/// </para>
+/// <para>
+/// Creating tasks after disposal has been requested is not prevented by this class.
+/// </para>
+/// </remarks>
+public class DisposableJoinableTaskFactory : DelegatingJoinableTaskFactory, IDisposable, System.IAsyncDisposable
+{
+    private readonly CancellationTokenSource disposalTokenSource = new();
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="DisposableJoinableTaskFactory"/> class.
+    /// </summary>
+    /// <param name="innerFactory">The factory instance to be wrapped. Must have an associated collection.</param>
+    public DisposableJoinableTaskFactory(JoinableTaskFactory innerFactory)
+        : base(innerFactory)
+    {
+        Requires.Argument(this.Collection is not null, nameof(innerFactory), "A collection must be associated with the factory.");
+
+        // Get it now, since after the CTS is disposed, it throws when we try to access the token.
+        this.DisposalToken = this.disposalTokenSource.Token;
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="DisposableJoinableTaskFactory"/> class.
+    /// </summary>
+    /// <param name="joinableTaskContext">The <see cref="JoinableTaskContext"/> used to construct the <see cref="JoinableTaskFactory"/>.</param>
+    /// <remarks>
+    /// This constructor creates a <see cref="JoinableTaskFactory"/> using <see cref="JoinableTaskContext.CreateFactory(JoinableTaskCollection)"/>.
+    /// </remarks>
+    public DisposableJoinableTaskFactory(JoinableTaskContext joinableTaskContext)
+        : this(Requires.NotNull(joinableTaskContext).CreateFactory(Requires.NotNull(joinableTaskContext).CreateCollection()))
+    {
+    }
+
+    /// <summary>
+    /// Gets a disposal token that <em>should</em> be used by tasks created by this factory to know when they should stop doing work.
+    /// </summary>
+    /// <remarks>
+    /// This token is canceled when the factory is disposed.
+    /// </remarks>
+    public CancellationToken DisposalToken { get; }
+
+    /// <summary>
+    /// Gets the collection to which created tasks belong until they complete.
+    /// </summary>
+    protected new JoinableTaskCollection Collection => base.Collection!;
+
+    /// <inheritdoc/>
+    public void Dispose()
+    {
+        if (!this.disposalTokenSource.IsCancellationRequested)
+        {
+            this.disposalTokenSource.Cancel();
+            this.disposalTokenSource.Dispose();
+        }
+
+        this.Context.Factory.Run(() => this.Collection.JoinTillEmptyAsync());
+    }
+
+    /// <inheritdoc/>
+    public async ValueTask DisposeAsync()
+    {
+        if (!this.disposalTokenSource.IsCancellationRequested)
+        {
+            this.disposalTokenSource.Cancel();
+            this.disposalTokenSource.Dispose();
+        }
+
+        await this.Collection.JoinTillEmptyAsync();
+    }
+}

src/Microsoft.VisualStudio.Threading/JoinableTaskFactory.cs

@@ -93,7 +93,7 @@ internal SynchronizationContext? ApplicableJobSyncContext
     /// <summary>
     /// Gets the collection to which created tasks belong until they complete. May be null.
     /// </summary>
-    internal JoinableTaskCollection? Collection
+    protected internal JoinableTaskCollection? Collection
     {
         get { return this.jobCollection; }
     }

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

@@ -0,0 +1,84 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Threading.Tasks;
+
+public class DisposableJoinableTaskFactoryTests : TestBase
+{
+    private readonly JoinableTaskContext context;
+    private readonly JoinableTaskCollection joinableCollection;
+    private readonly JoinableTaskFactory asyncPump;
+
+    public DisposableJoinableTaskFactoryTests(ITestOutputHelper logger)
+        : base(logger)
+    {
+        this.context = JoinableTaskContext.CreateNoOpContext();
+        this.joinableCollection = this.context.CreateCollection();
+        this.asyncPump = this.context.CreateFactory(this.joinableCollection);
+    }
+
+    [Fact]
+    public void DisposeCancelsToken()
+    {
+        DisposableJoinableTaskFactory factory = new(this.asyncPump);
+
+        // The collection is already empty, so this completes immediately.
+        factory.Dispose();
+
+        Assert.True(factory.DisposalToken.IsCancellationRequested);
+    }
+
+    [Fact]
+    public async Task DisposeAsyncCancelsToken()
+    {
+        DisposableJoinableTaskFactory factory = new(this.asyncPump);
+
+        // The collection is already empty, so this completes immediately.
+        await factory.DisposeAsync();
+
+        Assert.True(factory.DisposalToken.IsCancellationRequested);
+    }
+
+    [Fact]
+    public void MultipleDisposalsDoNotThrow()
+    {
+        using DisposableJoinableTaskFactory factory = new(this.asyncPump);
+
+        factory.Dispose();
+        factory.Dispose();
+
+        Assert.True(factory.DisposalToken.IsCancellationRequested);
+    }
+
+    [Fact]
+    public async Task MultipleDisposeAsyncsDoesNotThrow()
+    {
+        DisposableJoinableTaskFactory factory = new(this.asyncPump);
+
+        await factory.DisposeAsync();
+        await factory.DisposeAsync();
+
+        Assert.True(factory.DisposalToken.IsCancellationRequested);
+    }
+
+    [Fact]
+    public void ConstructorRequiresCollection()
+    {
+        // A factory created with just a context has no collection.
+        JoinableTaskFactory factoryWithoutCollection = this.context.Factory;
+        Assert.ThrowsAny<Exception>(() => new DisposableJoinableTaskFactory(factoryWithoutCollection));
+    }
+
+    [Fact]
+    public async Task TaskObservesCancellationDuringDrain()
+    {
+        DisposableJoinableTaskFactory factory = new(this.asyncPump);
+
+        JoinableTask jt = factory.RunAsync(() => Task.Delay(UnexpectedTimeout, factory.DisposalToken));
+
+        await factory.DisposeAsync().AsTask().WithTimeout(UnexpectedTimeout);
+        Assert.True(jt.IsCompleted);
+        await Assert.ThrowsAsync<TaskCanceledException>(async () => await jt);
+    }
+}