Fix race condition in TaskHelpersExtensions.Catch

bradwilson · bradwilson · commit 507866de048f · 2012-04-04T12:27:31.000-07:00
diff --git a/src/Common/TaskHelpersExtensions.cs b/src/Common/TaskHelpersExtensions.cs
@@ -21,6 +21,7 @@ internal static class TaskHelpersExtensions
         /// </summary>
         internal static Task Catch(this Task task, Func<CatchInfo, CatchInfo.CatchResult> continuation, CancellationToken cancellationToken = default(CancellationToken))
         {
+            // Fast path for successful tasks, to prevent an extra TCS allocation
             if (task.Status == TaskStatus.RanToCompletion)
             {
                 return task;
@@ -40,10 +41,12 @@ internal static class TaskHelpersExtensions
         /// </summary>
         internal static Task<TResult> Catch<TResult>(this Task<TResult> task, Func<CatchInfo<TResult>, CatchInfo<TResult>.CatchResult> continuation, CancellationToken cancellationToken = default(CancellationToken))
         {
+            // Fast path for successful tasks, to prevent an extra TCS allocation
             if (task.Status == TaskStatus.RanToCompletion)
             {
                 return task;
             }
+
             return task.CatchImpl(() => continuation(new CatchInfo<TResult>(task)).Task, cancellationToken);
         }
 
@@ -55,8 +58,6 @@ private static Task<TResult> CatchImpl<TResult>(this Task task, Func<Task<TResul
             // Stay on the same thread if we can
             if (task.IsCompleted)
             {
-                Contract.Assert(task.IsFaulted || task.IsCanceled); // caller ensures 
-
                 if (task.IsFaulted)
                 {
                     try
@@ -80,8 +81,12 @@ private static Task<TResult> CatchImpl<TResult>(this Task task, Func<Task<TResul
                 {
                     return TaskHelpers.Canceled<TResult>();
                 }
-
-                Contract.Assert(task.Status != TaskStatus.RanToCompletion);                
+                if (task.Status == TaskStatus.RanToCompletion)
+                {
+                    TaskCompletionSource<TResult> tcs = new TaskCompletionSource<TResult>();
+                    tcs.TrySetFromTask(task);
+                    return tcs.Task;
+                }
             }
 
             // Split into a continuation method so that we don't create a closure unnecessarily
@@ -242,7 +247,7 @@ internal static Task Finally(this Task task, Action continuation)
         {
             // Stay on the same thread if we can
             if (task.IsCompleted)
-            {                
+            {
                 try
                 {
                     continuation();
@@ -251,7 +256,7 @@ internal static Task Finally(this Task task, Action continuation)
                 catch (Exception ex)
                 {
                     return TaskHelpers.FromError(ex);
-                }                
+                }
             }
 
             // Split into a continuation method so that we don't create a closure unnecessarily
@@ -267,7 +272,7 @@ internal static Task<TResult> Finally<TResult>(this Task<TResult> task, Action c
         {
             // Stay on the same thread if we can
             if (task.IsCompleted)
-            {               
+            {
                 try
                 {
                     continuation();
@@ -276,7 +281,7 @@ internal static Task<TResult> Finally<TResult>(this Task<TResult> task, Action c
                 catch (Exception ex)
                 {
                     return TaskHelpers.FromError<TResult>(ex);
-                }                
+                }
             }
 
             // Split into a continuation method so that we don't create a closure unnecessarily

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@ internal static class TaskHelpersExtensions`
`21`	`21`	`/// </summary>`
`22`	`22`	`internal static Task Catch(this Task task, Func<CatchInfo, CatchInfo.CatchResult> continuation, CancellationToken cancellationToken = default(CancellationToken))`
`23`	`23`	`{`
	`24`	`+ // Fast path for successful tasks, to prevent an extra TCS allocation`
`24`	`25`	`if (task.Status == TaskStatus.RanToCompletion)`
`25`	`26`	`{`
`26`	`27`	`return task;`
`@@ -40,10 +41,12 @@ internal static class TaskHelpersExtensions`
`40`	`41`	`/// </summary>`
`41`	`42`	`internal static Task<TResult> Catch<TResult>(this Task<TResult> task, Func<CatchInfo<TResult>, CatchInfo<TResult>.CatchResult> continuation, CancellationToken cancellationToken = default(CancellationToken))`
`42`	`43`	`{`
	`44`	`+ // Fast path for successful tasks, to prevent an extra TCS allocation`
`43`	`45`	`if (task.Status == TaskStatus.RanToCompletion)`
`44`	`46`	`{`
`45`	`47`	`return task;`
`46`	`48`	`}`
	`49`	`+`
`47`	`50`	`return task.CatchImpl(() => continuation(new CatchInfo<TResult>(task)).Task, cancellationToken);`
`48`	`51`	`}`
`49`	`52`
`@@ -55,8 +58,6 @@ private static Task<TResult> CatchImpl<TResult>(this Task task, Func<Task<TResul`
`55`	`58`	`// Stay on the same thread if we can`
`56`	`59`	`if (task.IsCompleted)`
`57`	`60`	`{`
`58`		`- Contract.Assert(task.IsFaulted \|\| task.IsCanceled); // caller ensures`
`59`		`-`
`60`	`61`	`if (task.IsFaulted)`
`61`	`62`	`{`
`62`	`63`	`try`
`@@ -80,8 +81,12 @@ private static Task<TResult> CatchImpl<TResult>(this Task task, Func<Task<TResul`
`80`	`81`	`{`
`81`	`82`	`return TaskHelpers.Canceled<TResult>();`
`82`	`83`	`}`
`83`		`-`
`84`		`- Contract.Assert(task.Status != TaskStatus.RanToCompletion);`
	`84`	`+ if (task.Status == TaskStatus.RanToCompletion)`
	`85`	`+ {`
	`86`	`+ TaskCompletionSource<TResult> tcs = new TaskCompletionSource<TResult>();`
	`87`	`+ tcs.TrySetFromTask(task);`
	`88`	`+ return tcs.Task;`
	`89`	`+ }`
`85`	`90`	`}`
`86`	`91`
`87`	`92`	`// Split into a continuation method so that we don't create a closure unnecessarily`
`@@ -242,7 +247,7 @@ internal static Task Finally(this Task task, Action continuation)`
`242`	`247`	`{`
`243`	`248`	`// Stay on the same thread if we can`
`244`	`249`	`if (task.IsCompleted)`
`245`		`- {`
	`250`	`+ {`
`246`	`251`	`try`
`247`	`252`	`{`
`248`	`253`	`continuation();`
`@@ -251,7 +256,7 @@ internal static Task Finally(this Task task, Action continuation)`
`251`	`256`	`catch (Exception ex)`
`252`	`257`	`{`
`253`	`258`	`return TaskHelpers.FromError(ex);`
`254`		`- }`
	`259`	`+ }`
`255`	`260`	`}`
`256`	`261`
`257`	`262`	`// Split into a continuation method so that we don't create a closure unnecessarily`
`@@ -267,7 +272,7 @@ internal static Task<TResult> Finally<TResult>(this Task<TResult> task, Action c`
`267`	`272`	`{`
`268`	`273`	`// Stay on the same thread if we can`
`269`	`274`	`if (task.IsCompleted)`
`270`		`- {`
	`275`	`+ {`
`271`	`276`	`try`
`272`	`277`	`{`
`273`	`278`	`continuation();`
`@@ -276,7 +281,7 @@ internal static Task<TResult> Finally<TResult>(this Task<TResult> task, Action c`
`276`	`281`	`catch (Exception ex)`
`277`	`282`	`{`
`278`	`283`	`return TaskHelpers.FromError<TResult>(ex);`
`279`		`- }`
	`284`	`+ }`
`280`	`285`	`}`
`281`	`286`
`282`	`287`	`// Split into a continuation method so that we don't create a closure unnecessarily`