Avoid F# built-in task CEs

Author: gusty

src/FSharpPlus/Extensions/Task.fs

@@ -60,10 +60,15 @@ module Task =
     let map (mapper: 'T -> 'U) (source: Task<'T>) : Task<'U> =
         let source = nullArgCheck (nameof source) source
 
-        backgroundTask {
-            let! r = source
-            return mapper r
-        }
+        if source.IsCompleted then
+            match source with
+            | Succeeded r -> try result (mapper r) with e -> Task.FromException<_> e
+            | Faulted exn -> FromExceptions exn
+            | Canceled    -> canceled
+        else
+            let tcs = TaskCompletionSource<'U> TaskCreationOptions.RunContinuationsAsynchronously
+            source |> continueTask tcs (fun r -> try tcs.SetResult (mapper r) with e -> tcs.SetException e)
+            tcs.Task
 
     /// <summary>Creates a task workflow from two workflows 'task1' and 'task2', mapping its results with 'mapper'.</summary>
     /// <remarks>Workflows are run in sequence.</remarks>
@@ -293,19 +298,13 @@ module Task =
     let join (source: Task<Task<'T>>) : Task<'T> =
         let source = nullArgCheck (nameof source) source
 
-        backgroundTask {
-            let! inner = source
-            return! inner
-        }
+        source.Unwrap()
 
     /// <summary>Creates a task workflow from 'source' workflow, mapping and flattening its result with 'f'.</summary>
     let bind (f: 'T -> Task<'U>) (source: Task<'T>) : Task<'U> =
         let source = nullArgCheck (nameof source) source
 
-        backgroundTask {
-            let! r = source
-            return! f r
-        }
+        source |> Unchecked.nonNull |> map f |> join
 
     /// <summary>Creates a task that ignores the result of the source task.</summary>
     /// <param name="source">The source Task.</param>
@@ -327,14 +326,37 @@ module Task =
             tcs.Task
 
     [<ObsoleteAttribute("Swap parameters")>]
-    let tryWith (body: unit -> Task<'T>) (compensation: exn -> Task<'T>) : Task<'T> = backgroundTask {
-        try return! body ()
-        with e -> return! compensation e }
+    let rec tryWith (body: unit -> Task<'T>) (compensation: exn -> Task<'T>) : Task<'T> =
+        let unwrapException (agg: AggregateException) =
+            if agg.InnerExceptions.Count = 1 then agg.InnerExceptions.[0]
+            else agg :> Exception
+        try Ok (body ()) with e -> Error e
+        |> function
+        | Ok task ->
+            if task.IsCompleted then
+                match task with
+                | Succeeded _ -> task
+                | Faulted aex -> compensation (unwrapException aex)
+                | Canceled    -> canceled
+            else
+                task.ContinueWith((fun (x: Task<'T>) -> tryWith (fun () -> x) compensation)).Unwrap ()
+        | Error (:? AggregateException as exn) -> compensation (unwrapException exn)
+        | Error exn                            -> compensation exn
 
     [<ObsoleteAttribute("Swap parameters")>]
-    let tryFinally (body: unit -> Task<'T>) (compensation : unit -> unit) : Task<'T> = backgroundTask {
-        try return! body ()
-        finally compensation () }
+    let tryFinally (body: unit -> Task<'T>) (compensation : unit -> unit) : Task<'T> =
+        let mutable ran = false
+        let compensation () =
+            if not ran then
+                ran <- true
+                compensation ()
+        try
+            let task = body ()
+            if task.IsCompleted then compensation (); task
+            else task.ContinueWith(fun (_: Task<'T>) -> compensation (); task).Unwrap ()
+        with _ ->
+            compensation ()
+            reraise ()
 
     /// Used to de-sugar use .. blocks in Computation Expressions.
     let using (disp: 'T when 'T :> IDisposable) (body: 'T -> Task<'U>) =

src/FSharpPlus/Extensions/ValueTask.fs

@@ -63,10 +63,15 @@ module ValueTask =
     /// <param name="source">The source ValueTask workflow.</param>
     /// <returns>The resulting ValueTask workflow.</returns>
     let map (mapper: 'T -> 'U) (source: ValueTask<'T>) : ValueTask<'U> =
-        backgroundTask {
-            let! r = source
-            return mapper r
-        } |> ValueTask<'U>
+        if source.IsCompleted then
+            match source with
+            | Succeeded r -> try result (mapper r) with e -> ValueTask.FromException<_> e
+            | Faulted exn -> FromExceptions exn
+            | Canceled    -> canceled
+        else
+            let tcs = TaskCompletionSource<'U> TaskCreationOptions.RunContinuationsAsynchronously
+            source |> continueTask tcs (fun r -> try tcs.SetResult (mapper r) with e -> tcs.SetException e)
+            tcs.Task |> ValueTask<'U>
 
 
     /// <summary>Creates a ValueTask workflow from two workflows 'x' and 'y', mapping its results with 'f'.</summary>
@@ -271,17 +276,11 @@ module ValueTask =
 
     /// Flattens two nested ValueTask into one.
     let join (source: ValueTask<ValueTask<'T>>) : ValueTask<'T> =
-        backgroundTask {
-            let! inner = source
-            return! inner
-        } |> ValueTask<'T>    
+        (source |> map (fun x -> x.AsTask())).AsTask().Unwrap () |> ValueTask<'T>
 
     /// <summary>Creates a ValueTask workflow from 'source' workflow, mapping and flattening its result with 'f'.</summary>
     let bind (f: 'T -> ValueTask<'U>) (source: ValueTask<'T>) : ValueTask<'U> =
-        backgroundTask {
-            let! r = source
-            return! f r
-        } |> ValueTask<'U>
+        source |> map f |> join
 
     /// <summary>Creates a ValueTask that ignores the result of the source ValueTask.</summary>
     /// <param name="source">The source ValueTask.</param>
@@ -302,17 +301,38 @@ module ValueTask =
 
     /// Used to de-sugar try .. with .. blocks in Computation Expressions.
     let inline tryWith ([<InlineIfLambda>]compensation: exn -> ValueTask<'T>) ([<InlineIfLambda>]body: unit -> ValueTask<'T>) : ValueTask<'T> =
-        backgroundTask {
-            try return! body ()
-            with e -> return! compensation e
-        } |> ValueTask<'T>
+        let unwrapException (agg: AggregateException) =
+            if agg.InnerExceptions.Count = 1 then agg.InnerExceptions.[0]
+            else agg :> Exception
+        try Ok (body ()) with e -> Error e
+        |> function
+        | Ok task ->
+            if task.IsCompleted then
+                match task with
+                | Succeeded _ -> task
+                | Faulted aex -> compensation (unwrapException aex)
+                | Canceled    -> canceled
+            else
+                task.AsTask().ContinueWith((fun (x: Task<'T>) -> Task.tryWith (compensation >> fun x -> x.AsTask()) (fun () -> x))).Unwrap () |> ValueTask<'T>
+        | Error (:? AggregateException as exn) -> compensation (unwrapException exn)
+        | Error exn                            -> compensation exn
+        
 
     /// Used to de-sugar try .. finally .. blocks in Computation Expressions.
     let inline tryFinally ([<InlineIfLambda>]compensation : unit -> unit) ([<InlineIfLambda>]body: unit -> ValueTask<'T>) : ValueTask<'T> =
-        backgroundTask {
-            try return! body ()
-            finally compensation ()
-        } |> ValueTask<'T>
+        let mutable ran = false
+        let compensation () =
+            if not ran then
+                ran <- true
+                compensation ()
+        try
+            let task = body().AsTask ()
+            if task.IsCompleted then compensation (); task
+            else task.ContinueWith(fun (_: Task<'T>) -> compensation (); task).Unwrap ()
+        with _ ->
+            compensation ()
+            reraise ()
+        |> ValueTask<'T>
 
     /// Used to de-sugar use .. blocks in Computation Expressions.
     let inline using (disp: 'T when 'T :> IDisposable) ([<InlineIfLambda>]body: 'T -> ValueTask<'U>) =