fix oce filtering in try .. with

Author: majocha

src/Compiler/Utilities/Cancellable.fs

@@ -16,18 +16,11 @@ open System.Threading
 
 [<Sealed>]
 type Cancellable =
-    static let tokenHolder = AsyncLocal<CancellationToken voption>()
+    static let tokenHolder = AsyncLocal<CancellationToken>()
 
-    static let guard =
-        String.IsNullOrWhiteSpace(Environment.GetEnvironmentVariable("DISABLE_CHECKANDTHROW_ASSERT"))
+    static member HasCancellationToken = tokenHolder.Value <> CancellationToken.None
 
-    static let ensureToken msg =
-        tokenHolder.Value
-        |> ValueOption.defaultWith (fun () -> if guard then failwith msg else CancellationToken.None)
-
-    static member HasCancellationToken = tokenHolder.Value.IsSome
-
-    static member Token = ensureToken "Token not available outside of Cancellable computation."
+    static member Token = tokenHolder.Value
 
     static member UseToken() =
         async {
@@ -37,23 +30,22 @@ type Cancellable =
 
     static member UsingToken(ct) =
         let oldCt = tokenHolder.Value
-        tokenHolder.Value <- ValueSome ct
+        tokenHolder.Value <- ct
 
         { new IDisposable with
             member _.Dispose() = tokenHolder.Value <- oldCt
         }
 
     static member CheckAndThrow() =
-        let token = ensureToken "CheckAndThrow invoked outside of Cancellable computation."
-        token.ThrowIfCancellationRequested()
+        tokenHolder.Value.ThrowIfCancellationRequested()
 
     static member TryCheckAndThrow() =
-        match tokenHolder.Value with
-        | ValueNone -> ()
-        | ValueSome token -> token.ThrowIfCancellationRequested()
+        tokenHolder.Value.ThrowIfCancellationRequested()
 
 namespace Internal.Utilities.Library.CancellableImplementation
 
+type Cancellable = FSharp.Compiler.Cancellable
+
 open System
 open System.Threading
 
@@ -62,7 +54,6 @@ open FSharp.Core.CompilerServices.StateMachineHelpers
 open Microsoft.FSharp.Core.CompilerServices
 open System.Runtime.CompilerServices
 open System.Runtime.ExceptionServices
-open System.Diagnostics
 
 type ITrampolineInvocation =
     abstract member MoveNext: unit -> bool
@@ -80,7 +71,7 @@ type PendingInvocation =
     | Immediate of ITrampolineInvocation
 
 [<Sealed>]
-type Trampoline(cancellationToken: CancellationToken) =
+type Trampoline() =
 
     let mutable bindDepth = 0
 
@@ -101,11 +92,6 @@ type Trampoline(cancellationToken: CancellationToken) =
             edi.Throw()
         | _ -> ()
 
-    member this.IsCancelled = cancellationToken.IsCancellationRequested
-
-    member this.ThrowIfCancellationRequested() =
-        cancellationToken.ThrowIfCancellationRequested()
-
     member this.ShoudBounce = bindDepth % bindDepthLimit = 0
 
     member this.SetDelayed(invocation) = pending.Push(Delayed invocation)
@@ -142,8 +128,8 @@ type Trampoline(cancellationToken: CancellationToken) =
 
     static member Current = current.Value.Value
 
-    static member Install ct =
-        current.Value <- ValueSome <| Trampoline ct
+    static member Install() =
+        current.Value <- ValueSome <| Trampoline()
 
 type ITrampolineInvocation<'T> =
     inherit ITrampolineInvocation
@@ -174,12 +160,15 @@ type Cancellable<'T>(clone: unit -> ITrampolineInvocation<'T>) =
 [<AutoOpen>]
 module CancellableCode =
 
-    let inline filterCancellation (catch: exn -> CancellableCode<_, _>) exn =
-        CancellableCode(fun sm -> Trampoline.Current.IsCancelled || (catch exn).Invoke(&sm))
+    let inline filterCancellation (catch: exn -> CancellableCode<_, _>) (exn: exn) =
+        CancellableCode(fun sm ->
+            match exn with
+            | :? OperationCanceledException as oce when oce.CancellationToken = Cancellable.Token -> raise exn
+            | _ -> (catch exn).Invoke(&sm))
 
     let inline throwIfCancellationRequested (code: CancellableCode<_, _>) =
         CancellableCode(fun sm ->
-            Trampoline.Current.ThrowIfCancellationRequested()
+            Cancellable.Token.ThrowIfCancellationRequested()
             code.Invoke(&sm))
 
 type CancellableBuilder() =
@@ -194,6 +183,7 @@ type CancellableBuilder() =
         CancellableCode<'T, _>(fun sm ->
             sm.Data <- value
             true)
+        |> throwIfCancellationRequested
 
     member inline _.Combine
         (code1: CancellableCode<'TOverall, unit>, code2: CancellableCode<'TOverall, 'T>)
@@ -307,7 +297,6 @@ type CancellableBuilder() =
 
 namespace Internal.Utilities.Library
 
-open System
 open System.Threading
 
 type Cancellable<'T> = CancellableImplementation.Cancellable<'T>
@@ -322,19 +311,19 @@ module Cancellable =
 
     let run (code: Cancellable<_>) =
         let invocation = code.GetInvocation()
-        Trampoline.Install FSharp.Compiler.Cancellable.Token
+        Trampoline.Install()
         Trampoline.Current.RunImmediate invocation
         invocation.Result
 
     let runWithoutCancellation code =
-        use _ = FSharp.Compiler.Cancellable.UsingToken CancellationToken.None
+        use _ = Cancellable.UsingToken CancellationToken.None
         run code
 
     let toAsync code =
         async {
-            use! _holder = FSharp.Compiler.Cancellable.UseToken()
+            use! _holder = Cancellable.UseToken()
             return run code
         }
 
     let token () =
-        cancellable { FSharp.Compiler.Cancellable.Token }
+        cancellable { Cancellable.Token }

src/Compiler/Utilities/Cancellable.fsi

@@ -38,8 +38,6 @@ type internal Trampoline =
     member RunImmediate: ITrampolineInvocation -> unit
     member ReplayException: unit -> unit
     static member Current: Trampoline
-    member IsCancelled: bool
-    member ThrowIfCancellationRequested: unit -> unit
     member ShoudBounce: bool
 
 [<NoEquality; NoComparison>]

tests/FSharp.Compiler.Service.Tests/ModuleReaderCancellationTests.fs

@@ -173,11 +173,6 @@ open Ns1.Ns2
 let t: T = T()
 """
 
-
-[<Fact>]
-let ``CheckAndThrow is not allowed to throw outside of cancellable`` () =
-    Assert.Throws<Exception>(fun () -> Cancellable.CheckAndThrow())
-
 [<Fact>]
 let ``Type defs 01 - assembly import`` () =
     let source = source1