diff --git a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/TaskAwaiter.cs b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/TaskAwaiter.cs
index cb7b2c27d4b2b..36ae29e77cdfa 100644
--- a/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/TaskAwaiter.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/TaskAwaiter.cs
@@ -156,15 +156,27 @@ private static void HandleNonSuccessAndDebuggerNotification(Task task)
/// The awaited task.
///
///
+ ///
[StackTraceHidden]
- internal static void ValidateEnd(Task task, AwaitBehavior awaitBehavior, CancellationToken cancellationToken)
+ internal static void ValidateEnd(Task task, AwaitBehavior awaitBehavior, CancellationToken cancellationToken, bool disposeCts)
{
- // Fast checks that can be inlined.
- if (task.IsWaitNotificationEnabledOrNotRanToCompletion)
+ try
{
- // If either the end await bit is set or we're not completed successfully,
- // fall back to the slower path.
- HandleNonSuccessAndDebuggerNotification(task, awaitBehavior, cancellationToken);
+ if (task.IsWaitNotificationEnabledOrNotRanToCompletion)
+ {
+ // If either the end await bit is set or we're not completed successfully,
+ // fall back to the slower path.
+ HandleNonSuccessAndDebuggerNotification(task, awaitBehavior, cancellationToken);
+ }
+ }
+ finally
+ {
+ // can probably avoid try/finally, using in the PoC for brevity
+ if (disposeCts)
+ {
+ Debug.Assert(cancellationToken._source != null);
+ cancellationToken._source.Dispose();
+ }
}
}
@@ -625,10 +637,11 @@ public readonly struct ConfiguredCancelableTaskAwaitable
/// true to attempt to marshal the continuation back to the original context captured; otherwise, false.
///
///
- internal ConfiguredCancelableTaskAwaitable(Task task, AwaitBehavior awaitBehavior, CancellationToken cancellationToken)
+ ///
+ internal ConfiguredCancelableTaskAwaitable(Task task, AwaitBehavior awaitBehavior, CancellationToken cancellationToken, bool disposeCancellationTokenSource)
{
Debug.Assert(task != null, "Constructing an awaitable requires a task to await.");
- m_configuredTaskAwaiter = new ConfiguredCancelableTaskAwaitable.ConfiguredCancelableTaskAwaiter(task, awaitBehavior, cancellationToken);
+ m_configuredTaskAwaiter = new ConfiguredCancelableTaskAwaitable.ConfiguredCancelableTaskAwaiter(task, awaitBehavior, cancellationToken, disposeCancellationTokenSource);
}
/// Gets an awaiter for this awaitable.
@@ -651,6 +664,8 @@ public ConfiguredCancelableTaskAwaitable.ConfiguredCancelableTaskAwaiter GetAwai
internal readonly AwaitBehavior m_awaitBehavior;
/// The cancellation token.
internal readonly CancellationToken m_cancellationToken;
+ /// The cancellation token source is owned by the awaiter.
+ internal readonly bool m_disposeCancellationTokenSource;
/// Initializes the .
/// The to await.
@@ -659,12 +674,14 @@ public ConfiguredCancelableTaskAwaitable.ConfiguredCancelableTaskAwaiter GetAwai
/// when BeginAwait is called; otherwise, false.
///
/// The to await.
- internal ConfiguredCancelableTaskAwaiter(Task task, AwaitBehavior awaitBehavior, CancellationToken cancellationToken)
+ ///
+ internal ConfiguredCancelableTaskAwaiter(Task task, AwaitBehavior awaitBehavior, CancellationToken cancellationToken, bool disposeCancellationTokenSource)
{
Debug.Assert(task != null, "Constructing an awaiter requires a task to await.");
m_task = task;
m_cancellationToken = cancellationToken;
m_awaitBehavior = awaitBehavior;
+ m_disposeCancellationTokenSource = disposeCancellationTokenSource;
}
/// Gets whether the task being awaited is completed.
@@ -705,8 +722,7 @@ public void UnsafeOnCompleted(Action continuation)
[StackTraceHidden]
public void GetResult()
{
- // task exceptions take precedence over token cancellation
- TaskAwaiter.ValidateEnd(m_task, m_awaitBehavior, m_cancellationToken);
+ TaskAwaiter.ValidateEnd(m_task, m_awaitBehavior, m_cancellationToken, m_disposeCancellationTokenSource);
}
}
}
@@ -810,7 +826,7 @@ public void UnsafeOnCompleted(Action continuation)
[StackTraceHidden]
public TResult GetResult()
{
- TaskAwaiter.ValidateEnd(m_task, m_awaitBehavior, m_cancellationToken);
+ TaskAwaiter.ValidateEnd(m_task, m_awaitBehavior, m_cancellationToken, disposeCts: false);
Debug.Assert(m_task.IsWaitNotificationEnabledOrNotRanToCompletion ||
!m_task.IsCompletedSuccessfully && (m_awaitBehavior & AwaitBehavior.NoThrow) == AwaitBehavior.NoThrow,
"Should only be used when the task completed successfully and there's no wait notification enabled OR " +
diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/CancellationToken.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/CancellationToken.cs
index c4ad9c658cc97..6bd14fc540aaa 100644
--- a/src/libraries/System.Private.CoreLib/src/System/Threading/CancellationToken.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/Threading/CancellationToken.cs
@@ -32,7 +32,7 @@ public readonly struct CancellationToken
// The backing TokenSource.
// if null, it implicitly represents the same thing as new CancellationToken(false).
// When required, it will be instantiated to reflect this.
- private readonly CancellationTokenSource? _source;
+ internal readonly CancellationTokenSource? _source;
// !! warning. If more fields are added, the assumptions in CreateLinkedToken may no longer be valid
///
diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs
index 3ff60db514e02..a21fc23aa9b85 100644
--- a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs
@@ -2452,7 +2452,16 @@ public ConfiguredTaskAwaitable ConfigureAwait(bool continueOnCapturedContext)
/// An object used to await this task.
public ConfiguredCancelableTaskAwaitable ConfigureAwait(CancellationToken cancellationToken)
{
- return new ConfiguredCancelableTaskAwaitable(this, AwaitBehavior.Default, cancellationToken);
+ return new ConfiguredCancelableTaskAwaitable(this, AwaitBehavior.Default, cancellationToken, disposeCancellationTokenSource: false);
+ }
+
+ /// Configures an awaiter used to await this .
+ ///
+ /// An object used to await this task.
+ public ConfiguredCancelableTaskAwaitable ConfigureAwait(TimeSpan timeout)
+ {
+ var cts = new CancellationTokenSource(timeout);
+ return new ConfiguredCancelableTaskAwaitable(this, AwaitBehavior.Default, cts.Token, disposeCancellationTokenSource: true);
}
/// Configures an awaiter used to await this .
@@ -2463,7 +2472,7 @@ public ConfiguredCancelableTaskAwaitable ConfigureAwait(CancellationToken cancel
/// An object used to await this task.
public ConfiguredCancelableTaskAwaitable ConfigureAwait(bool continueOnCapturedContext, CancellationToken cancellationToken)
{
- return new ConfiguredCancelableTaskAwaitable(this, continueOnCapturedContext ? AwaitBehavior.Default : AwaitBehavior.NoCapturedContext, cancellationToken);
+ return new ConfiguredCancelableTaskAwaitable(this, continueOnCapturedContext ? AwaitBehavior.Default : AwaitBehavior.NoCapturedContext, cancellationToken, disposeCancellationTokenSource: false);
}
/// Configures an awaiter used to await this .
@@ -2473,7 +2482,7 @@ public ConfiguredCancelableTaskAwaitable ConfigureAwait(bool continueOnCapturedC
/// An object used to await this task.
public ConfiguredCancelableTaskAwaitable ConfigureAwait(AwaitBehavior awaitBehavior)
{
- return new ConfiguredCancelableTaskAwaitable(this, awaitBehavior, default);
+ return new ConfiguredCancelableTaskAwaitable(this, awaitBehavior, default, disposeCancellationTokenSource: false);
}
/// Configures an awaiter used to await this .
@@ -2484,7 +2493,7 @@ public ConfiguredCancelableTaskAwaitable ConfigureAwait(AwaitBehavior awaitBehav
/// An object used to await this task.
public ConfiguredCancelableTaskAwaitable ConfigureAwait(AwaitBehavior awaitBehavior, CancellationToken cancellationToken)
{
- return new ConfiguredCancelableTaskAwaitable(this, awaitBehavior, cancellationToken);
+ return new ConfiguredCancelableTaskAwaitable(this, awaitBehavior, cancellationToken, disposeCancellationTokenSource: false);
}
///
@@ -5616,8 +5625,12 @@ internal CancelableTaskWrapper(Task task, CancellationToken token)
}, this);
}
- private void Cleanup() => _registration.Dispose();
+ private void Cleanup()
+ {
+ _registration.Dispose();
+ }
}
+
#endregion
#region WhenAll