Cancellable: built-in stack guard

src/Compiler/Checking/CheckDeclarations.fs

@@ -5349,7 +5349,6 @@ let rec TcModuleOrNamespaceElementNonMutRec (cenv: cenv) parent typeNames scopem
               // Now typecheck. 
               let! moduleContents, topAttrsNew, envAtEnd = 
                 TcModuleOrNamespaceElements cenv (Parent (mkLocalModuleRef moduleEntity)) endm envForModule xml None [] moduleDefs
-                |> cenv.stackGuard.GuardCancellable
 
               // Get the inferred type of the decls and record it in the modul. 
               moduleEntity.entity_modul_type <- MaybeLazy.Strict moduleTyAcc.Value
@@ -5440,7 +5439,6 @@ let rec TcModuleOrNamespaceElementNonMutRec (cenv: cenv) parent typeNames scopem
 
           let! moduleContents, topAttrs, envAtEnd = 
             TcModuleOrNamespaceElements cenv parent endm envNS xml mutRecNSInfo [] defs
-            |> cenv.stackGuard.GuardCancellable
 
           MutRecBindingChecking.TcMutRecDefns_UpdateNSContents nsInfo 
           let env, openDecls = 
@@ -5476,14 +5474,11 @@ let rec TcModuleOrNamespaceElementNonMutRec (cenv: cenv) parent typeNames scopem
  }
 
 /// The non-mutually recursive case for a sequence of declarations
-and [<TailCall>] TcModuleOrNamespaceElementsNonMutRec cenv parent typeNames endm (defsSoFar, env, envAtEnd) (moreDefs: SynModuleDecl list) (ct: CancellationToken) =
-
-    if ct.IsCancellationRequested then
-        ValueOrCancelled.Cancelled(OperationCanceledException ct)
-    else
+and TcModuleOrNamespaceElementsNonMutRec cenv parent typeNames endm (defsSoFar, env, envAtEnd) (moreDefs: SynModuleDecl list) =
+    cancellable {
         match moreDefs with
         | [] ->
-            ValueOrCancelled.Value (List.rev defsSoFar, envAtEnd)
+            return List.rev defsSoFar, envAtEnd
         | firstDef :: otherDefs ->
             // Lookahead one to find out the scope of the next declaration.
             let scopem =
@@ -5492,14 +5487,9 @@ and [<TailCall>] TcModuleOrNamespaceElementsNonMutRec cenv parent typeNames endm
                 else
                     unionRanges (List.head otherDefs).Range endm
 
-            let result = Cancellable.run ct (TcModuleOrNamespaceElementNonMutRec cenv parent typeNames scopem env firstDef |> cenv.stackGuard.GuardCancellable)
-
-            match result with
-            | ValueOrCancelled.Cancelled x ->
-                ValueOrCancelled.Cancelled x
-            | ValueOrCancelled.Value(firstDef, env, envAtEnd) ->
-                TcModuleOrNamespaceElementsNonMutRec cenv parent typeNames endm ((firstDef :: defsSoFar), env, envAtEnd) otherDefs ct
-
+            let! firstDef, env, envAtEnd = TcModuleOrNamespaceElementNonMutRec cenv parent typeNames scopem env firstDef
+            return! TcModuleOrNamespaceElementsNonMutRec cenv parent typeNames endm ((firstDef :: defsSoFar), env, envAtEnd) otherDefs
+    }
 
 and TcModuleOrNamespaceElements cenv parent endm env xml mutRecNSInfo openDecls0 synModuleDecls =
   cancellable {
@@ -5524,21 +5514,15 @@ and TcModuleOrNamespaceElements cenv parent endm env xml mutRecNSInfo openDecls0
         return (moduleContents, topAttrsNew, envAtEnd)
 
     | None ->
-        let! ct = Cancellable.token ()
-        let result = TcModuleOrNamespaceElementsNonMutRec cenv parent typeNames endm ([], env, env) synModuleDecls ct
-
-        match result with
-        | ValueOrCancelled.Value(compiledDefs, envAtEnd) ->
-            // Apply the functions for each declaration to build the overall expression-builder
-            let moduleDefs = List.collect p13 compiledDefs
-            let moduleDefs = match openDecls0 with [] -> moduleDefs | _ -> TMDefOpens openDecls0 :: moduleDefs
-            let moduleContents = TMDefs moduleDefs
-
-            // Collect up the attributes that are global to the file
-            let topAttrsNew = List.collect p33 compiledDefs
-            return (moduleContents, topAttrsNew, envAtEnd)
-        | ValueOrCancelled.Cancelled x -> 
-            return! Cancellable(fun _ -> ValueOrCancelled.Cancelled x)
+        let! compiledDefs, envAtEnd = TcModuleOrNamespaceElementsNonMutRec cenv parent typeNames endm ([], env, env) synModuleDecls
+        // Apply the functions for each declaration to build the overall expression-builder
+        let moduleDefs = List.collect p13 compiledDefs
+        let moduleDefs = match openDecls0 with [] -> moduleDefs | _ -> TMDefOpens openDecls0 :: moduleDefs
+        let moduleContents = TMDefs moduleDefs
+
+        // Collect up the attributes that are global to the file
+        let topAttrsNew = List.collect p33 compiledDefs
+        return (moduleContents, topAttrsNew, envAtEnd)
   }
 
 
@@ -5775,7 +5759,6 @@ let CheckOneImplFile
         let defs = [ for x in implFileFrags -> SynModuleDecl.NamespaceFragment x ]
         let! moduleContents, topAttrs, envAtEnd = 
             TcModuleOrNamespaceElements cenv ParentNone qualNameOfFile.Range envinner PreXmlDoc.Empty None openDecls0 defs
-            |> cenv.stackGuard.GuardCancellable
 
         let implFileTypePriorToSig = moduleTyAcc.Value
 

src/Compiler/Driver/CompilerImports.fs

@@ -2255,7 +2255,7 @@ and [<Sealed>] TcImports
             r: AssemblyResolution
         ) : Async<(_ * (unit -> AvailableImportedAssembly list)) option> =
         async {
-            do! Cancellable.UseToken()
+            use! _holder = Cancellable.UseToken()
             CheckDisposed()
             let m = r.originalReference.Range
             let fileName = r.resolvedPath
@@ -2327,6 +2327,7 @@ and [<Sealed>] TcImports
         async {
             CheckDisposed()
 
+            use! _holder = Cancellable.UseToken()
 
             let tcConfig = tcConfigP.Get ctok
 

src/Compiler/Facilities/DiagnosticsLogger.fs

@@ -891,7 +891,6 @@ type StackGuard(maxDepth: int, name: string) =
 
         try
             if depth % maxDepth = 0 then
-
                 async {
                     do! Async.SwitchToNewThread()
                     Thread.CurrentThread.Name <- $"F# Extra Compilation Thread for {name} (depth {depth})"
@@ -903,10 +902,6 @@ type StackGuard(maxDepth: int, name: string) =
         finally
             depth <- depth - 1
 
-    [<DebuggerHidden; DebuggerStepThrough>]
-    member x.GuardCancellable(original: Cancellable<'T>) =
-        Cancellable(fun ct -> x.Guard(fun () -> Cancellable.run ct original))
-
     static member val DefaultDepth =
 #if DEBUG
         GetEnvInteger "FSHARP_DefaultStackGuardDepth" 50

src/Compiler/Facilities/DiagnosticsLogger.fsi

@@ -462,8 +462,6 @@ type StackGuard =
         [<CallerLineNumber; Optional; DefaultParameterValue(0)>] line: int ->
             'T
 
-    member GuardCancellable: Internal.Utilities.Library.Cancellable<'T> -> Internal.Utilities.Library.Cancellable<'T>
-
     static member GetDepthOption: string -> int
 
 /// This represents the global state established as each task function runs as part of the build.

src/Compiler/Utilities/Cancellable.fs

@@ -21,7 +21,7 @@ type Cancellable =
     static member UseToken() =
         async {
             let! ct = Async.CancellationToken
-            tokenHolder.Value <- ValueSome ct
+            return Cancellable.UsingToken ct
         }
 
     static member UsingToken(ct) =
@@ -47,44 +47,53 @@ open System
 open System.Threading
 open FSharp.Compiler
 
-#if !FSHARPCORE_USE_PACKAGE
 open FSharp.Core.CompilerServices.StateMachineHelpers
-#endif
 
 [<RequireQualifiedAccess; Struct>]
 type ValueOrCancelled<'TResult> =
     | Value of result: 'TResult
     | Cancelled of ``exception``: OperationCanceledException
 
 [<Struct>]
-type Cancellable<'T> = Cancellable of (CancellationToken -> ValueOrCancelled<'T>)
+type Cancellable<'T> = Cancellable of (CancellationToken * int -> ValueOrCancelled<'T>)
 
 module Cancellable =
 
-    let inline run (ct: CancellationToken) (Cancellable oper) =
+    let maxDepth = 100
+
+    let handleCheckAndThrow (ct: CancellationToken, depth) oper =
+        try
+            oper (ct, depth)
+        with
+        | :? OperationCanceledException as e when ct.IsCancellationRequested -> ValueOrCancelled.Cancelled e
+        | :? OperationCanceledException as e -> InvalidOperationException("Wrong cancellation token", e) |> raise
+
+    let run (ct: CancellationToken, depth: int) (Cancellable oper) =
         if ct.IsCancellationRequested then
             ValueOrCancelled.Cancelled(OperationCanceledException ct)
+        else if depth % maxDepth = 0 then
+            async {
+                do! Async.SwitchToNewThread()
+                return handleCheckAndThrow (ct, depth + 1) oper
+            }
+            |> Async.RunSynchronously
         else
-            try
-                use _ = Cancellable.UsingToken(ct)
-                oper ct
-            with
-            | :? OperationCanceledException as e when ct.IsCancellationRequested -> ValueOrCancelled.Cancelled e
-            | :? OperationCanceledException as e -> InvalidOperationException("Wrong cancellation token", e) |> raise
+            handleCheckAndThrow (ct, depth + 1) oper
 
     let fold f acc seq =
-        Cancellable(fun ct ->
+        Cancellable(fun state ->
             let mutable acc = ValueOrCancelled.Value acc
 
             for x in seq do
                 match acc with
-                | ValueOrCancelled.Value accv -> acc <- run ct (f accv x)
+                | ValueOrCancelled.Value accv -> acc <- run state (f accv x)
                 | ValueOrCancelled.Cancelled _ -> ()
 
             acc)
 
     let runWithoutCancellation comp =
-        let res = run CancellationToken.None comp
+        use _ = Cancellable.UsingToken CancellationToken.None
+        let res = run (CancellationToken.None, 1) comp
 
         match res with
         | ValueOrCancelled.Cancelled _ -> failwith "unexpected cancellation"
@@ -93,7 +102,8 @@ module Cancellable =
     let toAsync c =
         async {
             let! ct = Async.CancellationToken
-            let res = run ct c
+            use! _holder = Cancellable.UseToken()
+            let res = run (ct, 1) c
 
             return!
                 Async.FromContinuations(fun (cont, _econt, ccont) ->
@@ -102,54 +112,51 @@ module Cancellable =
                     | ValueOrCancelled.Cancelled ce -> ccont ce)
         }
 
-    let token () = Cancellable(ValueOrCancelled.Value)
+    let token () =
+        Cancellable(fun (ct, _) -> ValueOrCancelled.Value ct)
 
 type CancellableBuilder() =
 
     member inline _.Delay([<InlineIfLambda>] f) =
-        Cancellable(fun ct ->
+        Cancellable(fun state ->
             let (Cancellable g) = f ()
-            g ct)
+            g state)
 
     member inline _.Bind(comp, [<InlineIfLambda>] k) =
-        Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+        Cancellable(fun state ->
+
             __debugPoint ""
-#endif
 
-            match Cancellable.run ct comp with
-            | ValueOrCancelled.Value v1 -> Cancellable.run ct (k v1)
+            match Cancellable.run state comp with
+            | ValueOrCancelled.Value v1 -> Cancellable.run state (k v1)
             | ValueOrCancelled.Cancelled err1 -> ValueOrCancelled.Cancelled err1)
 
     member inline _.BindReturn(comp, [<InlineIfLambda>] k) =
-        Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+        Cancellable(fun state ->
+
             __debugPoint ""
-#endif
 
-            match Cancellable.run ct comp with
+            match Cancellable.run state comp with
             | ValueOrCancelled.Value v1 -> ValueOrCancelled.Value(k v1)
             | ValueOrCancelled.Cancelled err1 -> ValueOrCancelled.Cancelled err1)
 
     member inline _.Combine(comp1, comp2) =
-        Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+        Cancellable(fun state ->
+
             __debugPoint ""
-#endif
 
-            match Cancellable.run ct comp1 with
-            | ValueOrCancelled.Value() -> Cancellable.run ct comp2
+            match Cancellable.run state comp1 with
+            | ValueOrCancelled.Value() -> Cancellable.run state comp2
             | ValueOrCancelled.Cancelled err1 -> ValueOrCancelled.Cancelled err1)
 
     member inline _.TryWith(comp, [<InlineIfLambda>] handler) =
-        Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+        Cancellable(fun state ->
+
             __debugPoint ""
-#endif
 
             let compRes =
                 try
-                    match Cancellable.run ct comp with
+                    match Cancellable.run state comp with
                     | ValueOrCancelled.Value res -> ValueOrCancelled.Value(Choice1Of2 res)
                     | ValueOrCancelled.Cancelled exn -> ValueOrCancelled.Cancelled exn
                 with err ->
@@ -159,19 +166,19 @@ type CancellableBuilder() =
             | ValueOrCancelled.Value res ->
                 match res with
                 | Choice1Of2 r -> ValueOrCancelled.Value r
-                | Choice2Of2 err -> Cancellable.run ct (handler err)
+                | Choice2Of2 err -> Cancellable.run state (handler err)
             | ValueOrCancelled.Cancelled err1 -> ValueOrCancelled.Cancelled err1)
 
     member inline _.Using(resource, [<InlineIfLambda>] comp) =
-        Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+        Cancellable(fun state ->
+
             __debugPoint ""
-#endif
+
             let body = comp resource
 
             let compRes =
                 try
-                    match Cancellable.run ct body with
+                    match Cancellable.run state body with
                     | ValueOrCancelled.Value res -> ValueOrCancelled.Value(Choice1Of2 res)
                     | ValueOrCancelled.Cancelled exn -> ValueOrCancelled.Cancelled exn
                 with err ->
@@ -187,14 +194,13 @@ type CancellableBuilder() =
             | ValueOrCancelled.Cancelled err1 -> ValueOrCancelled.Cancelled err1)
 
     member inline _.TryFinally(comp, [<InlineIfLambda>] compensation) =
-        Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+        Cancellable(fun state ->
+
             __debugPoint ""
-#endif
 
             let compRes =
                 try
-                    match Cancellable.run ct comp with
+                    match Cancellable.run state comp with
                     | ValueOrCancelled.Value res -> ValueOrCancelled.Value(Choice1Of2 res)
                     | ValueOrCancelled.Cancelled exn -> ValueOrCancelled.Cancelled exn
                 with err ->
@@ -212,7 +218,8 @@ type CancellableBuilder() =
     member inline _.Return v =
         Cancellable(fun _ -> ValueOrCancelled.Value v)
 
-    member inline _.ReturnFrom(v: Cancellable<'T>) = v
+    member inline _.ReturnFrom(v: Cancellable<'T>) =
+        Cancellable(fun state -> Cancellable.run state v)
 
     member inline _.Zero() =
         Cancellable(fun _ -> ValueOrCancelled.Value())

src/Compiler/Utilities/Cancellable.fsi

@@ -5,10 +5,8 @@ open System.Threading
 
 [<Sealed>]
 type Cancellable =
-    static member internal UseToken: unit -> Async<unit>
-
-    /// For use in testing only. Cancellable.token should be set only by the cancellable computation.
-    static member internal UsingToken: CancellationToken -> IDisposable
+    /// Capture the current cancellation token for use in CheckAndThrow.
+    static member internal UseToken: unit -> Async<IDisposable>
 
     static member HasCancellationToken: bool
     static member Token: CancellationToken
@@ -31,12 +29,12 @@ type internal ValueOrCancelled<'TResult> =
 /// A cancellable computation may be cancelled via a CancellationToken, which is propagated implicitly.
 /// If cancellation occurs, it is propagated as data rather than by raising an OperationCanceledException.
 [<Struct>]
-type internal Cancellable<'T> = Cancellable of (CancellationToken -> ValueOrCancelled<'T>)
+type internal Cancellable<'T> = Cancellable of (CancellationToken * int -> ValueOrCancelled<'T>)
 
 module internal Cancellable =
 
     /// Run a cancellable computation using the given cancellation token
-    val inline run: ct: CancellationToken -> Cancellable<'T> -> ValueOrCancelled<'T>
+    val run: CancellationToken * int -> Cancellable<'T> -> ValueOrCancelled<'T>
 
     val fold: f: ('State -> 'T -> Cancellable<'State>) -> acc: 'State -> seq: seq<'T> -> Cancellable<'State>
 

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

@@ -22,7 +22,6 @@ let mutable private wasCancelled = false
 let runCancelFirstTime f =
     let mutable requestCount = 0
     fun () ->
-        use _ = Cancellable.UsingToken cts.Token
         if requestCount = 0 then
             cts.Cancel()
 
@@ -150,7 +149,6 @@ let referenceReaderProject getPreTypeDefs (cancelOnModuleAccess: bool) (options:
 let parseAndCheck path source options =
     cts <- new CancellationTokenSource()
     wasCancelled <- false
-    use _ = Cancellable.UsingToken cts.Token
 
     try
         match Async.RunSynchronously(checker.ParseAndCheckFileInProject(path, 0, SourceText.ofString source, options), cancellationToken = cts.Token) with