Proper Task.zip implementation

Author: gusty

src/FSharpPlus/Extensions/Task.fs

@@ -7,6 +7,7 @@ namespace FSharpPlus
 module Task =
 
     open System
+    open System.Threading
     open System.Threading.Tasks
 
     let private (|Canceled|Faulted|Completed|) (t: Task<'a>) =
@@ -135,26 +136,91 @@ module Task =
             tcs.Task
 
     /// <summary>Creates a task workflow from two workflows 'x' and 'y', mapping its results with 'f'.</summary>
-    /// <remarks>Similar to lift2 but although workflows are started in sequence they might end independently in different order.</remarks>
-    /// <param name="f">The mapping function.</param>
-    /// <param name="x">First task workflow.</param>
-    /// <param name="y">Second task workflow.</param>
-    let map2 f x y = task {
-        let! x' = x
-        let! y' = y
-        return f x' y' }
+    /// <remarks>Similar to lift2 but although workflows are started in sequence they might end independently in different order
+    /// and all errors are collected.
+    /// </remarks>
+    /// <param name="mapper">The mapping function.</param>
+    /// <param name="task1">First task workflow.</param>
+    /// <param name="task2">Second task workflow.</param>
+    let map2 mapper (task1: Task<'T1>) (task2: Task<'T2>) : Task<'U> =
+        if task1.Status = TaskStatus.RanToCompletion && task2.Status = TaskStatus.RanToCompletion then
+            try Task.FromResult (mapper task1.Result task2.Result)
+            with e ->
+                let tcs = TaskCompletionSource<_> ()
+                tcs.SetException e
+                tcs.Task
+        else
+        let tcs = TaskCompletionSource<_> ()
+        let r1 = ref Unchecked.defaultof<_>
+        let r2 = ref Unchecked.defaultof<_>
+        let mutable cancelled = false
+        let failures = [|IReadOnlyCollection.empty; IReadOnlyCollection.empty|]
+        let pending = ref 2
+
+        let trySet () =
+            if Interlocked.Decrement pending = 0 then
+                let noFailures = Array.forall IReadOnlyCollection.isEmpty failures
+                if noFailures && not cancelled then
+                    try tcs.SetResult (mapper r1.Value r2.Value)
+                    with e -> tcs.SetException e
+                elif noFailures then tcs.SetCanceled ()
+                else tcs.SetException (failures |> Seq.map AggregateException |> Seq.reduce Exception.add).InnerExceptions
+
+        let k (v: ref<'k>) i t =
+            match t with
+            | Canceled    -> cancelled <- true
+            | Faulted e   -> failures[i] <- e.InnerExceptions
+            | Completed r -> v.Value <- r
+            trySet ()
+
+        task1.ContinueWith (k r1 0) |> ignore
+        task2.ContinueWith (k r2 1) |> ignore
+        tcs.Task
 
     /// <summary>Creates a task workflow from three workflows 'x', 'y' and z, mapping its results with 'f'.</summary>
-    /// <remarks>Similar to lift3 but although workflows are started in sequence they might end independently in different order.</remarks>
-    /// <param name="f">The mapping function.</param>
-    /// <param name="x">First task workflow.</param>
-    /// <param name="y">Second task workflow.</param>
-    /// <param name="z">Third task workflow.</param>
-    let map3 f x y z = task {
-        let! x' = x
-        let! y' = y
-        let! z' = z
-        return f x' y' z' }
+    /// <remarks>Similar to lift3 but although workflows are started in sequence they might end independently in different order
+    /// and all errors are collected.
+    /// </remarks>
+    /// <param name="mapper">The mapping function.</param>
+    /// <param name="task1">First task workflow.</param>
+    /// <param name="task2">Second task workflow.</param>
+    /// <param name="task3">Third task workflow.</param>
+    let map3 mapper (task1: Task<'T1>) (task2: Task<'T2>) (task3: Task<'T3>) : Task<'U> =
+        if task1.Status = TaskStatus.RanToCompletion && task2.Status = TaskStatus.RanToCompletion && task3.Status = TaskStatus.RanToCompletion then
+            try Task.FromResult (mapper task1.Result task2.Result task3.Result)
+            with e ->
+                let tcs = TaskCompletionSource<_> ()
+                tcs.SetException e
+                tcs.Task
+        else
+        let tcs = TaskCompletionSource<_> ()
+        let r1 = ref Unchecked.defaultof<_>
+        let r2 = ref Unchecked.defaultof<_>
+        let r3 = ref Unchecked.defaultof<_>
+        let mutable cancelled = false
+        let failures = [|IReadOnlyCollection.empty<exn>; IReadOnlyCollection.empty; IReadOnlyCollection.empty|]
+        let pending = ref 3
+
+        let trySet () =
+            if Interlocked.Decrement pending = 0 then
+                let noFailures = Array.forall isNull failures
+                if noFailures && not cancelled then
+                    try tcs.SetResult (mapper r1.Value r2.Value r3.Value)
+                    with e -> tcs.SetException e
+                elif noFailures then tcs.SetCanceled ()
+                else tcs.SetException (failures |> Seq.concat |> Seq.fold Exception.add (AggregateException ())).InnerExceptions
+
+        let k (v: ref<'k>) i t =
+            match t with
+            | Canceled    -> cancelled <- true
+            | Faulted e   -> failures[i] <- e.InnerExceptions
+            | Completed r -> v.Value <- r
+            trySet ()
+
+        task1.ContinueWith (k r1 0) |> ignore
+        task2.ContinueWith (k r2 1) |> ignore
+        task3.ContinueWith (k r3 2) |> ignore
+        tcs.Task
 
     /// <summary>Creates a task workflow that is the result of applying the resulting function of a task workflow
     /// to the resulting value of another task workflow</summary>
@@ -242,11 +308,16 @@ module Task =
             tcs.Task
 
     /// <summary>Creates a task workflow from two workflows 'x' and 'y', tupling its results.</summary>
-    /// <remarks>Similar to zipSequentially but although workflows are started in sequence they might end independently in different order.</remarks>
-    let zip x y = task {
-        let! x' = x
-        let! y' = y
-        return x', y' }
+    /// <remarks>Similar to zipSequentially but although workflows are started in sequence they might end independently in different order
+    /// and all errors are collected.
+    /// </remarks>
+    let zip (task1: Task<'T1>) (task2: Task<'T2>) = map2 (fun x y -> x, y) task1 task2
+
+    /// <summary>Creates a task workflow from two workflows 'x', 'y' and 'z', tupling its results.</summary>
+    /// <remarks>Similar to zipSequentially but although workflows are started in sequence they might end independently in different order
+    /// and all errors are collected.
+    /// </remarks>
+    let zip3 (task1: Task<'T1>) (task2: Task<'T2>) (task3: Task<'T3>) = map3 (fun x y z -> x, y, z) task1 task2 task3
 
     /// Flattens two nested tasks into one.
     let join (source: Task<Task<'T>>) : Task<'T> = source.Unwrap()

tests/FSharpPlus.Tests/Task.fs

@@ -12,6 +12,7 @@ module Task =
     exception TestException of string
 
     module TaskTests =
+        open System.Threading
 
         let createTask isFailed delay value =
             if not isFailed && delay = 0 then Task.FromResult value
@@ -39,11 +40,17 @@ module Task =
             let b = Task.zipSequentially x1 x2
             require b.IsCompleted "Task.zipSequentially didn't short-circuit"
 
+            let b1 = Task.zip3 x1 x2 x3
+            require b1.IsCompleted "Task.zip3 didn't short-circuit"
+
             let c = Task.lift2 (+) x1 x2
             require c.IsCompleted "Task.lift2 didn't short-circuit"
 
             let d = Task.lift3 (fun x y z -> x + y + z) x1 x2 x3
-            require d.IsCompleted "Task.lift3 didn't short-circiut"
+            require d.IsCompleted "Task.lift3 didn't short-circuit"
+
+            let d2 = Task.map3 (fun x y z -> x + y + z) x1 x2 x3
+            require d2.IsCompleted "Task.map3 didn't short-circuit"
 
         [<Test>]
         let erroredTasks () =
@@ -151,6 +158,42 @@ module Task =
             r19.Exception.InnerExceptions |> areEquivalent [TestException "Ouch, can't create: 2"]
             let r20 = Task.lift3 (mapping3 false) (x1 ()) (x2 ()) (e3 ())
             r20.Exception.InnerExceptions |> areEquivalent [TestException "Ouch, can't create: 3"]
+
+        [<Test>]
+        let testTaskZip () =
+            let t1 = createTask true 0 1
+            let t2 = createTask true 0 2
+            let t3 = createTask true 0 3
+
+            let c = new CancellationToken true
+            let t4 = Task.FromCanceled<int> c
+
+            let t5 = createTask false 0 5
+            let t6 = createTask false 0 6
+
+            let t12 =    Task.WhenAll [t1; t2]
+            let t12t12 = Task.WhenAll [t12; t12]
+            let t33    = Task.WhenAll [t3; t3]
+
+            t12.Wait ()
+            t12t12.Wait ()
+            t33.Wait ()
+
+            let t12123 = Task.zip3 t12t12 t33 t4
+            let ac1 =
+                try
+                    t12123.Exception.InnerExceptions |> Seq.map (fun x -> int (Char.GetNumericValue x.Message.[35]))
+                with e ->
+                    failwithf "Failure in testTaskZip. Task status is %A . Exception is %A" t12123.Status e
+
+            CollectionAssert.AreEquivalent ([1; 2; 1; 2; 3], ac1, "Task.zip(3) should add only non already existing exceptions.")
+
+            let t13 = Task.zip3 (Task.zip t1 t3) t4 (Task.zip t5 t6)
+            Assert.AreEqual (true, t13.IsFaulted, "Task.zip(3) between a value, an exception and a cancellation -> exception wins.")
+            let ac2 = t13.Exception.InnerExceptions |> Seq.map (fun x -> int (Char.GetNumericValue x.Message.[35]))
+            CollectionAssert.AreEquivalent ([1; 3], ac2, "Task.zip between 2 exceptions => both exceptions returned, even after combining with cancellation and values.")
+
+
 
     module TaskBuilderTests =