From 39739706965e7e53b7263be9532099dbe7e4354a Mon Sep 17 00:00:00 2001 From: Vlad Zarytovskii Date: Wed, 12 Jul 2023 17:33:39 +0200 Subject: [PATCH] Remove state machine allocations from cancellable task on token request --- .../FSharp.Editor/Common/CancellableTasks.fs | 145 +++++++++++++----- 1 file changed, 106 insertions(+), 39 deletions(-) diff --git a/vsintegration/src/FSharp.Editor/Common/CancellableTasks.fs b/vsintegration/src/FSharp.Editor/Common/CancellableTasks.fs index 11ef1ebe152..c4d606f2b43 100644 --- a/vsintegration/src/FSharp.Editor/Common/CancellableTasks.fs +++ b/vsintegration/src/FSharp.Editor/Common/CancellableTasks.fs @@ -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) @@ -545,6 +544,79 @@ module CancellableTasks = ) + /// + /// 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. + /// + [] + static member inline BindDynamic<'TResult1, 'TResult2, 'Awaiter, 'TOverall + when Awaiter<'Awaiter, 'TResult1>> + ( + sm: byref>>, + getAwaiter: 'Awaiter, + [] 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 + + /// Creates an CancellableTask that runs computation, and when + /// computation generates a result T, runs binder res. + /// + /// 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. + /// + /// The computation to provide an unbound result. + /// The function to bind the result of computation. + /// + /// An CancellableTask that performs a monadic bind on the result + /// of computation. + [] + member inline _.Bind<'TResult1, 'TResult2, 'Awaiter, 'TOverall + when Awaiter<'Awaiter, 'TResult1>> + ( + getAwaiter: 'Awaiter, + [] 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)) + + /// Delegates to the input computation. /// /// The existence of this method permits the use of return! in the @@ -560,15 +632,29 @@ module CancellableTasks = : CancellableTaskCode<_, _> = this.Bind(getAwaiter, this.Return) + [] + member inline this.ReturnFrom<'TResult1, 'TResult2, 'Awaiter, 'TOverall when Awaiter<'Awaiter, 'TResult1>> (getAwaiter: 'Awaiter) : CancellableTaskCode<_, _> = + this.Bind(getAwaiter, (fun v -> this.Return v)) + [] member inline this.BindReturn<'TResult1, 'TResult2, 'Awaiter, 'TOverall when Awaiter<'Awaiter, 'TResult1>> ( [] getAwaiter: CancellationToken -> 'Awaiter, - f + [] mapper: 'TResult1 -> 'TResult2 ) : CancellableTaskCode<'TResult2, 'TResult2> = - this.Bind(getAwaiter, (fun v -> this.Return(f v))) + this.Bind(getAwaiter, (fun v -> this.Return(mapper v))) + + + [] + member inline this.BindReturn<'TResult1, 'TResult2, 'Awaiter, 'TOverall + when Awaiter<'Awaiter, 'TResult1>> + ( + getAwaiter: 'Awaiter, + [] mapper: 'TResult1 -> 'TResult2 + ) : CancellableTaskCode<'TResult2, 'TResult2> = + this.Bind(getAwaiter, (fun v -> this.Return(mapper v))) /// Allows the computation expression to turn other types into CancellationToken -> 'Awaiter /// @@ -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 /// Allows the computation expression to turn other types into CancellationToken -> 'Awaiter @@ -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 + /// Allows the computation expression to turn other types into CancellationToken -> 'Awaiter /// @@ -622,10 +705,7 @@ module CancellableTasks = when Awaitable<'Awaitable, 'Awaiter, 'TResult1>> ([] task: CancellationToken -> 'Awaitable) : CancellationToken -> 'Awaiter = - (fun ct -> - task ct - |> Awaitable.getAwaiter - ) + (fun ct -> Awaitable.getAwaiter (task ct)) /// Allows the computation expression to turn other types into CancellationToken -> 'Awaiter @@ -638,10 +718,7 @@ module CancellableTasks = when Awaitable<'Awaitable, 'Awaiter, 'TResult1>> ([] task: unit -> 'Awaitable) : CancellationToken -> 'Awaiter = - (fun _ct -> - task () - |> Awaitable.getAwaiter - ) + (fun _ct -> Awaitable.getAwaiter (task ())) /// Creates an CancellableTask that runs binder(resource). @@ -681,24 +758,19 @@ module CancellableTasks = /// Return an asynchronous computation that will wait for the given task to complete and return /// its result. - static member inline AwaitCancellableTask(t: CancellableTask<'T>) = + static member inline AwaitCancellableTask([] t: CancellableTask<'T>) = async { let! ct = Async.CancellationToken - - return! - t ct - |> Async.AwaitTask + return! Async.AwaitTask (t ct) } /// Return an asynchronous computation that will wait for the given task to complete and return /// its result. - static member inline AwaitCancellableTask(t: CancellableTask) = + static member inline AwaitCancellableTask([] t: CancellableTask) = async { let! ct = Async.CancellationToken - return! - t ct - |> Async.AwaitTask + return! Async.AwaitTask (t ct) } /// Runs an asynchronous computation, starting on the current operating system thread. @@ -767,14 +839,14 @@ module CancellableTasks = /// /// CancellationToken -> 'Awaiter member inline this.Source(computation: Async<'TResult1>) = - this.Source(Async.AsCancellableTask(computation)) + this.Source(fun ct -> Async.AsCancellableTask(computation) ct) /// Allows the computation expression to turn other types into CancellationToken -> 'Awaiter /// /// This turns a CancellableTask<'T> into a CancellationToken -> 'Awaiter. /// /// CancellationToken -> 'Awaiter - member inline _.Source(awaiter: TaskAwaiter<'TResult1>) = (fun _ct -> awaiter) + member inline _.Source(awaiter: TaskAwaiter<'TResult1>) = awaiter /// /// A set of extension methods making it possible to bind against in async computations. @@ -783,13 +855,13 @@ module CancellableTasks = module AsyncExtenions = type Control.AsyncBuilder with - member inline this.Bind(t: CancellableTask<'T>, [] binder: ('T -> Async<'U>)) : Async<'U> = + member inline this.Bind([] t: CancellableTask<'T>, [] binder: ('T -> Async<'U>)) : Async<'U> = this.Bind(Async.AwaitCancellableTask t, binder) member inline this.ReturnFrom([] t: CancellableTask<'T>) : Async<'T> = this.ReturnFrom(Async.AwaitCancellableTask t) - member inline this.Bind([] t: CancellableTask, binder: (unit -> Async<'U>)) : Async<'U> = + member inline this.Bind([] t: CancellableTask, [] binder: (unit -> Async<'U>)) : Async<'U> = this.Bind(Async.AwaitCancellableTask t, binder) member inline this.ReturnFrom([] t: CancellableTask) : Async = @@ -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". /// - let getCurrentCancellationToken () = - cancellableTask.Run( - CancellableTaskCode<_, _>(fun sm -> - sm.Data.Result <- sm.Data.CancellationToken - true - ) - ) + let inline getCurrentCancellationToken () : CancellationToken -> Task = + fun ct -> Task.FromResult ct /// Lifts an item to a CancellableTask. /// The item to be the result of the CancellableTask.