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.