Cancel task continuations when async contexts are cancelled

src/fsharp/FSharp.Core/async.fs

@@ -984,7 +984,15 @@ namespace Microsoft.FSharp.Control
                     else
                         ctxt.cont completedTask.Result) |> unfake
 
-            task.ContinueWith(Action<Task<'T>>(continuation)) |> ignore |> fake
+            let cancelContinuation (_: Task) : unit =
+                ctxt.trampolineHolder.ExecuteWithTrampoline (fun () ->
+                    ctxt.OnCancellation ()
+                ) |> unfake
+
+            task
+                .ContinueWith(Action<Task<'T>>(continuation), ctxt.token)
+                .ContinueWith(Action<Task>(cancelContinuation), TaskContinuationOptions.OnlyOnCanceled)
+            |> ignore |> fake
 
         [<DebuggerHidden>]
         let taskContinueWithUnit (task: Task) (ctxt: AsyncActivation<unit>) useCcontForTaskCancellation =
@@ -1003,7 +1011,15 @@ namespace Microsoft.FSharp.Control
                     else
                         ctxt.cont ()) |> unfake
 
-            task.ContinueWith(Action<Task>(continuation)) |> ignore |> fake
+            let cancelContinuation (_: Task) : unit =
+                ctxt.trampolineHolder.ExecuteWithTrampoline (fun () ->
+                    ctxt.OnCancellation ()
+                ) |> unfake
+
+            task
+                .ContinueWith(Action<Task>(continuation), ctxt.token)
+                .ContinueWith(Action<Task>(cancelContinuation), TaskContinuationOptions.OnlyOnCanceled)
+            |> ignore |> fake
 
         [<Sealed; AutoSerializable(false)>]
         type AsyncIAsyncResult<'T>(callback: System.AsyncCallback, state:obj) =

tests/FSharp.Core.UnitTests/FSharp.Core/Microsoft.FSharp.Control/AsyncType.fs

@@ -143,7 +143,7 @@ type AsyncType() =
         Assert.AreEqual(s, t.Result)
 
     [<Test>]
-    member this.StartAsTaskCancellation () =
+    member this.StartAsTaskCancelAsync () =
         let cts = new CancellationTokenSource()
         let tcs = TaskCompletionSource<unit>()
         let a = async {
@@ -155,12 +155,32 @@ type AsyncType() =
         use t : Task<unit> =
 #endif
             Async.StartAsTask(a, cancellationToken = cts.Token)
+
+        try
+            this.WaitASec t
+        with :? AggregateException as a ->
+            match a.InnerException with
+            | :? TaskCanceledException as t -> ()
+            | _ -> reraise()
+        Assert.IsTrue (t.IsCompleted, "Task is not completed")
+
+    [<Test>]
+    member this.StartAsTaskCancelTask () =
+        let tcs = TaskCompletionSource<unit>()
+        let a = async {
+            do! tcs.Task |> Async.AwaitTask }
+#if !NET46
+        let t : Task<unit> =
+#else
+        use t : Task<unit> =
+#endif
+            Async.StartAsTask(a)
 
         // Should not finish
         try
             let result = t.Wait(300)
             Assert.IsFalse (result)
-        with :? AggregateException -> Assert.Fail "Task should not finish, jet"
+        with :? AggregateException -> Assert.Fail "Task should not finish, yet"
 
         tcs.SetCanceled()
 
@@ -382,7 +402,7 @@ type AsyncType() =
         Async.RunSynchronously(a, 1000) |> Assert.IsTrue        
 
     [<Test>]
-    member this.AwaitTaskCancellation () =
+    member this.AwaitTaskTaskCancellation () =
         let test() = async {
             let tcs = new System.Threading.Tasks.TaskCompletionSource<unit>()
             tcs.SetCanceled()
@@ -392,8 +412,22 @@ type AsyncType() =
             with :? System.OperationCanceledException -> return true
         }
 
-        Async.RunSynchronously(test()) |> Assert.IsTrue   
-
+        Async.RunSynchronously(test()) |> Assert.IsTrue
+
+    [<Test>]
+    member this.AwaitTaskAsyncCancellation () =
+        let tcs = new System.Threading.Tasks.TaskCompletionSource<unit>()
+        let test = Async.AwaitTask tcs.Task
+
+        use cts = new CancellationTokenSource()
+        cts.CancelAfter(250)
+        try
+            Async.RunSynchronously(test, cancellationToken=cts.Token) |> ignore
+            Assert.Fail("Expected async to throw")
+        with
+        | :? TaskCanceledException -> Assert.Fail("Did not expect TaskCanceledException")
+        | :? System.OperationCanceledException -> ()
+
     [<Test>]
     member this.AwaitTaskCancellationUntyped () =
         let test() = async {