prototype awaitables with timeouts

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/TaskAwaiter.cs

@@ -156,15 +156,27 @@ private static void HandleNonSuccessAndDebuggerNotification(Task task)
         /// <param name="task">The awaited task.</param>
         /// <param name="awaitBehavior"></param>
         /// <param name="cancellationToken"></param>
+        /// <param name="disposeCts"></param>
         [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.
         /// </param>
         /// <param name="cancellationToken"></param>
-        internal ConfiguredCancelableTaskAwaitable(Task task, AwaitBehavior awaitBehavior, CancellationToken cancellationToken)
+        /// <param name="disposeCancellationTokenSource"></param>
+        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);
         }
 
         /// <summary>Gets an awaiter for this awaitable.</summary>
@@ -651,6 +664,8 @@ public ConfiguredCancelableTaskAwaitable.ConfiguredCancelableTaskAwaiter GetAwai
             internal readonly AwaitBehavior m_awaitBehavior;
             /// <summary>The cancellation token.</summary>
             internal readonly CancellationToken m_cancellationToken;
+            /// <summary>The cancellation token source is owned by the awaiter.</summary>
+            internal readonly bool m_disposeCancellationTokenSource;
 
             /// <summary>Initializes the <see cref="ConfiguredCancelableTaskAwaiter"/>.</summary>
             /// <param name="task">The <see cref="System.Threading.Tasks.Task"/> to await.</param>
@@ -659,12 +674,14 @@ public ConfiguredCancelableTaskAwaitable.ConfiguredCancelableTaskAwaiter GetAwai
             /// when BeginAwait is called; otherwise, false.
             /// </param>
             /// <param name="cancellationToken">The <see cref="System.Threading.Tasks.Task"/> to await.</param>
-            internal ConfiguredCancelableTaskAwaiter(Task task, AwaitBehavior awaitBehavior, CancellationToken cancellationToken)
+            /// <param name="disposeCancellationTokenSource"></param>
+            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;
             }
 
             /// <summary>Gets whether the task being awaited is completed.</summary>
@@ -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 " +

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
 
         /// <summary>

src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs

@@ -2452,7 +2452,16 @@ public ConfiguredTaskAwaitable ConfigureAwait(bool continueOnCapturedContext)
         /// <returns>An object used to await this task.</returns>
         public ConfiguredCancelableTaskAwaitable ConfigureAwait(CancellationToken cancellationToken)
         {
-            return new ConfiguredCancelableTaskAwaitable(this, AwaitBehavior.Default, cancellationToken);
+            return new ConfiguredCancelableTaskAwaitable(this, AwaitBehavior.Default, cancellationToken, disposeCancellationTokenSource: false);
+        }
+
+        /// <summary>Configures an awaiter used to await this <see cref="System.Threading.Tasks.Task"/>.</summary>
+        /// <param name="timeout"></param>
+        /// <returns>An object used to await this task.</returns>
+        public ConfiguredCancelableTaskAwaitable ConfigureAwait(TimeSpan timeout)
+        {
+            var cts = new CancellationTokenSource(timeout);
+            return new ConfiguredCancelableTaskAwaitable(this, AwaitBehavior.Default, cts.Token, disposeCancellationTokenSource: true);
         }
 
         /// <summary>Configures an awaiter used to await this <see cref="System.Threading.Tasks.Task"/>.</summary>
@@ -2463,7 +2472,7 @@ public ConfiguredCancelableTaskAwaitable ConfigureAwait(CancellationToken cancel
         /// <returns>An object used to await this task.</returns>
         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);
         }
 
         /// <summary>Configures an awaiter used to await this <see cref="System.Threading.Tasks.Task"/>.</summary>
@@ -2473,7 +2482,7 @@ public ConfiguredCancelableTaskAwaitable ConfigureAwait(bool continueOnCapturedC
         /// <returns>An object used to await this task.</returns>
         public ConfiguredCancelableTaskAwaitable ConfigureAwait(AwaitBehavior awaitBehavior)
         {
-            return new ConfiguredCancelableTaskAwaitable(this, awaitBehavior, default);
+            return new ConfiguredCancelableTaskAwaitable(this, awaitBehavior, default, disposeCancellationTokenSource: false);
         }
 
         /// <summary>Configures an awaiter used to await this <see cref="System.Threading.Tasks.Task"/>.</summary>
@@ -2484,7 +2493,7 @@ public ConfiguredCancelableTaskAwaitable ConfigureAwait(AwaitBehavior awaitBehav
         /// <returns>An object used to await this task.</returns>
         public ConfiguredCancelableTaskAwaitable ConfigureAwait(AwaitBehavior awaitBehavior, CancellationToken cancellationToken)
         {
-            return new ConfiguredCancelableTaskAwaitable(this, awaitBehavior, cancellationToken);
+            return new ConfiguredCancelableTaskAwaitable(this, awaitBehavior, cancellationToken, disposeCancellationTokenSource: false);
         }
 
         /// <summary>
@@ -5616,8 +5625,12 @@ internal CancelableTaskWrapper(Task task, CancellationToken token)
                 }, this);
             }
 
-            private void Cleanup() => _registration.Dispose();
+            private void Cleanup()
+            {
+                _registration.Dispose();
+            }
         }
+
         #endregion
 
         #region WhenAll