Cancellable: set token from node/async in features code

src/Compiler/Facilities/BuildGraph.fs

@@ -17,12 +17,14 @@ let wrapThreadStaticInfo computation =
     async {
         let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
         let phase = DiagnosticsThreadStatics.BuildPhase
+        let ct = Cancellable.Token
 
         try
             return! computation
         finally
             DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
             DiagnosticsThreadStatics.BuildPhase <- phase
+            Cancellable.Token <- ct
     }
 
 type Async<'T> with
@@ -125,20 +127,23 @@ type NodeCode private () =
     static member RunImmediate(computation: NodeCode<'T>, ct: CancellationToken) =
         let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
         let phase = DiagnosticsThreadStatics.BuildPhase
+        let ct2 = Cancellable.Token
 
         try
             try
                 let work =
                     async {
                         DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
                         DiagnosticsThreadStatics.BuildPhase <- phase
+                        Cancellable.Token <- ct2
                         return! computation |> Async.AwaitNodeCode
                     }
 
                 Async.StartImmediateAsTask(work, cancellationToken = ct).Result
             finally
                 DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
                 DiagnosticsThreadStatics.BuildPhase <- phase
+                Cancellable.Token <- ct2
         with :? AggregateException as ex when ex.InnerExceptions.Count = 1 ->
             raise (ex.InnerExceptions[0])
 
@@ -148,19 +153,22 @@ type NodeCode private () =
     static member StartAsTask_ForTesting(computation: NodeCode<'T>, ?ct: CancellationToken) =
         let diagnosticsLogger = DiagnosticsThreadStatics.DiagnosticsLogger
         let phase = DiagnosticsThreadStatics.BuildPhase
+        let ct2 = Cancellable.Token
 
         try
             let work =
                 async {
                     DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
                     DiagnosticsThreadStatics.BuildPhase <- phase
+                    Cancellable.Token <- ct2
                     return! computation |> Async.AwaitNodeCode
                 }
 
             Async.StartAsTask(work, cancellationToken = defaultArg ct CancellationToken.None)
         finally
             DiagnosticsThreadStatics.DiagnosticsLogger <- diagnosticsLogger
             DiagnosticsThreadStatics.BuildPhase <- phase
+            Cancellable.Token <- ct2
 
     static member CancellationToken = cancellationToken
 

src/Compiler/Interactive/fsi.fs

@@ -4089,6 +4089,7 @@ type FsiInteractionProcessor
             ?cancellationToken: CancellationToken
         ) =
         let cancellationToken = defaultArg cancellationToken CancellationToken.None
+        use _ = Cancellable.UsingToken(cancellationToken)
 
         if tokenizer.LexBuffer.IsPastEndOfStream then
             let stepStatus =
@@ -4217,6 +4218,7 @@ type FsiInteractionProcessor
 
     member _.EvalInteraction(ctok, sourceText, scriptFileName, diagnosticsLogger, ?cancellationToken) =
         let cancellationToken = defaultArg cancellationToken CancellationToken.None
+        use _ = Cancellable.UsingToken(cancellationToken)
         use _ = UseBuildPhase BuildPhase.Interactive
         use _ = UseDiagnosticsLogger diagnosticsLogger
         use _scope = SetCurrentUICultureForThread fsiOptions.FsiLCID
@@ -4893,6 +4895,7 @@ type FsiEvaluationSession
             SpawnInteractiveServer(fsi, fsiOptions, fsiConsoleOutput)
 
         use _ = UseBuildPhase BuildPhase.Interactive
+        use _ = Cancellable.UsingToken(CancellationToken.None)
 
         if fsiOptions.Interact then
             // page in the type check env

src/Compiler/Service/FSharpCheckerResults.fs

@@ -3560,6 +3560,9 @@ type FsiInteractiveChecker(legacyReferenceResolver, tcConfig: TcConfig, tcGlobal
 
     member _.ParseAndCheckInteraction(sourceText: ISourceText, ?userOpName: string) =
         cancellable {
+            let! ct = Cancellable.token ()
+            use _ = Cancellable.UsingToken(ct)
+
             let userOpName = defaultArg userOpName "Unknown"
             let fileName = Path.Combine(tcConfig.implicitIncludeDir, "stdin.fsx")
             let suggestNamesForErrors = true // Will always be true, this is just for readability

src/Compiler/Service/service.fs

@@ -500,6 +500,9 @@ type BackgroundCompiler
                         Activity.Tags.cache, cache.ToString()
                     |]
 
+            let! ct = Async.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             if cache then
                 let hash = sourceText.GetHashCode() |> int64
 
@@ -541,6 +544,9 @@ type BackgroundCompiler
                     "BackgroundCompiler.GetBackgroundParseResultsForFileInProject"
                     [| Activity.Tags.fileName, fileName; Activity.Tags.userOpName, userOpName |]
 
+            let! ct = NodeCode.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let! builderOpt, creationDiags = getOrCreateBuilder (options, userOpName)
 
             match builderOpt with
@@ -690,6 +696,9 @@ type BackgroundCompiler
                         Activity.Tags.userOpName, userOpName
                     |]
 
+            let! ct = NodeCode.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let! cachedResults =
                 node {
                     let! builderOpt, creationDiags = getAnyBuilder (options, userOpName)
@@ -732,6 +741,9 @@ type BackgroundCompiler
                         Activity.Tags.userOpName, userOpName
                     |]
 
+            let! ct = NodeCode.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let! builderOpt, creationDiags = getOrCreateBuilder (options, userOpName)
 
             match builderOpt with
@@ -760,6 +772,9 @@ type BackgroundCompiler
                         Activity.Tags.userOpName, userOpName
                     |]
 
+            let! ct = NodeCode.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let! builderOpt, creationDiags = getOrCreateBuilder (options, userOpName)
 
             match builderOpt with
@@ -815,6 +830,9 @@ type BackgroundCompiler
                         Activity.Tags.userOpName, userOpName
                     |]
 
+            let! ct = NodeCode.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let! builderOpt, _ = getOrCreateBuilder (options, userOpName)
 
             match builderOpt with
@@ -834,6 +852,9 @@ type BackgroundCompiler
                         Activity.Tags.userOpName, userOpName
                     |]
 
+            let! ct = NodeCode.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let! builderOpt, creationDiags = getOrCreateBuilder (options, userOpName)
 
             match builderOpt with
@@ -974,6 +995,9 @@ type BackgroundCompiler
                         Activity.Tags.userOpName, userOpName
                     |]
 
+            let! ct = NodeCode.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let! builderOpt, _ = getOrCreateBuilder (options, userOpName)
 
             match builderOpt with
@@ -1016,6 +1040,9 @@ type BackgroundCompiler
     /// Parse and typecheck the whole project (the implementation, called recursively as project graph is evaluated)
     member private _.ParseAndCheckProjectImpl(options, userOpName) =
         node {
+            let! ct = NodeCode.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let! builderOpt, creationDiags = getOrCreateBuilder (options, userOpName)
 
             match builderOpt with
@@ -1149,6 +1176,9 @@ type BackgroundCompiler
             // Do we assume .NET Framework references for scripts?
             let assumeDotNetFramework = defaultArg assumeDotNetFramework true
 
+            let! ct = Cancellable.token ()
+            use _ = Cancellable.UsingToken(ct)
+
             let extraFlags =
                 if previewEnabled then
                     [| "--langversion:preview" |]
@@ -1269,6 +1299,9 @@ type BackgroundCompiler
                 |]
 
         async {
+            let! ct = Async.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let! ct = Async.CancellationToken
             // If there was a similar entry (as there normally will have been) then re-establish an empty builder .  This
             // is a somewhat arbitrary choice - it will have the effect of releasing memory associated with the previous
@@ -1514,6 +1547,9 @@ type FSharpChecker
         use _ = Activity.start "FSharpChecker.Compile" [| Activity.Tags.userOpName, _userOpName |]
 
         async {
+            let! ct = Async.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             let ctok = CompilationThreadToken()
             return CompileHelpers.compileFromArgs (ctok, argv, legacyReferenceResolver, None, None)
         }
@@ -1633,6 +1669,9 @@ type FSharpChecker
         let userOpName = defaultArg userOpName "Unknown"
 
         node {
+            let! ct = NodeCode.CancellationToken
+            use _ = Cancellable.UsingToken(ct)
+
             if fastCheck <> Some true || not captureIdentifiersWhenParsing then
                 return! backgroundCompiler.FindReferencesInFile(fileName, options, symbol, canInvalidateProject, userOpName)
             else

src/Compiler/Utilities/Cancellable.fs

@@ -2,47 +2,32 @@ namespace FSharp.Compiler
 
 open System
 open System.Threading
-open Internal.Utilities.Library
 
 [<Sealed>]
 type Cancellable =
     [<ThreadStatic; DefaultValue>]
-    static val mutable private tokens: CancellationToken list
+    static val mutable private token: CancellationToken
 
-    static let disposable =
-        { new IDisposable with
-            member this.Dispose() =
-                Cancellable.Tokens <- Cancellable.Tokens |> List.tail
-        }
+    static member UsingToken(ct) =
+        let oldCt = Cancellable.token
 
-    static member Tokens
-        with private get () =
-            match box Cancellable.tokens with
-            | Null -> []
-            | _ -> Cancellable.tokens
-        and private set v = Cancellable.tokens <- v
+        Cancellable.token <- ct
 
-    static member UsingToken(ct) =
-        Cancellable.Tokens <- ct :: Cancellable.Tokens
-        disposable
+        { new IDisposable with
+            member this.Dispose() = Cancellable.token <- oldCt
+        }
 
-    static member Token =
-        match Cancellable.Tokens with
-        | [] -> CancellationToken.None
-        | token :: _ -> token
+    static member Token
+        with get () = Cancellable.token
+        and internal set v = Cancellable.token <- v
 
-    /// There may be multiple tokens if `UsingToken` is called multiple times, producing scoped structure.
-    /// We're interested in the current, i.e. the most recent, one.
     static member CheckAndThrow() =
-        match Cancellable.Tokens with
-        | [] -> ()
-        | token :: _ -> token.ThrowIfCancellationRequested()
+        Cancellable.token.ThrowIfCancellationRequested()
 
 namespace Internal.Utilities.Library
 
 open System
 open System.Threading
-open FSharp.Compiler
 
 #if !FSHARPCORE_USE_PACKAGE
 open FSharp.Core.CompilerServices.StateMachineHelpers
@@ -63,7 +48,6 @@ module Cancellable =
             ValueOrCancelled.Cancelled(OperationCanceledException ct)
         else
             try
-                use _ = Cancellable.UsingToken(ct)
                 oper ct
             with :? OperationCanceledException as e ->
                 ValueOrCancelled.Cancelled(OperationCanceledException e.CancellationToken)

src/Compiler/Utilities/Cancellable.fsi

@@ -7,6 +7,7 @@ open System.Threading
 type Cancellable =
     static member internal UsingToken: CancellationToken -> IDisposable
     static member Token: CancellationToken
+    static member internal Token: CancellationToken with set
     static member CheckAndThrow: unit -> unit
 
 namespace Internal.Utilities.Library