Add TaskCompletionSource.SetFromTask

Author: stephentoub

src/libraries/System.Net.Requests/src/System/Net/TaskExtensions.cs

@@ -18,20 +18,7 @@ public static TaskCompletionSource<TResult> ToApm<TResult>(
 
             task.ContinueWith(completedTask =>
             {
-                bool shouldInvokeCallback = false;
-
-                if (completedTask.IsFaulted)
-                {
-                    shouldInvokeCallback = tcs.TrySetException(completedTask.Exception!.InnerExceptions);
-                }
-                else if (completedTask.IsCanceled)
-                {
-                    shouldInvokeCallback = tcs.TrySetCanceled();
-                }
-                else
-                {
-                    shouldInvokeCallback = tcs.TrySetResult(completedTask.Result);
-                }
+                bool shouldInvokeCallback = tcs.TrySetFromTask(completedTask);
 
                 // Only invoke the callback if it exists AND we were able to transition the TCS
                 // to the terminal state. If we couldn't transition the task it is because it was

src/libraries/System.Private.CoreLib/src/Resources/Strings.resx

@@ -3527,6 +3527,9 @@
   <data name="Task_WaitMulti_NullTask" xml:space="preserve">
     <value>The tasks array included at least one null element.</value>
   </data>
+  <data name="Task_MustBeCompleted" xml:space="preserve">
+    <value>The provided task must have already completed.</value>
+  </data>
   <data name="TaskT_ConfigureAwait_InvalidOptions" xml:space="preserve">
     <value>Task&lt;TResult&gt;.ConfigureAwait does not support ConfigureAwaitOptions.SuppressThrowing. To suppress throwing, instead cast the Task&lt;TResult&gt; to its base class Task and await that with SuppressThrowing.</value>
   </data>
@@ -4286,4 +4289,4 @@
   <data name="Reflection_Disabled" xml:space="preserve">
     <value>This operation is not available because the reflection support was disabled at compile time.</value>
   </data>
-</root>
+</root>

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

@@ -285,5 +285,75 @@ public bool TrySetCanceled(CancellationToken cancellationToken)
 
             return rval;
         }
+
+        /// <summary>
+        /// Transition the underlying <see cref="Task{TResult}"/> into the same completion state as the specified <paramref name="completedTask"/>.
+        /// </summary>
+        /// <param name="completedTask">The completed task whose completion status (including exception or cancellation information) should be copied to the underlying task.</param>
+        /// <exception cref="ArgumentNullException"><paramref name="completedTask"/> is <see langword="null"/>.</exception>
+        /// <exception cref="ArgumentException"><paramref name="completedTask"/> is not completed.</exception>
+        /// <exception cref="InvalidOperationException">
+        /// The underlying <see cref="Task{TResult}"/> is already in one of the three final states:
+        /// <see cref="TaskStatus.RanToCompletion"/>, <see cref="TaskStatus.Faulted"/>, or <see cref="TaskStatus.Canceled"/>.
+        /// </exception>
+        /// <remarks>
+        /// This operation will return false if the <see cref="Task{TResult}"/> is already in one of the three final states:
+        /// <see cref="TaskStatus.RanToCompletion"/>, <see cref="TaskStatus.Faulted"/>, or <see cref="TaskStatus.Canceled"/>.
+        /// </remarks>
+        public void SetFromTask(Task completedTask)
+        {
+            if (!TrySetFromTask(completedTask))
+            {
+                ThrowHelper.ThrowInvalidOperationException(ExceptionResource.TaskT_TransitionToFinal_AlreadyCompleted);
+            }
+        }
+
+        /// <summary>
+        /// Attempts to transition the underlying <see cref="Task{TResult}"/> into the same completion state as the specified <paramref name="completedTask"/>.
+        /// </summary>
+        /// <param name="completedTask">The completed task whose completion status (including exception or cancellation information) should be copied to the underlying task.</param>
+        /// <returns><see langword="true"/> if the operation was successful; otherwise, <see langword="false"/>.</returns>
+        /// <exception cref="ArgumentNullException"><paramref name="completedTask"/> is <see langword="null"/>.</exception>
+        /// <exception cref="ArgumentException"><paramref name="completedTask"/> is not completed.</exception>
+        /// <remarks>
+        /// This operation will return false if the <see cref="Task{TResult}"/> is already in one of the three final states:
+        /// <see cref="TaskStatus.RanToCompletion"/>, <see cref="TaskStatus.Faulted"/>, or <see cref="TaskStatus.Canceled"/>.
+        /// </remarks>
+        public bool TrySetFromTask(Task completedTask)
+        {
+            ArgumentNullException.ThrowIfNull(completedTask);
+            if (!completedTask.IsCompleted)
+            {
+                throw new ArgumentException(SR.Task_MustBeCompleted, nameof(completedTask));
+            }
+
+            // Try to transition to the appropriate final state based on the state of completedTask.
+            bool result = false;
+            switch (completedTask.Status)
+            {
+                case TaskStatus.RanToCompletion:
+                    result = _task.TrySetResult();
+                    break;
+
+                case TaskStatus.Canceled:
+                    result = _task.TrySetCanceled(completedTask.CancellationToken, completedTask.GetCancellationExceptionDispatchInfo());
+                    break;
+
+                case TaskStatus.Faulted:
+                    result = _task.TrySetException(completedTask.GetExceptionDispatchInfos());
+                    break;
+            }
+
+            // If we successfully transitioned to a final state, we're done. If we didn't, it's possible a concurrent operation
+            // is still in the process of completing the task, and callers of this method expect the task to already be fully
+            // completed when this method returns. As such, we spin until the task is completed, and then return whether this
+            // call successfully did the transition.
+            if (!result && !_task.IsCompleted)
+            {
+                _task.SpinUntilCompleted();
+            }
+
+            return result;
+        }
     }
 }

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

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Generic;
+using System.Runtime.ExceptionServices;
 
 namespace System.Threading.Tasks
 {
@@ -286,5 +287,75 @@ public bool TrySetCanceled(CancellationToken cancellationToken)
 
             return rval;
         }
+
+        /// <summary>
+        /// Transition the underlying <see cref="Task{TResult}"/> into the same completion state as the specified <paramref name="completedTask"/>.
+        /// </summary>
+        /// <param name="completedTask">The completed task whose completion status (including result, exception, or cancellation information) should be copied to the underlying task.</param>
+        /// <exception cref="ArgumentNullException"><paramref name="completedTask"/> is <see langword="null"/>.</exception>
+        /// <exception cref="ArgumentException"><paramref name="completedTask"/> is not completed.</exception>
+        /// <exception cref="InvalidOperationException">
+        /// The underlying <see cref="Task{TResult}"/> is already in one of the three final states:
+        /// <see cref="TaskStatus.RanToCompletion"/>, <see cref="TaskStatus.Faulted"/>, or <see cref="TaskStatus.Canceled"/>.
+        /// </exception>
+        /// <remarks>
+        /// This operation will return false if the <see cref="Task{TResult}"/> is already in one of the three final states:
+        /// <see cref="TaskStatus.RanToCompletion"/>, <see cref="TaskStatus.Faulted"/>, or <see cref="TaskStatus.Canceled"/>.
+        /// </remarks>
+        public void SetFromTask(Task<TResult> completedTask)
+        {
+            if (!TrySetFromTask(completedTask))
+            {
+                ThrowHelper.ThrowInvalidOperationException(ExceptionResource.TaskT_TransitionToFinal_AlreadyCompleted);
+            }
+        }
+
+        /// <summary>
+        /// Attempts to transition the underlying <see cref="Task{TResult}"/> into the same completion state as the specified <paramref name="completedTask"/>.
+        /// </summary>
+        /// <param name="completedTask">The completed task whose completion status (including result, exception, or cancellation information) should be copied to the underlying task.</param>
+        /// <returns><see langword="true"/> if the operation was successful; otherwise, <see langword="false"/>.</returns>
+        /// <exception cref="ArgumentNullException"><paramref name="completedTask"/> is <see langword="null"/>.</exception>
+        /// <exception cref="ArgumentException"><paramref name="completedTask"/> is not completed.</exception>
+        /// <remarks>
+        /// This operation will return false if the <see cref="Task{TResult}"/> is already in one of the three final states:
+        /// <see cref="TaskStatus.RanToCompletion"/>, <see cref="TaskStatus.Faulted"/>, or <see cref="TaskStatus.Canceled"/>.
+        /// </remarks>
+        public bool TrySetFromTask(Task<TResult> completedTask)
+        {
+            ArgumentNullException.ThrowIfNull(completedTask);
+            if (!completedTask.IsCompleted)
+            {
+                throw new ArgumentException(SR.Task_MustBeCompleted, nameof(completedTask));
+            }
+
+            // Try to transition to the appropriate final state based on the state of completedTask.
+            bool result = false;
+            switch (completedTask.Status)
+            {
+                case TaskStatus.RanToCompletion:
+                    result = _task.TrySetResult(completedTask.Result);
+                    break;
+
+                case TaskStatus.Canceled:
+                    result = _task.TrySetCanceled(completedTask.CancellationToken, completedTask.GetCancellationExceptionDispatchInfo());
+                    break;
+
+                case TaskStatus.Faulted:
+                    result = _task.TrySetException(completedTask.GetExceptionDispatchInfos());
+                    break;
+            }
+
+            // If we successfully transitioned to a final state, we're done. If we didn't, it's possible a concurrent operation
+            // is still in the process of completing the task, and callers of this method expect the task to already be fully
+            // completed when this method returns. As such, we spin until the task is completed, and then return whether this
+            // call successfully did the transition.
+            if (!result && !_task.IsCompleted)
+            {
+                _task.SpinUntilCompleted();
+            }
+
+            return result;
+        }
     }
 }

src/libraries/System.Runtime/ref/System.Runtime.cs

@@ -15352,13 +15352,15 @@ public TaskCompletionSource(System.Threading.Tasks.TaskCreationOptions creationO
         public System.Threading.Tasks.Task Task { get { throw null; } }
         public void SetCanceled() { }
         public void SetCanceled(System.Threading.CancellationToken cancellationToken) { }
+        public void SetFromTask(System.Threading.Tasks.Task completedTask) { throw null; }
         public void SetException(System.Collections.Generic.IEnumerable<System.Exception> exceptions) { }
         public void SetException(System.Exception exception) { }
         public void SetResult() { }
         public bool TrySetCanceled() { throw null; }
         public bool TrySetCanceled(System.Threading.CancellationToken cancellationToken) { throw null; }
         public bool TrySetException(System.Collections.Generic.IEnumerable<System.Exception> exceptions) { throw null; }
         public bool TrySetException(System.Exception exception) { throw null; }
+        public bool TrySetFromTask(System.Threading.Tasks.Task completedTask) { throw null; }
         public bool TrySetResult() { throw null; }
     }
     public partial class TaskCompletionSource<TResult>
@@ -15370,11 +15372,13 @@ public TaskCompletionSource(System.Threading.Tasks.TaskCreationOptions creationO
         public System.Threading.Tasks.Task<TResult> Task { get { throw null; } }
         public void SetCanceled() { }
         public void SetCanceled(System.Threading.CancellationToken cancellationToken) { }
+        public void SetFromTask(System.Threading.Tasks.Task<TResult> completedTask) { throw null; }
         public void SetException(System.Collections.Generic.IEnumerable<System.Exception> exceptions) { }
         public void SetException(System.Exception exception) { }
         public void SetResult(TResult result) { }
         public bool TrySetCanceled() { throw null; }
         public bool TrySetCanceled(System.Threading.CancellationToken cancellationToken) { throw null; }
+        public bool TrySetFromTask(System.Threading.Tasks.Task<TResult> completedTask) { throw null; }
         public bool TrySetException(System.Collections.Generic.IEnumerable<System.Exception> exceptions) { throw null; }
         public bool TrySetException(System.Exception exception) { throw null; }
         public bool TrySetResult(TResult result) { throw null; }

src/libraries/System.Runtime/tests/System.Threading.Tasks.Tests/Task/TaskCompletionSourceTResultTests.cs

@@ -202,5 +202,114 @@ private static void AssertCompletedTcsFailsToCompleteAgain<T>(TaskCompletionSour
             Assert.False(tcs.TrySetCanceled());
             Assert.False(tcs.TrySetCanceled(default));
         }
+
+        [Fact]
+        public void SetFromTask_InvalidArgument_Throws()
+        {
+            TaskCompletionSource<object> tcs = new();
+            AssertExtensions.Throws<ArgumentNullException>("completedTask", () => tcs.SetFromTask(null));
+            AssertExtensions.Throws<ArgumentException>("completedTask", () => tcs.SetFromTask(new TaskCompletionSource<object>().Task));
+            Assert.False(tcs.Task.IsCompleted);
+
+            tcs.SetResult(null);
+            Assert.True(tcs.Task.IsCompletedSuccessfully);
+
+            AssertExtensions.Throws<ArgumentNullException>("completedTask", () => tcs.SetFromTask(null));
+            AssertExtensions.Throws<ArgumentException>("completedTask", () => tcs.SetFromTask(new TaskCompletionSource<object>().Task));
+            Assert.True(tcs.Task.IsCompletedSuccessfully);
+        }
+
+        [Fact]
+        public void SetFromTask_AlreadyCompleted_ReturnsFalseOrThrows()
+        {
+            object result = new();
+            TaskCompletionSource<object> tcs = new();
+            tcs.SetResult(result);
+
+            Assert.False(tcs.TrySetFromTask(Task.FromResult(new object())));
+            Assert.False(tcs.TrySetFromTask(Task.FromException<object>(new Exception())));
+            Assert.False(tcs.TrySetFromTask(Task.FromCanceled<object>(new CancellationToken(canceled: true))));
+
+            Assert.Throws<InvalidOperationException>(() => tcs.SetFromTask(Task.FromResult(new object())));
+            Assert.Throws<InvalidOperationException>(() => tcs.SetFromTask(Task.FromException<object>(new Exception())));
+            Assert.Throws<InvalidOperationException>(() => tcs.SetFromTask(Task.FromCanceled<object>(new CancellationToken(canceled: true))));
+
+            Assert.True(tcs.Task.IsCompletedSuccessfully);
+            Assert.Same(result, tcs.Task.Result);
+        }
+
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public void SetFromTask_CompletedSuccessfully(bool tryMethod)
+        {
+            TaskCompletionSource<object> tcs = new();
+            Task<object> source = Task.FromResult(new object());
+
+            if (tryMethod)
+            {
+                Assert.True(tcs.TrySetFromTask(source));
+            }
+            else
+            {
+                tcs.SetFromTask(source);
+            }
+
+            Assert.Same(source.Result, tcs.Task.Result);
+        }
+
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public void SetFromTask_Faulted(bool tryMethod)
+        {
+            TaskCompletionSource<object> tcs = new();
+
+            var source = new TaskCompletionSource<object>();
+            source.SetException([new FormatException(), new DivideByZeroException()]);
+
+            if (tryMethod)
+            {
+                Assert.True(tcs.TrySetFromTask(source.Task));
+            }
+            else
+            {
+                tcs.SetFromTask(source.Task);
+            }
+
+            Assert.True(tcs.Task.IsFaulted);
+            Assert.True(tcs.Task.Exception.InnerExceptions.Count == 2);
+        }
+
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public void SetFromTask_Canceled(bool tryMethod)
+        {
+            TaskCompletionSource<object> tcs = new();
+
+            var cts = new CancellationTokenSource();
+            cts.Cancel();
+            Task<object> source = Task.FromCanceled<object>(cts.Token);
+
+            if (tryMethod)
+            {
+                Assert.True(tcs.TrySetFromTask(source));
+            }
+            else
+            {
+                tcs.SetFromTask(source);
+            }
+
+            Assert.True(tcs.Task.IsCanceled);
+            try
+            {
+                tcs.Task.GetAwaiter().GetResult();
+            }
+            catch (OperationCanceledException oce)
+            {
+                Assert.Equal(cts.Token, oce.CancellationToken);
+            }
+        }
     }
 }