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