feat: support `async fn` in macros with coroutine implementation #3540

wyfo · 2023-10-23T12:49:21Z

~~Fixes #1632~~

Here is a draft for the implementation of asyncio coroutine and async fn support. It lacks exhaustive documentation and test, the goal is at first to review the implementation.

davidhewitt

Thank you! I think this will need to go through several cycles of review as I get my head around it, so please bear with me 🙏

tests/test_coroutine.rs

src/impl_/coroutine.rs

src/coroutine.rs

adamreichold · 2023-10-26T08:53:01Z

Just as a process note, I would be glad if we could aim for at least two maintainer reviews here as this appears to be a fundamental expansion of the scope of PyO3. Of course, not necessarily mine, just more than one to increase the bus factor just a little bit.

pyo3-macros-backend/src/method.rs

pyo3-macros-backend/src/utils.rs

src/coroutine.rs

adamreichold · 2023-10-26T09:04:39Z

src/coroutine.rs

+    }
+}
+
+pub(crate) fn iter_result(result: IterNextOutput<PyObject, PyObject>) -> PyResult<PyObject> {


This will make our handling of iterators even more of a mess, won't it? So this is basically what we want to do with all iterators but it would apply only to coroutines? Maybe we should align all of this if we go ahead?

(Sorry if this not clear, but this is mainly directed at @davidhewitt, not so much at @wyfo.)

I honestly think __next__ should allow to return PyResult<PyObject> in addition to PyResult<IterNextOutput<PyObject, PyObject>> or PyResult<Option<PyObject>>.
It would remove the need for this small utilitary function.
But I did want to touch too much parts of PyO3 in this first draft.

We have a long discussion in #3190 about reworking __next__ and __anext__. Really just needs someone to write up a bunch of test cases / examples so we can commit to a design. I keep meaning to do it, but am sort of trying to avoid breaking too much at the same time as #3382

Cargo.toml

src/coroutine.rs

wyfo · 2023-10-26T12:47:13Z

I've updated the PR with the feedbacks.

I've also completely refactored the waker to lazily initialize the Python future (so there is no object allocated when the future is already ready), but I've used unsafe code.
I can also replace the combination AtomicBool + UnsafeCell<MaybeUninit<PyObject>> by a Mutex + an enum, using only safe code thus, but it will obviously impact performance.

Moreover, I've added a method to release the GIL while polling the future (it was the first request when I released pyo3-async).

Obviously, I will need to work on test now, especially with the unsafe code added (if you want to keep it).

src/coroutine.rs

adamreichold · 2023-10-26T16:04:20Z

I've also completely refactored the waker to lazily initialize the Python future (so there is no object allocated when the future is already ready), but I've used unsafe code.
I can also replace the combination AtomicBool + UnsafeCell<MaybeUninit> by a Mutex + an enum, using only safe code thus, but it will obviously impact performance.

Since you acquire the GIL anyway, did you consider our own GILOnceCell to handle the lazy initialization? I am not sure if it is up to the more involved lifecycle but maybe we could extend it slightly to cover this use case?

adamreichold · 2023-10-26T16:05:20Z

(Otherwise, I would prefer an initial safe version using Mutex and an enum which can optimize in a follow-up PR.)

wyfo · 2023-10-26T21:54:44Z

Since you acquire the GIL anyway, did you consider our own GILOnceCell to handle the lazy initialization? I am not if it is up to the more involved lifecycle but maybe we could extend it slightly to cover this use case?

I've considered it, but it had no reinitialization methods, and this unsafe algorithm was the first thing that came into my mind. I've added a take method to GILOnceCell and will use it.

wyfo · 2023-10-27T22:30:18Z

New version with GILOnceCell in the waker!
I've also added some tests (fixing a bug of Future.set_result called after cancellation thanks to them).

I've tried the throw_callback feature in a test, and it works fine. However, I'm not fully sure of the design yet.

codspeed-hq · 2023-10-27T23:14:02Z

CodSpeed Performance Report

Merging #3540 will improve performances by 22.03%

_{Comparing wyfo:coroutine (627841f) with main (abe518d)}

Summary

⚡ 2 improvements
✅ 76 untouched benchmarks

Benchmarks breakdown

	Benchmark	`main`	`wyfo:coroutine`	Change
⚡	`not_a_list_via_downcast`	153.9 ns	126.1 ns	+22.03%
⚡	`list_via_downcast`	153.9 ns	126.1 ns	+22.03%

davidhewitt

Thank you for the continued iteration here, I'm feeling relatively optimistic that we can get to async fn support from this approach. That said, it's definitely an expansion in PyO3's complexity so I agree we'll need a few maintainer passes, and I'd prefer we merged functionality in baby steps where we can be confident in each piece. Sorry I didn't manage to keep up with the threads from the last few days.

I've been thinking a bit about how we make sure that this functionality is a solid foundation which the rest of the ecosystem can build on.

I think that pyo3-asyncio might be somewhat compatible with its existing TaskLocals and other helpers, however I guess it might want to return a custom object instead of our Coroutine. (I haven't thought particularly hard about this.)

Similarly I guess the way that users could support other async runtimes is by supplying their own Coroutine implementation?

From that, it seems to me like we might want to have pyo3(coroutine_type = X) as an option on async functions to switch the returned concrete type. Plus maybe a global api pyo3::set_coroutine_type::<X>() to do this once for the whole extension module.

However I think really we don't need to replace the whole coroutine type, just the waker? (After all, all async def functions in Python return the same coroutine type...)

src/coroutine.rs

davidhewitt · 2023-10-28T13:01:11Z

src/coroutine.rs

+    }
+}
+
+pub(crate) fn iter_result(result: IterNextOutput<PyObject, PyObject>) -> PyResult<PyObject> {


We have a long discussion in #3190 about reworking __next__ and __anext__. Really just needs someone to write up a bunch of test cases / examples so we can commit to a design. I keep meaning to do it, but am sort of trying to avoid breaking too much at the same time as #3382

src/coroutine.rs

src/sync.rs

src/coroutine.rs

tests/test_coroutine.rs

wyfo · 2023-10-29T13:59:02Z

I think that pyo3-asyncio might be somewhat compatible with its existing TaskLocals and other helpers, however I guess it might want to return a custom object instead of our Coroutine. (I haven't thought particularly hard about this.)

pyo3-asyncio is completely different, as future_into_py simply returns an asyncio.Future. IMO, this PR is orthogonal to pyo3-asyncio, and poy3 Coroutine will simply replace pyo3-asyncio future_into_py in the future, as Coroutine can also be used to wrap tokio/async-std JoinHandle.

Similarly I guess the way that users could support other async runtimes is by supplying their own Coroutine implementation?

Do you talk about other Python async runtime? Indeed, it's about coroutine implementation, and more particularly the waker (the current coroutine code would still require some adaptations to make it generic, but nothing complicated).
I will update pyo3-async with the improvements from this PR to test.

davidhewitt · 2023-10-29T14:19:52Z

IMO, this PR is orthogonal to pyo3-asyncio, and poy3 Coroutine will simply replace pyo3-asyncio future_into_py in the future, as Coroutine can also be used to wrap tokio/async-std JoinHandle.

Yes, I agree that mostly this PR will not change things for pyo3-asyncio. What would be nice is for us to understand here how existing users of pyo3-asyncio will use this syntax. E.g. how much pyo3-asyncio will need to change its docs, if it's just a case of adding async fn and removing future_into_py calls, or if there need to be new abstractions in that library to allow migration. cc @awestlake87 in case you have any immediate thoughts on this.

Having at least a fork or branch of pyo3-asyncio with some examples using this async fn syntax would be very instructive for me. I am happy to make time to try this myself, but it might take at least a couple of weeks to find that opportunity 😄

wyfo · 2023-10-29T14:23:59Z

New version here, with a rework of the cancellation handling! Now I'm quite satisfied with it (but I had to replace futures-task by futures-util).

I've also updated macros to take in account CoroutineCancel parameter (without additional attribute) the same way Python parameter is. You can see an example in the test, it's quite easy to use now.

@davidhewitt I took the liberty of resolving some of your conversations, but one remains open, I would like to have your opinion about #3540 (comment)

P.S. I've rebased onto main to use #3556.

wyfo · 2023-11-19T23:59:31Z

I've still added a last commit to this PR with a small improvement I had forgotten in my previous list: panic handling while polling the future, in order to properly close the coroutine (not closing it led to surprising side effect when the coroutine was dropped with the not-awaited warning feature)

wyfo · 2023-11-20T09:32:35Z

I need to update the code because of #3578

davidhewitt

Thanks very much, new docs look good. Just two final points to finish off.

Also, there's a lot of commits in this PR history now. We like to try to keep the history tidy, but GH merge queue sadly doesn't give us the option to squash. Can you please squash manually into one commit when you fixup the docs? (Or several commits if desired.)

guide/src/SUMMARY.md

davidhewitt · 2023-11-21T08:16:59Z

guide/src/async-await.md

Please add this page to the doctests in lib.rs (and the examples may need to have some adjustments to compile).

Other than that, this doc looks good to me 👍

davidhewitt

Yep I was aware about the draft chapters, but I think just slightly better to not have them at all until we're ready to have them.

Thanks for the many revisions here, I think many users will be very excited to have this!

wyfo · 2023-11-22T12:58:37Z

I didn't know about emscripten, but it seems that asyncio is not supported on this target.

Should I add some #[cfg(not(target = "wasm32-unknown-emscripten"))] in the code?

davidhewitt · 2023-11-22T15:27:12Z

Hmm good question. I believe that pyodide might support asyncio, in which case we could potentially benefit from testing that.

For now how about marking the tests which fail with #[cfg_attr(target = "wasm32-unknown-emscripten", ignore)], and let's make an issue to investigate / solve that in a follow-up?

wyfo · 2023-11-22T18:05:21Z

tests/test_coroutine.rs

@@ -0,0 +1,98 @@
+#![cfg(feature = "macros")]
+#![cfg(not(target_arch = "wasm32"))]


@davidhewitt I've added this line to deal with emscripten for now. Is it ok for you? Can I let you open the issue? I don't really know anything about wasm.

davidhewitt · 2023-11-22T18:35:45Z

It's not my area of expertise also, though I've had a couple of rounds of understanding it for PyO3's sake. I've added it to #2582

davidhewitt · 2023-11-22T19:04:04Z

Hmm, guide entry isn't MSRV compatible. I'm sort of okay with that, however that does create potential distribution problems for users who follow the guide. Is there a futures alternative we can pull in as a dev dependency?

wyfo · 2023-11-22T19:19:18Z

Actually, the failing example is the workaround for allow_threads, and I don't really want to use futures::pin_mut! in it because nobody is using it today since std::pin::pin! has been released.
I will mark the example as ignore for now. Anyway, it's a workaround waiting other features to be implemented, so it should disappear soon, and it doesn't even use pyo3.

wyfo · 2023-11-22T20:40:42Z

🥳 🍾

davidhewitt reviewed Oct 23, 2023

View reviewed changes