Cancellable: hide UsingToken and simplify inlined run (#18309)

* simplify * add comment * already shipped
dotnet · Feb 12, 2025 · 2ea701e · 2ea701e
1 parent 79e8531
commit 2ea701e
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 28 deletions.
diff --git a/src/Compiler/Utilities/Cancellable.fs b/src/Compiler/Utilities/Cancellable.fs
@@ -3,6 +3,16 @@ namespace FSharp.Compiler
 open System
 open System.Threading
 
+// This code provides two methods for handling cancellation in synchronous code:
+// 1. Explicitly, by calling Cancellable.CheckAndThrow().
+// 2. Implicitly, by wrapping the code in a cancellable computation.
+// The cancellable computation propagates the CancellationToken and checks for cancellation implicitly.
+// When it is impractical to use the cancellable computation, such as in deeply nested functions, Cancellable.CheckAndThrow() can be used.
+// It checks a CancellationToken local to the current async execution context, held in AsyncLocal.
+// Before calling Cancellable.CheckAndThrow(), this token must be set.
+// The token is guaranteed to be set during execution of cancellable computation.
+// Otherwise, it can be passed explicitly from the ambient async computation using Cancellable.UseToken().
+
 [<Sealed>]
 type Cancellable =
     static let tokenHolder = AsyncLocal<CancellationToken voption>()
@@ -47,9 +57,7 @@ open System
 open System.Threading
 open FSharp.Compiler
 
-#if !FSHARPCORE_USE_PACKAGE
 open FSharp.Core.CompilerServices.StateMachineHelpers
-#endif
 
 [<RequireQualifiedAccess; Struct>]
 type ValueOrCancelled<'TResult> =
@@ -65,12 +73,7 @@ module Cancellable =
         if ct.IsCancellationRequested then
             ValueOrCancelled.Cancelled(OperationCanceledException ct)
         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
+            oper ct
 
     let fold f acc seq =
         Cancellable(fun ct ->
@@ -84,6 +87,7 @@ module Cancellable =
             acc)
 
     let runWithoutCancellation comp =
+        use _ = Cancellable.UsingToken CancellationToken.None
         let res = run CancellationToken.None comp
 
         match res with
@@ -92,14 +96,19 @@ module Cancellable =
 
     let toAsync c =
         async {
+            use! _holder = Cancellable.UseToken()
+
             let! ct = Async.CancellationToken
-            let res = run ct c
 
             return!
-                Async.FromContinuations(fun (cont, _econt, ccont) ->
-                    match res with
-                    | ValueOrCancelled.Value v -> cont v
-                    | ValueOrCancelled.Cancelled ce -> ccont ce)
+                Async.FromContinuations(fun (cont, econt, ccont) ->
+                    try
+                        match run ct c with
+                        | ValueOrCancelled.Value v -> cont v
+                        | ValueOrCancelled.Cancelled ce -> ccont ce
+                    with
+                    | :? OperationCanceledException as ce when ct.IsCancellationRequested -> ccont ce
+                    | :? OperationCanceledException as e -> InvalidOperationException("Wrong cancellation token", e) |> econt)
         }
 
     let token () = Cancellable(ValueOrCancelled.Value)
@@ -113,39 +122,35 @@ type CancellableBuilder() =
 
     member inline _.Bind(comp, [<InlineIfLambda>] k) =
         Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+
             __debugPoint ""
-#endif
 
             match Cancellable.run ct comp with
             | ValueOrCancelled.Value v1 -> Cancellable.run ct (k v1)
             | ValueOrCancelled.Cancelled err1 -> ValueOrCancelled.Cancelled err1)
 
     member inline _.BindReturn(comp, [<InlineIfLambda>] k) =
         Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+
             __debugPoint ""
-#endif
 
             match Cancellable.run ct 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
+
             __debugPoint ""
-#endif
 
             match Cancellable.run ct comp1 with
             | ValueOrCancelled.Value() -> Cancellable.run ct comp2
             | ValueOrCancelled.Cancelled err1 -> ValueOrCancelled.Cancelled err1)
 
     member inline _.TryWith(comp, [<InlineIfLambda>] handler) =
         Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+
             __debugPoint ""
-#endif
 
             let compRes =
                 try
@@ -164,9 +169,9 @@ type CancellableBuilder() =
 
     member inline _.Using(resource: _ MaybeNull, [<InlineIfLambda>] comp) =
         Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+
             __debugPoint ""
-#endif
+
             let body = comp resource
 
             let compRes =
@@ -188,9 +193,8 @@ type CancellableBuilder() =
 
     member inline _.TryFinally(comp, [<InlineIfLambda>] compensation) =
         Cancellable(fun ct ->
-#if !FSHARPCORE_USE_PACKAGE
+
             __debugPoint ""
-#endif
 
             let compRes =
                 try

diff --git a/src/Compiler/Utilities/Cancellable.fsi b/src/Compiler/Utilities/Cancellable.fsi
@@ -7,9 +7,6 @@ open System.Threading
 type Cancellable =
     static member internal UseToken: unit -> Async<IDisposable>
 
-    /// For use in testing only. Cancellable.token should be set only by the cancellable computation.
-    static member internal UsingToken: CancellationToken -> IDisposable
-
     static member HasCancellationToken: bool
     static member Token: CancellationToken