Skip to content

Remove state machine allocations from cancellable task on token request #15592

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 17, 2023
Merged

Remove state machine allocations from cancellable task on token request #15592

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 106 additions & 39 deletions vsintegration/src/FSharp.Editor/Common/CancellableTasks.fs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,7 +365,6 @@ module CancellableTasks =
&& obj.ReferenceEquals(TaskScheduler.Current, TaskScheduler.Default))
then
fun (ct) ->
// Warning: this will always try to yield even if on thread pool already.
Task.Run<'T>((fun () -> CancellableTaskBuilder.RunDynamicAux (code) (ct)), ct)
else
CancellableTaskBuilder.RunDynamicAux(code)
Expand Down Expand Up @@ -545,6 +544,79 @@ module CancellableTasks =
)


/// <summary>
/// The entry point for the dynamic implementation of the corresponding operation. Do not use directly, only used when executing quotations that involve tasks or other reflective execution of F# code.
/// </summary>
[<NoEagerConstraintApplication>]
static member inline BindDynamic<'TResult1, 'TResult2, 'Awaiter, 'TOverall
when Awaiter<'Awaiter, 'TResult1>>
(
sm: byref<ResumableStateMachine<CancellableTaskStateMachineData<'TOverall>>>,
getAwaiter: 'Awaiter,
[<InlineIfLambda>] continuation: ('TResult1 -> CancellableTaskCode<'TOverall, 'TResult2>)
) : bool =

sm.Data.ThrowIfCancellationRequested()

let mutable awaiter = getAwaiter

let cont =
(CancellableTaskResumptionFunc<'TOverall>(fun sm ->
let result = Awaiter.getResult awaiter
(continuation result).Invoke(&sm)))

// shortcut to continue immediately
if Awaiter.isCompleted awaiter then
cont.Invoke(&sm)
else
sm.ResumptionDynamicInfo.ResumptionData <- (awaiter :> ICriticalNotifyCompletion)
sm.ResumptionDynamicInfo.ResumptionFunc <- cont
false

/// <summary>Creates an CancellableTask that runs computation, and when
/// computation generates a result T, runs binder res.</summary>
///
/// <remarks>A cancellation check is performed when the computation is executed.
///
/// The existence of this method permits the use of let! in the
/// cancellableTask { ... } computation expression syntax.</remarks>
///
/// <param name="getAwaiter">The computation to provide an unbound result.</param>
/// <param name="continuation">The function to bind the result of computation.</param>
///
/// <returns>An CancellableTask that performs a monadic bind on the result
/// of computation.</returns>
[<NoEagerConstraintApplication>]
member inline _.Bind<'TResult1, 'TResult2, 'Awaiter, 'TOverall
when Awaiter<'Awaiter, 'TResult1>>
(
getAwaiter: 'Awaiter,
[<InlineIfLambda>] continuation: ('TResult1 -> CancellableTaskCode<'TOverall, 'TResult2>)
) : CancellableTaskCode<'TOverall, 'TResult2> =

CancellableTaskCode<'TOverall, _>(
fun sm ->
if __useResumableCode then
sm.Data.ThrowIfCancellationRequested()

let mutable awaiter = getAwaiter
let mutable __stack_fin = true

if not (Awaiter.isCompleted awaiter) then
let __stack_yield_fin = ResumableCode.Yield().Invoke(&sm)
__stack_fin <- __stack_yield_fin

if __stack_fin then
let result = Awaiter.getResult awaiter

(continuation result).Invoke(&sm)
else
sm.Data.MethodBuilder.AwaitUnsafeOnCompleted(&awaiter, &sm)
false
else
CancellableTaskBuilderBase.BindDynamic<'TResult1, 'TResult2, 'Awaiter, 'TOverall>(&sm, getAwaiter, continuation))


/// <summary>Delegates to the input computation.</summary>
///
/// <remarks>The existence of this method permits the use of return! in the
Expand All @@ -560,15 +632,29 @@ module CancellableTasks =
: CancellableTaskCode<_, _> =
this.Bind(getAwaiter, this.Return)

[<NoEagerConstraintApplication>]
member inline this.ReturnFrom<'TResult1, 'TResult2, 'Awaiter, 'TOverall when Awaiter<'Awaiter, 'TResult1>> (getAwaiter: 'Awaiter) : CancellableTaskCode<_, _> =
this.Bind(getAwaiter, (fun v -> this.Return v))


[<NoEagerConstraintApplication>]
member inline this.BindReturn<'TResult1, 'TResult2, 'Awaiter, 'TOverall
when Awaiter<'Awaiter, 'TResult1>>
(
[<InlineIfLambda>] getAwaiter: CancellationToken -> 'Awaiter,
f
[<InlineIfLambda>] mapper: 'TResult1 -> 'TResult2
) : CancellableTaskCode<'TResult2, 'TResult2> =
this.Bind(getAwaiter, (fun v -> this.Return(f v)))
this.Bind(getAwaiter, (fun v -> this.Return(mapper v)))


[<NoEagerConstraintApplication>]
member inline this.BindReturn<'TResult1, 'TResult2, 'Awaiter, 'TOverall
when Awaiter<'Awaiter, 'TResult1>>
(
getAwaiter: 'Awaiter,
[<InlineIfLambda>] mapper: 'TResult1 -> 'TResult2
) : CancellableTaskCode<'TResult2, 'TResult2> =
this.Bind(getAwaiter, (fun v -> this.Return(mapper v)))

/// <summary>Allows the computation expression to turn other types into CancellationToken -> 'Awaiter</summary>
///
Expand All @@ -592,8 +678,8 @@ module CancellableTasks =
member inline _.Source<'TResult1, 'TResult2, 'Awaiter, 'TOverall
when Awaiter<'Awaiter, 'TResult1>>
(getAwaiter: 'Awaiter)
: CancellationToken -> 'Awaiter =
(fun _ct -> getAwaiter)
: 'Awaiter =
getAwaiter


/// <summary>Allows the computation expression to turn other types into CancellationToken -> 'Awaiter</summary>
Expand All @@ -605,12 +691,9 @@ module CancellableTasks =
member inline _.Source<'Awaitable, 'TResult1, 'TResult2, 'Awaiter, 'TOverall
when Awaitable<'Awaitable, 'Awaiter, 'TResult1>>
(task: 'Awaitable)
: CancellationToken -> 'Awaiter =
(fun (_ct: CancellationToken) ->
task
|> Awaitable.getAwaiter
)

: 'Awaiter =
Awaitable.getAwaiter task


/// <summary>Allows the computation expression to turn other types into CancellationToken -> 'Awaiter</summary>
///
Expand All @@ -622,10 +705,7 @@ module CancellableTasks =
when Awaitable<'Awaitable, 'Awaiter, 'TResult1>>
([<InlineIfLambda>] task: CancellationToken -> 'Awaitable)
: CancellationToken -> 'Awaiter =
(fun ct ->
task ct
|> Awaitable.getAwaiter
)
(fun ct -> Awaitable.getAwaiter (task ct))


/// <summary>Allows the computation expression to turn other types into CancellationToken -> 'Awaiter</summary>
Expand All @@ -638,10 +718,7 @@ module CancellableTasks =
when Awaitable<'Awaitable, 'Awaiter, 'TResult1>>
([<InlineIfLambda>] task: unit -> 'Awaitable)
: CancellationToken -> 'Awaiter =
(fun _ct ->
task ()
|> Awaitable.getAwaiter
)
(fun _ct -> Awaitable.getAwaiter (task ()))


/// <summary>Creates an CancellableTask that runs binder(resource).
Expand Down Expand Up @@ -681,24 +758,19 @@ module CancellableTasks =

/// <summary>Return an asynchronous computation that will wait for the given task to complete and return
/// its result.</summary>
static member inline AwaitCancellableTask(t: CancellableTask<'T>) =
static member inline AwaitCancellableTask([<InlineIfLambda>] t: CancellableTask<'T>) =
async {
let! ct = Async.CancellationToken

return!
t ct
|> Async.AwaitTask
return! Async.AwaitTask (t ct)
}

/// <summary>Return an asynchronous computation that will wait for the given task to complete and return
/// its result.</summary>
static member inline AwaitCancellableTask(t: CancellableTask) =
static member inline AwaitCancellableTask([<InlineIfLambda>] t: CancellableTask) =
async {
let! ct = Async.CancellationToken

return!
t ct
|> Async.AwaitTask
return! Async.AwaitTask (t ct)
}

/// <summary>Runs an asynchronous computation, starting on the current operating system thread.</summary>
Expand Down Expand Up @@ -767,14 +839,14 @@ module CancellableTasks =
///
/// <returns>CancellationToken -> 'Awaiter</returns>
member inline this.Source(computation: Async<'TResult1>) =
this.Source(Async.AsCancellableTask(computation))
this.Source(fun ct -> Async.AsCancellableTask(computation) ct)

/// <summary>Allows the computation expression to turn other types into CancellationToken -> 'Awaiter</summary>
///
/// <remarks>This turns a CancellableTask&lt;'T&gt; into a CancellationToken -> 'Awaiter.</remarks>
///
/// <returns>CancellationToken -> 'Awaiter</returns>
member inline _.Source(awaiter: TaskAwaiter<'TResult1>) = (fun _ct -> awaiter)
member inline _.Source(awaiter: TaskAwaiter<'TResult1>) = awaiter

/// <summary>
/// A set of extension methods making it possible to bind against <see cref='T:CancellableTask`1'/> in async computations.
Expand All @@ -783,13 +855,13 @@ module CancellableTasks =
module AsyncExtenions =
type Control.AsyncBuilder with

member inline this.Bind(t: CancellableTask<'T>, [<InlineIfLambda>] binder: ('T -> Async<'U>)) : Async<'U> =
member inline this.Bind([<InlineIfLambda>] t: CancellableTask<'T>, [<InlineIfLambda>] binder: ('T -> Async<'U>)) : Async<'U> =
this.Bind(Async.AwaitCancellableTask t, binder)

member inline this.ReturnFrom([<InlineIfLambda>] t: CancellableTask<'T>) : Async<'T> =
this.ReturnFrom(Async.AwaitCancellableTask t)

member inline this.Bind([<InlineIfLambda>] t: CancellableTask, binder: (unit -> Async<'U>)) : Async<'U> =
member inline this.Bind([<InlineIfLambda>] t: CancellableTask, [<InlineIfLambda>] binder: (unit -> Async<'U>)) : Async<'U> =
this.Bind(Async.AwaitCancellableTask t, binder)

member inline this.ReturnFrom([<InlineIfLambda>] t: CancellableTask) : Async<unit> =
Expand Down Expand Up @@ -824,13 +896,8 @@ module CancellableTasks =
/// This will print "2" 2 seconds from start, "3" 3 seconds from start, "5" 5 seconds from start, cease computation and then
/// followed by "Tasks Finished".
/// </example>
let getCurrentCancellationToken () =
cancellableTask.Run(
CancellableTaskCode<_, _>(fun sm ->
sm.Data.Result <- sm.Data.CancellationToken
true
)
)
let inline getCurrentCancellationToken () : CancellationToken -> Task<CancellationToken> =
fun ct -> Task.FromResult ct

/// <summary>Lifts an item to a CancellableTask.</summary>
/// <param name="item">The item to be the result of the CancellableTask.</param>
Expand Down