Skip to content
This repository was archived by the owner on Jan 23, 2023. It is now read-only.

Commit 2616871

Browse files
stephentoubstephentoub
committed
Reimplement TaskExtensions.Unwrap
The currently checked in implementation of TaskExtensions.Unwrap isn't what's in .NET 4.5/6 desktop. Rather, it's a slight improvement over what was shipped in .NET 4, before a bunch of optimizations were made to it for .NET 4.5. In .NET 4.5, we effectively built the implementation into Task itself, taking advantage of Task internals to minimize allocations and other overheads; the Unwrap extension methods then had internals access and could delegate to that implementation in mscorlib. For the separate implementation now in System.Threading.Tasks.dll in .NET Core, we don't have that internals access, and need to implement Unwrap on top of Task's public surface area. Even so, we can do much better than the checked in implementation, and this commit overhauls it. The current implementation incurs 17 allocations (!) per Unwrap, whereas mscorlib's implementation incurs only 1 (along with a commensurate reduction in the number of bytes allocated). This implementation gets the number down from 17 to 4, again with a commensurate size reduction. It also gets the single-threaded speed of the implementation (i.e. not taking into account the GC impact of the extra allocations) to effectively the same as the implementation in mscorlib, whereas the old implementation is measurably worse. It also addresses an inconsistency between the implementations. In .NET 4.5/6, when the inner task has a CancellationToken that causes the task to be canceled, that CancellationToken is stored into the proxy task, a fact visible due to that token propagating out of any OperationCanceledException thrown when await'ing the proxy task. The current implementation doesn't have that, but rather the pre-.NET 4.5 behavior where the token wasn't marshaled. This commit brings it back.
1 parent 578e4ff commit 2616871

File tree

1 file changed

+119
-108
lines changed

1 file changed

+119
-108
lines changed

‎src/System.Threading.Tasks/src/System/Threading/Tasks/TaskExtensions.CoreCLR.cs

Lines changed: 119 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
33

44
using System.Diagnostics;
5-
using System.Diagnostics.Contracts;
65

76
namespace System.Threading.Tasks
87
{
@@ -28,56 +27,13 @@ public static class TaskExtensions
2827
/// <returns>A Task that represents the asynchronous operation of the provided Task{Task}.</returns>
2928
public static Task Unwrap(this Task<Task> task)
3029
{
31-
if (task == null) throw new ArgumentNullException("task");
30+
if (task == null)
31+
throw new ArgumentNullException("task");
3232

33-
bool result;
34-
35-
// tcs.Task serves as a proxy for task.Result.
36-
// AttachedToParent is the only legal option for TCS-style task.
37-
var tcs = new TaskCompletionSource<Task>(task.CreationOptions & TaskCreationOptions.AttachedToParent);
38-
39-
// Set up some actions to take when task has completed.
40-
task.ContinueWith(delegate
41-
{
42-
switch (task.Status)
43-
{
44-
// If task did not run to completion, then record the cancellation/fault information
45-
// to tcs.Task.
46-
case TaskStatus.Canceled:
47-
case TaskStatus.Faulted:
48-
result = tcs.TrySetFromTask(task);
49-
Debug.Assert(result, "Unwrap(Task<Task>): Expected TrySetFromTask #1 to succeed");
50-
break;
51-
52-
case TaskStatus.RanToCompletion:
53-
// task.Result == null ==> proxy should be canceled.
54-
if (task.Result == null) tcs.TrySetCanceled();
55-
56-
// When task.Result completes, take some action to set the completion state of tcs.Task.
57-
else
58-
{
59-
task.Result.ContinueWith(_ =>
60-
{
61-
// Copy completion/cancellation/exception info from task.Result to tcs.Task.
62-
result = tcs.TrySetFromTask(task.Result);
63-
Debug.Assert(result, "Unwrap(Task<Task>): Expected TrySetFromTask #2 to succeed");
64-
}, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously | TaskContinuationOptions.DenyChildAttach, TaskScheduler.Default)
65-
.ContinueWith(antecedent =>
66-
{
67-
// Clean up if ContinueWith() operation fails due to TSE
68-
tcs.TrySetException(antecedent.Exception);
69-
}, CancellationToken.None, TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.DenyChildAttach, TaskScheduler.Default);
70-
}
71-
break;
72-
}
73-
}, TaskContinuationOptions.ExecuteSynchronously | TaskContinuationOptions.DenyChildAttach, TaskScheduler.Default).ContinueWith(antecedent =>
74-
{
75-
// Clean up if ContinueWith() operation fails due to TSE
76-
tcs.TrySetException(antecedent.Exception);
77-
}, CancellationToken.None, TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.DenyChildAttach, TaskScheduler.Default);
78-
79-
// Return this immediately as a proxy. When task.Result completes, or task is faulted/canceled,
80-
// the completion information will be transfered to tcs.Task.
33+
// Create a new Task to serve as a proxy for the actual inner task. Attach it
34+
// to the parent if the original was attached to the parent.
35+
var tcs = new TaskCompletionSource<VoidResult>(task.CreationOptions & TaskCreationOptions.AttachedToParent);
36+
TransferAsynchronously(tcs, task);
8137
return tcs.Task;
8238
}
8339

@@ -97,86 +53,141 @@ public static Task Unwrap(this Task<Task> task)
9753
/// <returns>A Task{TResult} that represents the asynchronous operation of the provided Task{Task{TResult}}.</returns>
9854
public static Task<TResult> Unwrap<TResult>(this Task<Task<TResult>> task)
9955
{
100-
if (task == null) throw new ArgumentNullException("task");
101-
102-
bool result;
56+
if (task == null)
57+
throw new ArgumentNullException("task");
10358

104-
// tcs.Task serves as a proxy for task.Result.
105-
// AttachedToParent is the only legal option for TCS-style task.
59+
// Create a new Task to serve as a proxy for the actual inner task. Attach it
60+
// to the parent if the original was attached to the parent.
10661
var tcs = new TaskCompletionSource<TResult>(task.CreationOptions & TaskCreationOptions.AttachedToParent);
62+
TransferAsynchronously(tcs, task);
63+
return tcs.Task;
64+
}
10765

108-
// Set up some actions to take when task has completed.
109-
task.ContinueWith(delegate
66+
/// <summary>
67+
/// Transfer the results of the <paramref name="outer"/> task's inner task to the <paramref name="completionSource"/>.
68+
/// </summary>
69+
/// <param name="completionSource">The completion source to which results should be transfered.</param>
70+
/// <param name="outer">
71+
/// The outer task that when completed will yield an inner task whose results we want marshaled to the <paramref name="completionSource"/>.
72+
/// </param>
73+
private static void TransferAsynchronously<TResult, TInner>(TaskCompletionSource<TResult> completionSource, Task<TInner> outer) where TInner : Task
74+
{
75+
Action callback = null;
76+
77+
// Create a continuation delegate. For performance reasons, we reuse the same delegate/closure across multiple
78+
// continuations; by using .ConfigureAwait(false).GetAwaiter().UnsafeOnComplete(action), in most cases
79+
// this delegate can be stored directly into the Task's continuation field, eliminating the need for additional
80+
// allocations. Thus, this whole Unwrap operation generally results in four allocations: one for the TaskCompletionSource,
81+
// one for the returned task, one for the delegate, and one for the closure. Since the delegate is used
82+
// across multiple continuations, we use the callback variable as well to indicate which continuation we're in:
83+
// if the callback is non-null, then we're processing the continuation for the outer task and use the callback
84+
// object as the continuation off of the inner task; if the callback is null, then we're processing the
85+
// inner task.
86+
callback = delegate
11087
{
111-
switch (task.Status)
88+
Debug.Assert(outer.IsCompleted);
89+
if (callback != null)
11290
{
113-
// If task did not run to completion, then record the cancellation/fault information
114-
// to tcs.Task.
115-
case TaskStatus.Canceled:
116-
case TaskStatus.Faulted:
117-
result = tcs.TrySetFromTask(task);
118-
Debug.Assert(result, "Unwrap(Task<Task<T>>): Expected TrySetFromTask #1 to succeed");
119-
break;
120-
121-
case TaskStatus.RanToCompletion:
122-
// task.Result == null ==> proxy should be canceled.
123-
if (task.Result == null) tcs.TrySetCanceled();
124-
125-
// When task.Result completes, take some action to set the completion state of tcs.Task.
126-
else
127-
{
128-
task.Result.ContinueWith(_ =>
91+
// Process the outer task's completion
92+
93+
// Clear out the callback field to indicate that any future invocations should
94+
// be for processing the inner task, but store away a local copy in case we need
95+
// to use it as the continuation off of the outer task.
96+
Action innerCallback = callback;
97+
callback = null;
98+
99+
bool result = true;
100+
switch (outer.Status)
101+
{
102+
case TaskStatus.Canceled:
103+
case TaskStatus.Faulted:
104+
// The outer task has completed as canceled or faulted; transfer that
105+
// status to the completion source, and we're done.
106+
result = completionSource.TrySetFromTask(outer);
107+
break;
108+
case TaskStatus.RanToCompletion:
109+
Task inner = outer.Result;
110+
if (inner == null)
129111
{
130-
// Copy completion/cancellation/exception info from task.Result to tcs.Task.
131-
result = tcs.TrySetFromTask(task.Result);
132-
Debug.Assert(result, "Unwrap(Task<Task<T>>): Expected TrySetFromTask #2 to succeed");
133-
},
134-
CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously | TaskContinuationOptions.DenyChildAttach, TaskScheduler.Default)
135-
.ContinueWith(antecedent =>
136-
{
137-
// Clean up if ContinueWith() operation fails due to TSE
138-
tcs.TrySetException(antecedent.Exception);
139-
}, CancellationToken.None, TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.DenyChildAttach, TaskScheduler.Default);
140-
}
141-
142-
break;
112+
// The outer task completed successfully, but with a null inner task;
113+
// cancel the completionSource, and we're done.
114+
result = completionSource.TrySetCanceled();
115+
}
116+
else if (inner.IsCompleted)
117+
{
118+
// The inner task also completed! Transfer the results, and we're done.
119+
result = completionSource.TrySetFromTask(inner);
120+
}
121+
else
122+
{
123+
// Run this delegate again once the inner task has completed.
124+
inner.ConfigureAwait(false).GetAwaiter().UnsafeOnCompleted(innerCallback);
125+
}
126+
break;
127+
}
128+
Debug.Assert(result);
143129
}
144-
}, TaskContinuationOptions.ExecuteSynchronously | TaskContinuationOptions.DenyChildAttach, TaskScheduler.Current).ContinueWith(antecedent =>
145-
{
146-
// Clean up if ContinueWith() operation fails due to TSE
147-
tcs.TrySetException(antecedent.Exception);
148-
}, CancellationToken.None, TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.DenyChildAttach, TaskScheduler.Default); ;
130+
else
131+
{
132+
// Process the inner task's completion. All we need to do is transfer its results
133+
// to the completion source.
134+
Debug.Assert(outer.Status == TaskStatus.RanToCompletion);
135+
Debug.Assert(outer.Result.IsCompleted);
136+
completionSource.TrySetFromTask(outer.Result);
137+
}
138+
};
149139

150-
// Return this immediately as a proxy. When task.Result completes, or task is faulted/canceled,
151-
// the completion information will be transfered to tcs.Task.
152-
return tcs.Task;
140+
// Kick things off by hooking up the callback as the task's continuation
141+
outer.ConfigureAwait(false).GetAwaiter().UnsafeOnCompleted(callback);
153142
}
154143

155-
// Transfer the completion status from "source" to "me".
156-
private static bool TrySetFromTask<TResult>(this TaskCompletionSource<TResult> me, Task source)
144+
/// <summary>Copies that ending state information from <paramref name="task"/> to <paramref name="completionSource"/>.</summary>
145+
private static bool TrySetFromTask<TResult>(this TaskCompletionSource<TResult> completionSource, Task task)
157146
{
158-
Debug.Assert(source.IsCompleted, "TrySetFromTask: Expected source to have completed.");
159-
bool rval = false;
147+
Debug.Assert(task.IsCompleted);
160148

161-
switch(source.Status)
149+
bool result = false;
150+
switch(task.Status)
162151
{
163152
case TaskStatus.Canceled:
164-
rval = me.TrySetCanceled();
153+
result = completionSource.TrySetCanceled(ExtractCancellationToken(task));
165154
break;
166155

167156
case TaskStatus.Faulted:
168-
rval = me.TrySetException(source.Exception.InnerExceptions);
157+
result = completionSource.TrySetException(task.Exception.InnerExceptions);
169158
break;
170159

171160
case TaskStatus.RanToCompletion:
172-
if(source is Task<TResult>)
173-
rval = me.TrySetResult( ((Task<TResult>)source).Result);
174-
else
175-
rval = me.TrySetResult(default(TResult));
161+
Task<TResult> resultTask = task as Task<TResult>;
162+
result = resultTask != null ?
163+
completionSource.TrySetResult(resultTask.Result) :
164+
completionSource.TrySetResult(default(TResult));
176165
break;
177166
}
167+
return result;
168+
}
178169

179-
return rval;
170+
/// <summary>Gets the CancellationToken associated with a canceled task.</summary>
171+
private static CancellationToken ExtractCancellationToken(Task task)
172+
{
173+
// With the public Task APIs as of .NET 4.6, the only way to extract a CancellationToken
174+
// that was associated with a Task is by await'ing it, catching the resulting
175+
// OperationCanceledException, and getting the token from the OCE.
176+
Debug.Assert(task.IsCanceled);
177+
try
178+
{
179+
task.GetAwaiter().GetResult();
180+
}
181+
catch (OperationCanceledException oce)
182+
{
183+
CancellationToken ct = oce.CancellationToken;
184+
if (ct.IsCancellationRequested)
185+
return ct;
186+
}
187+
return new CancellationToken(true);
180188
}
189+
190+
/// <summary>Dummy type to use as a void TResult.</summary>
191+
private struct VoidResult { }
181192
}
182-
}
193+
}

0 commit comments

Comments
 (0)