Skip to content

Commit

Permalink
Cancellable: set token from node/async in features code (dotnet#16348)
Browse files Browse the repository at this point in the history
* Cancellable: set token from node/async in features code

* Fantomas

* Fix wrong token

* Remove assert

* Set token in fsi

* Use cancellation token default value instead of option

* More cancellation token in fsi
  • Loading branch information
auduchinok authored Dec 5, 2023
1 parent 8cc0726 commit d4e3b26
Show file tree
Hide file tree
Showing 6 changed files with 65 additions and 27 deletions.
8 changes: 8 additions & 0 deletions src/Compiler/Facilities/BuildGraph.fs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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])

Expand All @@ -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

Expand Down
3 changes: 3 additions & 0 deletions src/Compiler/Interactive/fsi.fs
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions src/Compiler/Service/FSharpCheckerResults.fs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
39 changes: 39 additions & 0 deletions src/Compiler/Service/service.fs
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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" |]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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
Expand Down
38 changes: 11 additions & 27 deletions src/Compiler/Utilities/Cancellable.fs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down
1 change: 1 addition & 0 deletions src/Compiler/Utilities/Cancellable.fsi
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit d4e3b26

Please sign in to comment.