Skip to content

Commit fc125c3

Browse files
majochamajocha
committed
fix poolingValueTask trampolining by using AwaitOnCompleted
1 parent ddf9917 commit fc125c3

File tree

10 files changed

+295
-351
lines changed

10 files changed

+295
-351
lines changed

‎src/IcedTasks/CancellablePoolingValueTask.fs‎

Lines changed: 32 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -60,63 +60,56 @@ module CancellablePoolingValueTasks =
6060
{ new CancellableTaskBaseResumptionDynamicInfo<'T, _>(initialResumptionFunc) with
6161
member info.MoveNext(sm) =
6262
let current = state
63+
let mutable continuation = Stop
64+
65+
let hijackCheck () =
66+
if Trampoline.IncrementBindCount() then
67+
Bounce
68+
else
69+
Immediate
6370

6471
match current with
6572
| InitialYield ->
6673
state <- Running
67-
68-
if Trampoline.IncrementBindCount() then
69-
MethodBuilder.AwaitUnsafeOnCompleted(
70-
&sm.Data.MethodBuilder,
71-
Trampoline.AwaiterRef,
72-
&sm
73-
)
74-
else
75-
info.MoveNext(&sm)
74+
continuation <- hijackCheck ()
7675
| Running ->
7776
try
7877
let step = info.ResumptionFunc.Invoke(&sm)
7978

8079
if step then
8180
state <- SetResult
82-
83-
if Trampoline.IncrementBindCount() then
84-
// Yield before setting result to prevent stack overflow.
85-
MethodBuilder.AwaitUnsafeOnCompleted(
86-
&sm.Data.MethodBuilder,
87-
Trampoline.AwaiterRef,
88-
&sm
89-
)
90-
else
91-
info.MoveNext(&sm)
81+
continuation <- hijackCheck ()
9282
else
93-
match sm.ResumptionDynamicInfo.ResumptionData with
94-
| :? ICriticalNotifyCompletion as awaiter ->
95-
let mutable awaiter = awaiter
96-
97-
MethodBuilder.AwaitUnsafeOnCompleted(
98-
&sm.Data.MethodBuilder,
99-
&awaiter,
100-
&sm
101-
)
102-
| _ -> ()
83+
continuation <-
84+
Await(downcast sm.ResumptionDynamicInfo.ResumptionData)
10385
with exn ->
10486
state <- SetException(ExceptionCache.CaptureOrRetrieve exn)
105-
106-
if Trampoline.IncrementBindCount() then
107-
// Yield before setting exception to prevent stack overflow.
108-
MethodBuilder.AwaitUnsafeOnCompleted(
109-
&sm.Data.MethodBuilder,
110-
Trampoline.AwaiterRef,
111-
&sm
112-
)
113-
else
114-
info.MoveNext(&sm)
87+
continuation <- hijackCheck ()
11588
| SetResult ->
11689
MethodBuilder.SetResult(&sm.Data.MethodBuilder, sm.Data.Result)
11790
| SetException edi ->
11891
MethodBuilder.SetException(&sm.Data.MethodBuilder, edi.SourceException)
11992

93+
let continuation = continuation
94+
95+
match continuation with
96+
| Stop -> ()
97+
| Immediate -> info.MoveNext(&sm)
98+
| Bounce ->
99+
MethodBuilder.AwaitOnCompleted(
100+
&sm.Data.MethodBuilder,
101+
Trampoline.AwaiterRef,
102+
&sm
103+
)
104+
| Await awaiter ->
105+
let mutable awaiter = awaiter
106+
107+
MethodBuilder.AwaitUnsafeOnCompleted(
108+
&sm.Data.MethodBuilder,
109+
&awaiter,
110+
&sm
111+
)
112+
120113
member _.SetStateMachine(sm, state) =
121114
MethodBuilder.SetStateMachine(&sm.Data.MethodBuilder, state)
122115
}

‎src/IcedTasks/CancellableTask.fs‎

Lines changed: 32 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -58,66 +58,56 @@ module CancellableTasks =
5858
{ new CancellableTaskBaseResumptionDynamicInfo<'T, _>(initialResumptionFunc) with
5959
member info.MoveNext(sm) =
6060
let current = state
61+
let mutable continuation = Stop
62+
63+
let hijackCheck () =
64+
if Trampoline.IncrementBindCount() then
65+
Bounce
66+
else
67+
Immediate
6168

6269
match current with
6370
| InitialYield ->
6471
state <- Running
65-
66-
if Trampoline.IncrementBindCount() then
67-
MethodBuilder.AwaitUnsafeOnCompleted(
68-
&sm.Data.MethodBuilder,
69-
Trampoline.AwaiterRef,
70-
&sm
71-
)
72-
73-
else
74-
info.MoveNext(&sm)
72+
continuation <- hijackCheck ()
7573
| Running ->
7674
try
7775
let step = info.ResumptionFunc.Invoke(&sm)
7876

7977
if step then
8078
state <- SetResult
81-
82-
if Trampoline.IncrementBindCount() then
83-
// Yield before setting result to prevent stack overflow.
84-
MethodBuilder.AwaitUnsafeOnCompleted(
85-
&sm.Data.MethodBuilder,
86-
Trampoline.AwaiterRef,
87-
&sm
88-
)
89-
90-
else
91-
info.MoveNext(&sm)
79+
continuation <- hijackCheck ()
9280
else
93-
match sm.ResumptionDynamicInfo.ResumptionData with
94-
| :? ICriticalNotifyCompletion as awaiter ->
95-
let mutable awaiter = awaiter
96-
97-
MethodBuilder.AwaitUnsafeOnCompleted(
98-
&sm.Data.MethodBuilder,
99-
&awaiter,
100-
&sm
101-
)
102-
| _ -> ()
81+
continuation <-
82+
Await(downcast sm.ResumptionDynamicInfo.ResumptionData)
10383
with exn ->
10484
state <- SetException(ExceptionCache.CaptureOrRetrieve exn)
105-
106-
if Trampoline.IncrementBindCount() then
107-
// Yield before setting exception to prevent stack overflow.
108-
MethodBuilder.AwaitUnsafeOnCompleted(
109-
&sm.Data.MethodBuilder,
110-
Trampoline.AwaiterRef,
111-
&sm
112-
)
113-
114-
else
115-
info.MoveNext(&sm)
85+
continuation <- hijackCheck ()
11686
| SetResult ->
11787
MethodBuilder.SetResult(&sm.Data.MethodBuilder, sm.Data.Result)
11888
| SetException edi ->
11989
MethodBuilder.SetException(&sm.Data.MethodBuilder, edi.SourceException)
12090

91+
let continuation = continuation
92+
93+
match continuation with
94+
| Stop -> ()
95+
| Immediate -> info.MoveNext(&sm)
96+
| Bounce ->
97+
MethodBuilder.AwaitOnCompleted(
98+
&sm.Data.MethodBuilder,
99+
Trampoline.AwaiterRef,
100+
&sm
101+
)
102+
| Await awaiter ->
103+
let mutable awaiter = awaiter
104+
105+
MethodBuilder.AwaitUnsafeOnCompleted(
106+
&sm.Data.MethodBuilder,
107+
&awaiter,
108+
&sm
109+
)
110+
121111
member _.SetStateMachine(sm, state) =
122112
MethodBuilder.SetStateMachine(&sm.Data.MethodBuilder, state)
123113
}

‎src/IcedTasks/CancellableValueTask.fs‎

Lines changed: 32 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -60,61 +60,56 @@ module CancellableValueTasks =
6060
{ new CancellableTaskBaseResumptionDynamicInfo<'T, _>(initialResumptionFunc) with
6161
member info.MoveNext(sm) =
6262
let current = state
63+
let mutable continuation = Stop
64+
65+
let hijackCheck () =
66+
if Trampoline.IncrementBindCount() then
67+
Bounce
68+
else
69+
Immediate
6370

6471
match current with
6572
| InitialYield ->
6673
state <- Running
67-
68-
if Trampoline.IncrementBindCount() then
69-
MethodBuilder.AwaitUnsafeOnCompleted(
70-
&sm.Data.MethodBuilder,
71-
Trampoline.AwaiterRef,
72-
&sm
73-
)
74-
else
75-
info.MoveNext(&sm)
74+
continuation <- hijackCheck ()
7675
| Running ->
7776
try
7877
let step = info.ResumptionFunc.Invoke(&sm)
7978

8079
if step then
8180
state <- SetResult
82-
83-
if Trampoline.IncrementBindCount() then
84-
MethodBuilder.AwaitUnsafeOnCompleted(
85-
&sm.Data.MethodBuilder,
86-
Trampoline.AwaiterRef,
87-
&sm
88-
)
89-
else
90-
info.MoveNext(&sm)
81+
continuation <- hijackCheck ()
9182
else
92-
match sm.ResumptionDynamicInfo.ResumptionData with
93-
| :? ICriticalNotifyCompletion as awaiter ->
94-
let mutable awaiter = awaiter
95-
96-
MethodBuilder.AwaitUnsafeOnCompleted(
97-
&sm.Data.MethodBuilder,
98-
&awaiter,
99-
&sm
100-
)
101-
| _ -> ()
83+
continuation <-
84+
Await(downcast sm.ResumptionDynamicInfo.ResumptionData)
10285
with exn ->
10386
state <- SetException(ExceptionCache.CaptureOrRetrieve exn)
104-
105-
if Trampoline.IncrementBindCount() then
106-
MethodBuilder.AwaitUnsafeOnCompleted(
107-
&sm.Data.MethodBuilder,
108-
Trampoline.AwaiterRef,
109-
&sm
110-
)
111-
else
112-
info.MoveNext(&sm)
87+
continuation <- hijackCheck ()
11388
| SetResult ->
11489
MethodBuilder.SetResult(&sm.Data.MethodBuilder, sm.Data.Result)
11590
| SetException edi ->
11691
MethodBuilder.SetException(&sm.Data.MethodBuilder, edi.SourceException)
11792

93+
let continuation = continuation
94+
95+
match continuation with
96+
| Stop -> ()
97+
| Immediate -> info.MoveNext(&sm)
98+
| Bounce ->
99+
MethodBuilder.AwaitOnCompleted(
100+
&sm.Data.MethodBuilder,
101+
Trampoline.AwaiterRef,
102+
&sm
103+
)
104+
| Await awaiter ->
105+
let mutable awaiter = awaiter
106+
107+
MethodBuilder.AwaitUnsafeOnCompleted(
108+
&sm.Data.MethodBuilder,
109+
&awaiter,
110+
&sm
111+
)
112+
118113
member _.SetStateMachine(sm, state) =
119114
MethodBuilder.SetStateMachine(&sm.Data.MethodBuilder, state)
120115
}

‎src/IcedTasks/ColdTask.fs‎

Lines changed: 32 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -315,66 +315,56 @@ module ColdTasks =
315315
{ new ColdTaskResumptionDynamicInfo<'T>(initialResumptionFunc) with
316316
member info.MoveNext(sm) =
317317
let current = state
318+
let mutable continuation = Stop
319+
320+
let hijackCheck () =
321+
if Trampoline.IncrementBindCount() then
322+
Bounce
323+
else
324+
Immediate
318325

319326
match current with
320327
| InitialYield ->
321328
state <- Running
322-
323-
if Trampoline.IncrementBindCount() then
324-
MethodBuilder.AwaitUnsafeOnCompleted(
325-
&sm.Data.MethodBuilder,
326-
Trampoline.AwaiterRef,
327-
&sm
328-
)
329-
330-
else
331-
info.MoveNext(&sm)
329+
continuation <- hijackCheck ()
332330
| Running ->
333331
try
334332
let step = info.ResumptionFunc.Invoke(&sm)
335333

336334
if step then
337335
state <- SetResult
338-
339-
if Trampoline.IncrementBindCount() then
340-
// Yield before setting result to prevent stack overflow.
341-
MethodBuilder.AwaitUnsafeOnCompleted(
342-
&sm.Data.MethodBuilder,
343-
Trampoline.AwaiterRef,
344-
&sm
345-
)
346-
347-
else
348-
info.MoveNext(&sm)
336+
continuation <- hijackCheck ()
349337
else
350-
match sm.ResumptionDynamicInfo.ResumptionData with
351-
| :? ICriticalNotifyCompletion as awaiter ->
352-
let mutable awaiter = awaiter
353-
354-
MethodBuilder.AwaitUnsafeOnCompleted(
355-
&sm.Data.MethodBuilder,
356-
&awaiter,
357-
&sm
358-
)
359-
| _ -> ()
338+
continuation <-
339+
Await(downcast sm.ResumptionDynamicInfo.ResumptionData)
360340
with exn ->
361341
state <- SetException(ExceptionCache.CaptureOrRetrieve exn)
362-
363-
if Trampoline.IncrementBindCount() then
364-
// Yield before setting exception to prevent stack overflow.
365-
MethodBuilder.AwaitUnsafeOnCompleted(
366-
&sm.Data.MethodBuilder,
367-
Trampoline.AwaiterRef,
368-
&sm
369-
)
370-
371-
else
372-
info.MoveNext(&sm)
342+
continuation <- hijackCheck ()
373343
| SetResult ->
374344
MethodBuilder.SetResult(&sm.Data.MethodBuilder, sm.Data.Result)
375345
| SetException edi ->
376346
MethodBuilder.SetException(&sm.Data.MethodBuilder, edi.SourceException)
377347

348+
let continuation = continuation
349+
350+
match continuation with
351+
| Stop -> ()
352+
| Immediate -> info.MoveNext(&sm)
353+
| Bounce ->
354+
MethodBuilder.AwaitOnCompleted(
355+
&sm.Data.MethodBuilder,
356+
Trampoline.AwaiterRef,
357+
&sm
358+
)
359+
| Await awaiter ->
360+
let mutable awaiter = awaiter
361+
362+
MethodBuilder.AwaitUnsafeOnCompleted(
363+
&sm.Data.MethodBuilder,
364+
&awaiter,
365+
&sm
366+
)
367+
378368
member _.SetStateMachine(sm, state) =
379369
MethodBuilder.SetStateMachine(&sm.Data.MethodBuilder, state)
380370
}

0 commit comments

Comments
 (0)