From 7ea28524e3dde8da142ad0431d7e3ba4c3d6f483 Mon Sep 17 00:00:00 2001
From: Marcio Rinaldi <me@marciorinaldi.com>
Date: Fri, 23 Sep 2016 14:34:26 -0300
Subject: [PATCH] Dispose when async workflow is cancelled (#1555)

* Add test for #1436

* Dispose when async workflow is cancelled
---
 .../Microsoft.FSharp.Control/AsyncModule.fs   | 22 +++++++++++++++++++
 src/fsharp/FSharp.Core/control.fs             |  6 ++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/fsharp/FSharp.Core.Unittests/FSharp.Core/Microsoft.FSharp.Control/AsyncModule.fs b/src/fsharp/FSharp.Core.Unittests/FSharp.Core/Microsoft.FSharp.Control/AsyncModule.fs
index 9c2eae6437e..e13d5d01962 100644
--- a/src/fsharp/FSharp.Core.Unittests/FSharp.Core/Microsoft.FSharp.Control/AsyncModule.fs
+++ b/src/fsharp/FSharp.Core.Unittests/FSharp.Core/Microsoft.FSharp.Control/AsyncModule.fs
@@ -299,6 +299,28 @@ type AsyncModule() =
             Assert.IsTrue(isSet())
         for _i = 1 to 3 do test()
 
+    [<Test>]
+    member this.``OnCancel.RaceBetweenCancellationAndDispose``() = 
+        let flag = ref 0
+        let cts = new System.Threading.CancellationTokenSource()
+        let go = async {
+            use disp =
+                cts.Cancel()
+                { new IDisposable with
+                    override __.Dispose() = incr flag }
+            while true do
+                do! Async.Sleep 50
+            }
+        try
+            Async.RunSynchronously (go, cancellationToken = cts.Token)
+        with
+#if FX_NO_OPERATION_CANCELLED
+            _ -> ()
+#else
+            :? System.OperationCanceledException -> ()
+#endif
+        Assert.AreEqual(1, !flag)
+
     [<Test>]
     member this.``OnCancel.CancelThatWasSignalledBeforeRunningTheComputation``() = 
         let test() = 
diff --git a/src/fsharp/FSharp.Core/control.fs b/src/fsharp/FSharp.Core/control.fs
index 6f58a8c47e0..88798702646 100644
--- a/src/fsharp/FSharp.Core/control.fs
+++ b/src/fsharp/FSharp.Core/control.fs
@@ -909,7 +909,11 @@ namespace Microsoft.FSharp.Control
 
         /// Implement use/Dispose
         let usingA (r:'T :> IDisposable) (f:'T -> Async<'a>) : Async<'a> =
-            tryFinallyA (fun () -> Microsoft.FSharp.Core.LanguagePrimitives.IntrinsicFunctions.Dispose r) (callA f r)
+            let mutable x = 0
+            let disposeFunction _ =
+                if Interlocked.CompareExchange(&x, 1, 0) = 0 then
+                    Microsoft.FSharp.Core.LanguagePrimitives.IntrinsicFunctions.Dispose r
+            tryFinallyA disposeFunction (callA f r) |> whenCancelledA disposeFunction
 
         let ignoreA p = 
             bindA p (fun _ -> doneA)