From c905111347978e29c17eda67415448f3753ca5bb Mon Sep 17 00:00:00 2001 From: Andrew Arnott Date: Wed, 26 Jun 2024 08:47:43 -0600 Subject: [PATCH] Add `AsyncBarrier.SignalAndWait(CancellationToken)` overload Also optimize the pre-existing code. The new overload returns `ValueTask` to leave open the option of optimizing to zero-allocation use via `IValueTaskSource`. --- .../AsyncBarrier.cs | 62 ++++++++++++++----- .../net472/PublicAPI.Unshipped.txt | 1 + .../net6.0-windows/PublicAPI.Unshipped.txt | 1 + .../net6.0/PublicAPI.Unshipped.txt | 1 + .../netstandard2.0/PublicAPI.Unshipped.txt | 1 + .../AsyncBarrierTests.cs | 32 ++++++++++ 6 files changed, 84 insertions(+), 14 deletions(-) diff --git a/src/Microsoft.VisualStudio.Threading/AsyncBarrier.cs b/src/Microsoft.VisualStudio.Threading/AsyncBarrier.cs index 292cf0430..92d273663 100644 --- a/src/Microsoft.VisualStudio.Threading/AsyncBarrier.cs +++ b/src/Microsoft.VisualStudio.Threading/AsyncBarrier.cs @@ -24,7 +24,7 @@ public class AsyncBarrier /// /// The set of participants who have reached the barrier, with their awaiters that can resume those participants. /// - private readonly Stack> waiters; + private readonly Stack waiters; /// /// Initializes a new instance of the class. @@ -37,15 +37,22 @@ public AsyncBarrier(int participants) // Allocate the stack so no resizing is necessary. // We don't need space for the last participant, since we never have to store it. - this.waiters = new Stack>(participants - 1); + this.waiters = new Stack(participants - 1); } + /// + public Task SignalAndWait() => this.SignalAndWait(CancellationToken.None).AsTask(); + /// /// Signals that a participant is ready, and returns a Task /// that completes when all other participants have also signaled ready. /// - /// A Task, which will complete (or may already be completed) when the last participant calls this method. - public Task SignalAndWait() + /// + /// A token that signals the caller's lost interest in waiting. + /// The signal effect of the method is not canceled with the token. + /// + /// A task which will complete (or may already be completed) when the last participant calls this method. + public ValueTask SignalAndWait(CancellationToken cancellationToken) { lock (this.waiters) { @@ -55,24 +62,51 @@ public Task SignalAndWait() // Unleash everyone that preceded this one. while (this.waiters.Count > 0) { - Task.Factory.StartNew( - state => ((TaskCompletionSource)state!).SetResult(default(EmptyStruct)), - this.waiters.Pop(), - CancellationToken.None, - TaskCreationOptions.None, - TaskScheduler.Default); + Waiter waiter = this.waiters.Pop(); + waiter.CompletionSource.TrySetResult(default); + waiter.CancellationRegistration.Dispose(); } // And allow this one to continue immediately. - return Task.CompletedTask; + return new ValueTask(cancellationToken.IsCancellationRequested + ? Task.FromCanceled(cancellationToken) + : Task.CompletedTask); } else { // We need more folks. So suspend this caller. - var tcs = new TaskCompletionSource(); - this.waiters.Push(tcs); - return tcs.Task; + TaskCompletionSource tcs = new(TaskCreationOptions.RunContinuationsAsynchronously); + CancellationTokenRegistration ctr; + if (cancellationToken.CanBeCanceled) + { +#if NET + ctr = cancellationToken.Register( + static (tcs, ct) => ((TaskCompletionSource)tcs!).TrySetCanceled(ct), tcs); +#else + ctr = cancellationToken.Register( + static s => + { + var t = (Tuple, CancellationToken>)s!; + t.Item1.TrySetCanceled(t.Item2); + }, + Tuple.Create(tcs, cancellationToken)); +#endif + } + else + { + ctr = default; + } + + this.waiters.Push(new Waiter(tcs, ctr)); + return new ValueTask(tcs.Task); } } } + + private readonly struct Waiter(TaskCompletionSource completionSource, CancellationTokenRegistration cancellationRegistration) + { + internal readonly TaskCompletionSource CompletionSource => completionSource; + + internal readonly CancellationTokenRegistration CancellationRegistration => cancellationRegistration; + } } diff --git a/src/Microsoft.VisualStudio.Threading/net472/PublicAPI.Unshipped.txt b/src/Microsoft.VisualStudio.Threading/net472/PublicAPI.Unshipped.txt index e69de29bb..bb9b5f944 100644 --- a/src/Microsoft.VisualStudio.Threading/net472/PublicAPI.Unshipped.txt +++ b/src/Microsoft.VisualStudio.Threading/net472/PublicAPI.Unshipped.txt @@ -0,0 +1 @@ +Microsoft.VisualStudio.Threading.AsyncBarrier.SignalAndWait(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask \ No newline at end of file diff --git a/src/Microsoft.VisualStudio.Threading/net6.0-windows/PublicAPI.Unshipped.txt b/src/Microsoft.VisualStudio.Threading/net6.0-windows/PublicAPI.Unshipped.txt index e69de29bb..bb9b5f944 100644 --- a/src/Microsoft.VisualStudio.Threading/net6.0-windows/PublicAPI.Unshipped.txt +++ b/src/Microsoft.VisualStudio.Threading/net6.0-windows/PublicAPI.Unshipped.txt @@ -0,0 +1 @@ +Microsoft.VisualStudio.Threading.AsyncBarrier.SignalAndWait(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask \ No newline at end of file diff --git a/src/Microsoft.VisualStudio.Threading/net6.0/PublicAPI.Unshipped.txt b/src/Microsoft.VisualStudio.Threading/net6.0/PublicAPI.Unshipped.txt index e69de29bb..bb9b5f944 100644 --- a/src/Microsoft.VisualStudio.Threading/net6.0/PublicAPI.Unshipped.txt +++ b/src/Microsoft.VisualStudio.Threading/net6.0/PublicAPI.Unshipped.txt @@ -0,0 +1 @@ +Microsoft.VisualStudio.Threading.AsyncBarrier.SignalAndWait(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask \ No newline at end of file diff --git a/src/Microsoft.VisualStudio.Threading/netstandard2.0/PublicAPI.Unshipped.txt b/src/Microsoft.VisualStudio.Threading/netstandard2.0/PublicAPI.Unshipped.txt index e69de29bb..bb9b5f944 100644 --- a/src/Microsoft.VisualStudio.Threading/netstandard2.0/PublicAPI.Unshipped.txt +++ b/src/Microsoft.VisualStudio.Threading/netstandard2.0/PublicAPI.Unshipped.txt @@ -0,0 +1 @@ +Microsoft.VisualStudio.Threading.AsyncBarrier.SignalAndWait(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.ValueTask \ No newline at end of file diff --git a/test/Microsoft.VisualStudio.Threading.Tests/AsyncBarrierTests.cs b/test/Microsoft.VisualStudio.Threading.Tests/AsyncBarrierTests.cs index 983667e81..5b9c48559 100644 --- a/test/Microsoft.VisualStudio.Threading.Tests/AsyncBarrierTests.cs +++ b/test/Microsoft.VisualStudio.Threading.Tests/AsyncBarrierTests.cs @@ -42,6 +42,38 @@ public async Task ManyParticipantsAndSteps() await this.MultipleParticipantsHelperAsync(100, 50); } + [Fact] + public async Task SignalAndWait_PrecanceledButReady() + { + AsyncBarrier barrier = new(1); + CancellationToken precanceled = new(canceled: true); + OperationCanceledException ex = await Assert.ThrowsAnyAsync(async () => await barrier.SignalAndWait(precanceled)).WithCancellation(this.TimeoutToken); + Assert.Equal(precanceled, ex.CancellationToken); + } + + [Fact] + public async Task SignalAndWait_PrecanceledWhileWaiting() + { + AsyncBarrier barrier = new(2); + CancellationToken precanceled = new(canceled: true); + OperationCanceledException ex = await Assert.ThrowsAnyAsync(async () => await barrier.SignalAndWait(precanceled)).WithCancellation(this.TimeoutToken); + Assert.Equal(precanceled, ex.CancellationToken); + } + + [Fact] + public async Task SignalAndWait_CanceledLeavesSignalBehind() + { + AsyncBarrier barrier = new(2); + CancellationTokenSource cts = new(); + Task waiter1 = barrier.SignalAndWait(cts.Token).AsTask(); + cts.Cancel(); + OperationCanceledException ex = await Assert.ThrowsAnyAsync(() => waiter1).WithCancellation(this.TimeoutToken); + Assert.Equal(cts.Token, ex.CancellationToken); + + // Now test that the second awaiter gets in, even though the first was canceled. + await barrier.SignalAndWait().WithCancellation(this.TimeoutToken); + } + /// /// Verifies that with multiple threads constantly fulfilling the participant count /// and resetting and fulfilling it again, it still performs as expected.