gh-97696: DRAFT asyncio eager tasks factory prototype

[itamaro]

based on GH-98137 + feedback left there

still much left to do here, but something is working, so wanted to get early feedback!

with an opt build and the checked-in benchmarking script:

baseline (no eagerness)

./python.exe async_tree.py -s no_suspension -p
3.12.0a4+ (heads/gather-early-return-dirty:a2a7caf48d, Feb  3 2023, 16:14:11) [Clang 12.0.0 (clang-1200.0.32.28)]
/Users/itamaro/work/pyexe/main-opt/python.exe
Scenario: no_suspension
Time: 1.1034402861259878 s
Tasks created: 55989
Suspense called: 0

eager is ~2x faster in this scenario:

./python.exe async_tree.py -s no_suspension -p -e
3.12.0a4+ (heads/gather-early-return-dirty:a2a7caf48d, Feb  3 2023, 16:14:11) [Clang 12.0.0 (clang-1200.0.32.28)]
/Users/itamaro/work/pyexe/main-opt/python.exe
Scenario: no_suspension
Time: 0.5100075129885226 s
Tasks created: 3
Suspense called: 0

TODOs remaining:

• implement changes in C tasks too
• make needed changes in taskgroups
• handle task constructors and factories that don't support yield_result don't see a reason for this
• fix existing tests
• add new tests
• remove benchmarking script
• extend the async tree benchmark to cover taskgroups
• update pyperformance to support eager factory, run more benchmarks
• handle AsyncStopIteration?
• news and docs

to discuss:

1. do we want the task_factory arg I added to the runner? in this or a separate PR?
2. we are still leaving perf on the table by creating tasks and linking them to a gathering future in gather even when all future already completed. maybe something similar applies to taskgroups too, I didn't look yet. we can handle the case that all futures are completed to return synchronously (in a followup PR). any reason not to do it?
3. even more perf by reimplementing things in C (like gather)

• Issue: Add eager task creation API to asyncio #97696

[DinoV]

Do we really need a separate task_factory here? Ultimately this is just influencing the loop after it's created, couldn't that be rolled into the existing loop_factory mechanism?

[itamaro]

no we don't need this. it was convenient to add it for my testing. sure, the same result can be achieved with the loop factory, but this feels cleaner, so I figured I'd suggest it and see what others think

[DinoV]

The return value should be checked here, I assume if it's NULL we'll want to error out. But there should also presumably be a test case added which will hit that.

[itamaro]

I'd like to find the test case that would trip on this first :)

[DinoV]

This incref seems a little weird? The caller should be keeping this alive for the lifetime of the call, and if task_step2_impl wants to hold onto it then it seems it should do the inc ref.

[itamaro]

I ran into refcnt assertions and adding this made them go away 😬 I don't have a stronger justification for doing it here, so I'll need to dig deeper

[carljm]

from reading the code, I think the issue is on line 2911 in task_step_handle_result_impl, where the result is handed off to the task and there's currently a comment that "no incref is necessary." That comment made sense when result was fully a local variable of task_step_impl, so it could just hand off its owned reference to the task. but now that code is wrong (or at least, is assuming task_step_handle_result_impl steals the reference to its result argument, which isn't the typical/best approach.) So I think adding this incref is actually one "correct" option (as in, the refcounts all end up correct), but it's probably not the clearest / most maintainable option. The better option might be to add the incref down in task_step_handle_result_impl where it gives a reference to the task, and then add a decref of result at the end of task_step_impl after calling task_step_handle_result_impl (since task_step_handle_result_impl will no longer be stealing that reference.)

[itamaro]

good point, thanks @carljm !

tried moving the incref where you suggested and still ran into negative refcounts (in different cases though, like "different loop" test). it works if I incref right in the beginning of the function - we can go with that, or I can try make it more granular by chasing down all the branches that would need the incref the result

[carljm]

So currently the code in task_step_handle_result_impl was written with the assumption that it owns its reference to result. I didn't initially look down through its code far enough to see all the places that assumption manifests, but in addition to line 2911, there's also line 3004, and then below that there are a bunch of error exit cases that do Py_DECREF(result), which only makes sense under the assumption that the reference is owned.

So the choices are either to incref it right away, so that the assumption made by the rest of the existing code continues to be true, or to modify the code so it increfs on line 2911 and 3004 (I think those are the only spots I see?) where it stores a new reference, and so it doesn't decref on all the various error exits at the end of the func.

The latter choice is more typical, because generally doing it more lazily results in fewer reference counting operations (i.e. in one of the error-exit cases, no refcount operations should be needed at all, but the "incref early" option means there's a pointless incref and then decref), and also because it's typically less error prone to future changes (instead of having to remember to decref result on every exit path that does nothing with result, which is harder to remember precisely because you aren't doing anything with result, you only have to incref on the paths where you actually do store a new reference to result, which is pretty natural. But it does mean more changes to the current code.

[itamaro]

thanks for the explanation! super helpful :)
the last commit I just pushed (adding 2 increfs and removing a bunch of decrefs) seems to be working!