From 1bf4f4c30253db78e9e74b2a82957d32230d2c06 Mon Sep 17 00:00:00 2001 From: Joseph Perez Date: Thu, 30 Nov 2023 10:12:49 +0100 Subject: [PATCH] feat: add `coroutine::await_in_coroutine` to await awaitables in coroutine context --- guide/src/SUMMARY.md | 37 ++-- .../async-await/awaiting_python_awaitables.md | 62 ++++++ newsfragments/3611.added.md | 1 + pyo3-ffi/src/abstract_.rs | 6 +- src/coroutine.rs | 112 ++++++----- src/coroutine/asyncio.rs | 96 ++++++++++ src/coroutine/awaitable.rs | 161 ++++++++++++++++ src/coroutine/waker.rs | 179 ++++++++++-------- src/impl_/coroutine.rs | 2 +- src/lib.rs | 51 +++-- src/tests/common.rs | 20 +- tests/test_await_in_coroutine.rs | 177 +++++++++++++++++ tests/test_coroutine.rs | 43 ++--- 13 files changed, 730 insertions(+), 217 deletions(-) create mode 100644 guide/src/async-await/awaiting_python_awaitables.md create mode 100644 newsfragments/3611.added.md create mode 100644 src/coroutine/asyncio.rs create mode 100644 src/coroutine/awaitable.rs create mode 100644 tests/test_await_in_coroutine.rs diff --git a/guide/src/SUMMARY.md b/guide/src/SUMMARY.md index 4c22c26f587..f1e3eccea8a 100644 --- a/guide/src/SUMMARY.md +++ b/guide/src/SUMMARY.md @@ -6,24 +6,25 @@ - [Getting started](getting-started.md) - [Using Rust from Python](rust-from-python.md) - - [Python modules](module.md) - - [Python functions](function.md) - - [Function signatures](function/signature.md) - - [Error handling](function/error-handling.md) - - [Python classes](class.md) - - [Class customizations](class/protocols.md) - - [Basic object customization](class/object.md) - - [Emulating numeric types](class/numeric.md) - - [Emulating callable objects](class/call.md) + - [Python modules](module.md) + - [Python functions](function.md) + - [Function signatures](function/signature.md) + - [Error handling](function/error-handling.md) + - [Python classes](class.md) + - [Class customizations](class/protocols.md) + - [Basic object customization](class/object.md) + - [Emulating numeric types](class/numeric.md) + - [Emulating callable objects](class/call.md) - [Calling Python from Rust](python-from-rust.md) - - [Python object types](types.md) - - [Python exceptions](exception.md) - - [Calling Python functions](python-from-rust/function-calls.md) - - [Executing existing Python code](python-from-rust/calling-existing-code.md) + - [Python object types](types.md) + - [Python exceptions](exception.md) + - [Calling Python functions](python-from-rust/function-calls.md) + - [Executing existing Python code](python-from-rust/calling-existing-code.md) - [Type conversions](conversions.md) - - [Mapping of Rust types to Python types](conversions/tables.md) - - [Conversion traits](conversions/traits.md) + - [Mapping of Rust types to Python types](conversions/tables.md) + - [Conversion traits](conversions/traits.md) - [Using `async` and `await`](async-await.md) + - [Awaiting Python awaitables](async-await/awaiting_python_awaitables) - [Parallelism](parallelism.md) - [Debugging](debugging.md) - [Features reference](features.md) @@ -31,10 +32,10 @@ - [Performance](performance.md) - [Advanced topics](advanced.md) - [Building and distribution](building-and-distribution.md) - - [Supporting multiple Python versions](building-and-distribution/multiple-python-versions.md) + - [Supporting multiple Python versions](building-and-distribution/multiple-python-versions.md) - [Useful crates](ecosystem.md) - - [Logging](ecosystem/logging.md) - - [Using `async` and `await`](ecosystem/async-await.md) + - [Logging](ecosystem/logging.md) + - [Using `async` and `await`](ecosystem/async-await.md) - [FAQ and troubleshooting](faq.md) --- diff --git a/guide/src/async-await/awaiting_python_awaitables.md b/guide/src/async-await/awaiting_python_awaitables.md new file mode 100644 index 00000000000..febecc6a286 --- /dev/null +++ b/guide/src/async-await/awaiting_python_awaitables.md @@ -0,0 +1,62 @@ +# Awaiting Python awaitables + +Python awaitable can be awaited on Rust side +using [`await_in_coroutine`]({{#PYO3_DOCS_URL}}/pyo3/coroutine/function.await_in_coroutine). + +```rust +# # ![allow(dead_code)] +# #[cfg(feature = "experimental-async")] { +use pyo3::{prelude::*, coroutine::await_in_coroutine}; + +#[pyfunction] +async fn wrap_awaitable(awaitable: PyObject) -> PyResult { + Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?.await +} +# } +``` + +Behind the scene, `await_in_coroutine` calls the `__await__` method of the Python awaitable (or `__iter__` for +generator-based coroutine). + +## Restrictions + +As the name suggests, `await_in_coroutine` resulting future can only be awaited in coroutine context. Otherwise, it +panics. + +```rust +# # ![allow(dead_code)] +# #[cfg(feature = "experimental-async")] { +use pyo3::{prelude::*, coroutine::await_in_coroutine}; + +#[pyfunction] +fn block_on(awaitable: PyObject) -> PyResult { + let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?; + futures::executor::block_on(future) // ERROR: PyFuture must be awaited in coroutine context +} +# } +``` + +The future must also be the only one to be awaited at a time; it means that it's forbidden to await it in a `select!`. +Otherwise, it panics. + +```rust +# # ![allow(dead_code)] +# #[cfg(feature = "experimental-async")] { +use futures::FutureExt; +use pyo3::{prelude::*, coroutine::await_in_coroutine}; + +#[pyfunction] +async fn select(awaitable: PyObject) -> PyResult { + let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?; + futures::select_biased! { + _ = std::future::pending::<()>().fuse() => unreachable!(), + res = future.fuse() => res, // ERROR: Python awaitable mixed with Rust future + } +} +# } +``` + +These restrictions exist because awaiting a `await_in_coroutine` future strongly binds it to the +enclosing coroutine. The coroutine will then delegate its `send`/`throw`/`close` methods to the +awaited future. If it was awaited in a `select!`, `Coroutine::send` would no able to know if +the value passed would have to be delegated or not. diff --git a/newsfragments/3611.added.md b/newsfragments/3611.added.md new file mode 100644 index 00000000000..75aa9aee8fd --- /dev/null +++ b/newsfragments/3611.added.md @@ -0,0 +1 @@ +Add `coroutine::await_in_coroutine` to await awaitables in coroutine context diff --git a/pyo3-ffi/src/abstract_.rs b/pyo3-ffi/src/abstract_.rs index b5bf9cc3d35..3cc5cd346ad 100644 --- a/pyo3-ffi/src/abstract_.rs +++ b/pyo3-ffi/src/abstract_.rs @@ -129,7 +129,11 @@ extern "C" { pub fn PyIter_Next(arg1: *mut PyObject) -> *mut PyObject; #[cfg(all(not(PyPy), Py_3_10))] #[cfg_attr(PyPy, link_name = "PyPyIter_Send")] - pub fn PyIter_Send(iter: *mut PyObject, arg: *mut PyObject, presult: *mut *mut PyObject); + pub fn PyIter_Send( + iter: *mut PyObject, + arg: *mut PyObject, + presult: *mut *mut PyObject, + ) -> c_int; #[cfg_attr(PyPy, link_name = "PyPyNumber_Check")] pub fn PyNumber_Check(o: *mut PyObject) -> c_int; diff --git a/src/coroutine.rs b/src/coroutine.rs index a31d864e826..47c0378e6e0 100644 --- a/src/coroutine.rs +++ b/src/coroutine.rs @@ -11,21 +11,30 @@ use std::{ use pyo3_macros::{pyclass, pymethods}; use crate::{ - coroutine::{cancel::ThrowCallback, waker::AsyncioWaker}, - exceptions::{PyAttributeError, PyRuntimeError, PyStopIteration}, + coroutine::waker::CoroutineWaker, + exceptions::{PyAttributeError, PyGeneratorExit, PyRuntimeError, PyStopIteration}, + marker::Ungil, panic::PanicException, - types::{string::PyStringMethods, PyIterator, PyString}, - Bound, IntoPy, Py, PyAny, PyErr, PyObject, PyResult, Python, + types::{string::PyStringMethods, PyString}, + IntoPy, Py, PyErr, PyObject, PyResult, Python, }; -pub(crate) mod cancel; +mod asyncio; +mod awaitable; +mod cancel; mod waker; -use crate::marker::Ungil; -pub use cancel::CancelHandle; +pub use awaitable::await_in_coroutine; +pub use cancel::{CancelHandle, ThrowCallback}; const COROUTINE_REUSED_ERROR: &str = "cannot reuse already awaited coroutine"; +pub(crate) enum CoroOp { + Send(PyObject), + Throw(PyObject), + Close, +} + trait CoroutineFuture: Send { fn poll(self: Pin<&mut Self>, py: Python<'_>, waker: &Waker) -> Poll>; } @@ -69,7 +78,7 @@ pub struct Coroutine { qualname_prefix: Option<&'static str>, throw_callback: Option, future: Option>>, - waker: Option>, + waker: Option>, } impl Coroutine { @@ -104,58 +113,55 @@ impl Coroutine { } } - fn poll(&mut self, py: Python<'_>, throw: Option) -> PyResult { + fn poll_inner(&mut self, py: Python<'_>, mut op: CoroOp) -> PyResult { // raise if the coroutine has already been run to completion let future_rs = match self.future { Some(ref mut fut) => fut, None => return Err(PyRuntimeError::new_err(COROUTINE_REUSED_ERROR)), }; - // reraise thrown exception it - match (throw, &self.throw_callback) { - (Some(exc), Some(cb)) => cb.throw(exc), - (Some(exc), None) => { - self.close(); - return Err(PyErr::from_value_bound(exc.into_bound(py))); - } - (None, _) => {} + // if the future is not pending on a Python awaitable, + // execute throw callback or complete on close + if !matches!(self.waker, Some(ref w) if w.is_delegated(py)) { + match op { + send @ CoroOp::Send(_) => op = send, + CoroOp::Throw(exc) => match &self.throw_callback { + Some(cb) => { + cb.throw(exc.clone_ref(py)); + op = CoroOp::Send(py.None()); + } + None => return Err(PyErr::from_value_bound(exc.into_bound(py))), + }, + CoroOp::Close => return Err(PyGeneratorExit::new_err(py.None())), + }; } // create a new waker, or try to reset it in place if let Some(waker) = self.waker.as_mut().and_then(Arc::get_mut) { - waker.reset(); + waker.reset(op); } else { - self.waker = Some(Arc::new(AsyncioWaker::new())); + self.waker = Some(Arc::new(CoroutineWaker::new(op))); } - // poll the future and forward its results if ready + // poll the future and forward its results if ready; otherwise, yield from waker // polling is UnwindSafe because the future is dropped in case of panic let waker = Waker::from(self.waker.clone().unwrap()); let poll = || future_rs.as_mut().poll(py, &waker); match panic::catch_unwind(panic::AssertUnwindSafe(poll)) { - Ok(Poll::Ready(res)) => { - self.close(); - return Err(PyStopIteration::new_err(res?)); - } - Err(err) => { - self.close(); - return Err(PanicException::from_panic_payload(err)); - } - _ => {} + Err(err) => Err(PanicException::from_panic_payload(err)), + Ok(Poll::Ready(res)) => Err(PyStopIteration::new_err(res?)), + Ok(Poll::Pending) => match self.waker.as_ref().unwrap().yield_(py) { + Ok(to_yield) => Ok(to_yield), + Err(err) => Err(err), + }, } - // otherwise, initialize the waker `asyncio.Future` - if let Some(future) = self.waker.as_ref().unwrap().initialize_future(py)? { - // `asyncio.Future` must be awaited; fortunately, it implements `__iter__ = __await__` - // and will yield itself if its result has not been set in polling above - if let Some(future) = PyIterator::from_bound_object(&future.as_borrowed()) - .unwrap() - .next() - { - // future has not been leaked into Python for now, and Rust code can only call - // `set_result(None)` in `Wake` implementation, so it's safe to unwrap - return Ok(future.unwrap().into()); - } + } + + fn poll(&mut self, py: Python<'_>, op: CoroOp) -> PyResult { + let result = self.poll_inner(py, op); + if result.is_err() { + // the Rust future is dropped, and the field set to `None` + // to indicate the coroutine has been run to completion + drop(self.future.take()); } - // if waker has been waken during future polling, this is roughly equivalent to - // `await asyncio.sleep(0)`, so just yield `None`. - Ok(py.None().into_py(py)) + result } } @@ -180,18 +186,20 @@ impl Coroutine { } } - fn send(&mut self, py: Python<'_>, _value: &Bound<'_, PyAny>) -> PyResult { - self.poll(py, None) + fn send(&mut self, py: Python<'_>, value: PyObject) -> PyResult { + self.poll(py, CoroOp::Send(value)) } fn throw(&mut self, py: Python<'_>, exc: PyObject) -> PyResult { - self.poll(py, Some(exc)) + self.poll(py, CoroOp::Throw(exc)) } - fn close(&mut self) { - // the Rust future is dropped, and the field set to `None` - // to indicate the coroutine has been run to completion - drop(self.future.take()); + fn close(&mut self, py: Python<'_>) -> PyResult<()> { + match self.poll(py, CoroOp::Close) { + Ok(_) => Ok(()), + Err(err) if err.is_instance_of::(py) => Ok(()), + Err(err) => Err(err), + } } fn __await__(self_: Py) -> Py { @@ -199,6 +207,6 @@ impl Coroutine { } fn __next__(&mut self, py: Python<'_>) -> PyResult { - self.poll(py, None) + self.poll(py, CoroOp::Send(py.None())) } } diff --git a/src/coroutine/asyncio.rs b/src/coroutine/asyncio.rs new file mode 100644 index 00000000000..0049b56c0b8 --- /dev/null +++ b/src/coroutine/asyncio.rs @@ -0,0 +1,96 @@ +//! Coroutine implementation compatible with asyncio. +use pyo3_macros::pyfunction; + +use crate::{ + intern, + sync::GILOnceCell, + types::{PyAnyMethods, PyCFunction, PyIterator}, + wrap_pyfunction_bound, Bound, IntoPy, Py, PyAny, PyObject, PyResult, Python, +}; + +/// `asyncio.get_running_loop` +fn get_running_loop(py: Python<'_>) -> PyResult> { + static GET_RUNNING_LOOP: GILOnceCell = GILOnceCell::new(); + let import = || -> PyResult<_> { + let module = py.import_bound("asyncio")?; + Ok(module.getattr("get_running_loop")?.into()) + }; + GET_RUNNING_LOOP + .get_or_try_init(py, import)? + .bind(py) + .call0() +} + +/// Asyncio-compatible coroutine waker. +/// +/// Polling a Rust future yields an `asyncio.Future`, whose `set_result` method is called +/// when `Waker::wake` is called. +pub(super) struct AsyncioWaker { + event_loop: PyObject, + future: PyObject, +} + +impl AsyncioWaker { + pub(super) fn new(py: Python<'_>) -> PyResult { + let event_loop = get_running_loop(py)?.into_py(py); + let future = event_loop.call_method0(py, "create_future")?; + Ok(Self { event_loop, future }) + } + + pub(super) fn yield_(&self, py: Python<'_>) -> PyResult { + let __await__; + // `asyncio.Future` must be awaited; in normal case, it implements `__iter__ = __await__`, + // but `create_future` may have been overriden + let mut iter = match PyIterator::from_bound_object(self.future.bind(py)) { + Ok(iter) => iter, + Err(_) => { + __await__ = self.future.call_method0(py, intern!(py, "__await__"))?; + PyIterator::from_bound_object(__await__.bind(py))? + } + }; + // future has not been wakened (because `yield_waken` would have been called + // otherwise), so it is expected to yield itself + Ok(iter.next().expect("future didn't yield")?.into_py(py)) + } + + #[allow(clippy::unnecessary_wraps)] + pub(super) fn yield_waken(py: Python<'_>) -> PyResult { + Ok(py.None()) + } + + pub(super) fn wake(&self, py: Python<'_>) -> PyResult<()> { + static RELEASE_WAITER: GILOnceCell> = GILOnceCell::new(); + let release_waiter = RELEASE_WAITER.get_or_try_init(py, || { + wrap_pyfunction_bound!(release_waiter, py).map(Into::into) + })?; + // `Future.set_result` must be called in event loop thread, + // so it requires `call_soon_threadsafe` + let call_soon_threadsafe = self.event_loop.call_method1( + py, + intern!(py, "call_soon_threadsafe"), + (release_waiter, &self.future), + ); + if let Err(err) = call_soon_threadsafe { + // `call_soon_threadsafe` will raise if the event loop is closed; + // instead of catching an unspecific `RuntimeError`, check directly if it's closed. + let is_closed = self.event_loop.call_method0(py, "is_closed")?; + if !is_closed.extract(py)? { + return Err(err); + } + } + Ok(()) + } +} + +/// Call `future.set_result` if the future is not done. +/// +/// Future can be cancelled by the event loop before being wakened. +/// See +#[pyfunction(crate = "crate")] +fn release_waiter(future: &Bound<'_, PyAny>) -> PyResult<()> { + let done = future.call_method0(intern!(future.py(), "done"))?; + if !done.extract::()? { + future.call_method1(intern!(future.py(), "set_result"), (future.py().None(),))?; + } + Ok(()) +} diff --git a/src/coroutine/awaitable.rs b/src/coroutine/awaitable.rs new file mode 100644 index 00000000000..9eee03a0f07 --- /dev/null +++ b/src/coroutine/awaitable.rs @@ -0,0 +1,161 @@ +use std::{ + future::Future, + pin::Pin, + task::{Context, Poll}, +}; + +use super::waker::try_delegate; +use crate::{ + coroutine::CoroOp, + exceptions::{PyAttributeError, PyTypeError}, + intern, + sync::GILOnceCell, + types::{PyAnyMethods, PyIterator, PyTypeMethods}, + Bound, PyAny, PyErr, PyObject, PyResult, Python, +}; + +const NOT_IN_COROUTINE_CONTEXT: &str = "PyFuture must be awaited in coroutine context"; + +fn is_awaitable(obj: &Bound<'_, PyAny>) -> PyResult { + static IS_AWAITABLE: GILOnceCell = GILOnceCell::new(); + let import = || { + PyResult::Ok( + obj.py() + .import_bound("inspect")? + .getattr("isawaitable")? + .into(), + ) + }; + IS_AWAITABLE + .get_or_try_init(obj.py(), import)? + .call1(obj.py(), (obj,))? + .extract(obj.py()) +} + +pub(crate) enum YieldOrReturn { + Return(PyObject), + Yield(PyObject), +} + +pub(crate) fn delegate( + py: Python<'_>, + await_impl: PyObject, + op: &CoroOp, +) -> PyResult { + match op { + CoroOp::Send(obj) => { + cfg_if::cfg_if! { + if #[cfg(all(Py_3_10, not(PyPy), not(Py_LIMITED_API)))] { + let mut result = std::ptr::null_mut(); + match unsafe { crate::ffi::PyIter_Send(await_impl.as_ptr(), obj.as_ptr(), &mut result) } + { + -1 => Err(PyErr::take(py).unwrap()), + 0 => Ok(YieldOrReturn::Return(unsafe { + PyObject::from_owned_ptr(py, result) + })), + 1 => Ok(YieldOrReturn::Yield(unsafe { + PyObject::from_owned_ptr(py, result) + })), + _ => unreachable!(), + } + } else { + let send = intern!(py, "send"); + if obj.is_none(py) || !await_impl.bind(py).hasattr(send).unwrap_or(false) { + await_impl.call_method0(py, intern!(py, "__next__")) + } else { + await_impl.call_method1(py, send, (obj,)) + } + .map(YieldOrReturn::Yield) + } + } + } + CoroOp::Throw(exc) => { + let throw = intern!(py, "throw"); + if await_impl.bind(py).hasattr(throw).unwrap_or(false) { + await_impl + .call_method1(py, throw, (exc,)) + .map(YieldOrReturn::Yield) + } else { + Err(PyErr::from_value_bound(exc.bind(py).clone())) + } + } + CoroOp::Close => { + let close = intern!(py, "close"); + if await_impl.bind(py).hasattr(close).unwrap_or(false) { + await_impl + .call_method0(py, close) + .map(YieldOrReturn::Return) + } else { + Ok(YieldOrReturn::Return(py.None())) + } + } + } +} + +struct AwaitImpl(PyObject); + +impl AwaitImpl { + fn new(obj: &Bound<'_, PyAny>) -> PyResult { + let __await__ = intern!(obj.py(), "__await__"); + match obj.call_method0(__await__) { + Ok(iter) => Ok(Self(iter.unbind())), + Err(err) if err.is_instance_of::(obj.py()) => { + if obj.hasattr(__await__)? { + Err(err) + } else if is_awaitable(obj)? { + Ok(Self( + PyIterator::from_bound_object(obj)?.unbind().into_any(), + )) + } else { + Err(PyTypeError::new_err(format!( + "object {tp} can't be used in 'await' expression", + tp = obj.get_type().name()? + ))) + } + } + Err(err) => Err(err), + } + } +} + +impl Future for AwaitImpl { + type Output = PyResult; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match try_delegate(cx.waker(), self.0.clone()) { + Some(poll) => poll, + None => panic!("{}", NOT_IN_COROUTINE_CONTEXT), + } + } +} + +/// Allows awaiting arbitrary Python awaitable inside PyO3 coroutine context, e.g. async pyfunction. +/// +/// Awaiting the resulting future will panic if it's not done in coroutine context. +/// However, the future can be instantiated outside of coroutine context. +/// +/// ```rust +/// use pyo3::{coroutine::await_in_coroutine, prelude::*, py_run, wrap_pyfunction_bound}; +/// +/// # fn main() { +/// #[pyfunction] +/// async fn wrap_awaitable(awaitable: PyObject) -> PyResult { +/// let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?; +/// future.await +/// } +/// Python::with_gil(|py| { +/// let wrap_awaitable = wrap_pyfunction_bound!(wrap_awaitable, py).unwrap(); +/// let test = r#" +/// import asyncio +/// assert asyncio.run(wrap_awaitable(asyncio.sleep(1, result=42))) == 42 +/// "#; +/// py_run!(py, wrap_awaitable, test); +/// }) +/// # } +/// ``` +/// ```rust +pub fn await_in_coroutine( + obj: &Bound<'_, PyAny>, +) -> PyResult> + Send + Sync + 'static> { + AwaitImpl::new(obj) +} diff --git a/src/coroutine/waker.rs b/src/coroutine/waker.rs index fc7c54e1f5a..4e9592c25c3 100644 --- a/src/coroutine/waker.rs +++ b/src/coroutine/waker.rs @@ -1,106 +1,119 @@ -use crate::sync::GILOnceCell; -use crate::types::any::PyAnyMethods; -use crate::types::PyCFunction; -use crate::{intern, wrap_pyfunction_bound, Bound, Py, PyAny, PyObject, PyResult, Python}; -use pyo3_macros::pyfunction; -use std::sync::Arc; -use std::task::Wake; +use std::{ + cell::Cell, + sync::Arc, + task::{Poll, Wake, Waker}, +}; -/// Lazy `asyncio.Future` wrapper, implementing [`Wake`] by calling `Future.set_result`. -/// -/// asyncio future is let uninitialized until [`initialize_future`][1] is called. -/// If [`wake`][2] is called before future initialization (during Rust future polling), -/// [`initialize_future`][1] will return `None` (it is roughly equivalent to `asyncio.sleep(0)`) -/// -/// [1]: AsyncioWaker::initialize_future -/// [2]: AsyncioWaker::wake -pub struct AsyncioWaker(GILOnceCell>); +use crate::{ + coroutine::{ + asyncio::AsyncioWaker, + awaitable::{delegate, YieldOrReturn}, + CoroOp, + }, + exceptions::PyStopIteration, + intern, + sync::GILOnceCell, + types::PyAnyMethods, + Bound, PyObject, PyResult, Python, +}; -impl AsyncioWaker { - pub(super) fn new() -> Self { - Self(GILOnceCell::new()) +const MIXED_AWAITABLE_AND_FUTURE_ERROR: &str = "Python awaitable mixed with Rust future"; + +enum State { + Pending(AsyncioWaker), + Waken, + Delegated(PyObject), +} + +pub(super) struct CoroutineWaker { + state: GILOnceCell, + op: CoroOp, +} + +impl CoroutineWaker { + pub(super) fn new(op: CoroOp) -> Self { + Self { + state: GILOnceCell::new(), + op, + } } - pub(super) fn reset(&mut self) { - self.0.take(); + pub(super) fn reset(&mut self, op: CoroOp) { + self.state.take(); + self.op = op; } - pub(super) fn initialize_future<'py>( - &self, - py: Python<'py>, - ) -> PyResult>> { - let init = || LoopAndFuture::new(py).map(Some); - let loop_and_future = self.0.get_or_try_init(py, init)?.as_ref(); - Ok(loop_and_future.map(|LoopAndFuture { future, .. }| future.bind(py))) + pub(super) fn is_delegated(&self, py: Python<'_>) -> bool { + matches!(self.state.get(py), Some(State::Delegated(_))) + } + + pub(super) fn yield_(&self, py: Python<'_>) -> PyResult { + let init = || PyResult::Ok(State::Pending(AsyncioWaker::new(py)?)); + let state = self.state.get_or_try_init(py, init)?; + match state { + State::Pending(waker) => waker.yield_(py), + State::Waken => AsyncioWaker::yield_waken(py), + State::Delegated(obj) => Ok(obj.clone_ref(py)), + } + } + + fn delegate(&self, py: Python<'_>, await_impl: PyObject) -> Poll> { + match delegate(py, await_impl, &self.op) { + Ok(YieldOrReturn::Yield(obj)) => { + let delegated = self.state.set(py, State::Delegated(obj)); + assert!(delegated.is_ok(), "{}", MIXED_AWAITABLE_AND_FUTURE_ERROR); + Poll::Pending + } + Ok(YieldOrReturn::Return(obj)) => Poll::Ready(Ok(obj)), + Err(err) if err.is_instance_of::(py) => Poll::Ready( + err.value_bound(py) + .getattr(intern!(py, "value")) + .map(Bound::unbind), + ), + Err(err) => Poll::Ready(Err(err)), + } } } -impl Wake for AsyncioWaker { +impl Wake for CoroutineWaker { fn wake(self: Arc) { self.wake_by_ref() } fn wake_by_ref(self: &Arc) { - Python::with_gil(|gil| { - if let Some(loop_and_future) = self.0.get_or_init(gil, || None) { - loop_and_future - .set_result(gil) - .expect("unexpected error in coroutine waker"); - } - }); + Python::with_gil(|gil| match WAKER_HACK.with(|cell| cell.take()) { + Some(WakerHack::Argument(await_impl)) => WAKER_HACK.with(|cell| { + let res = self.delegate(gil, await_impl); + cell.set(Some(WakerHack::Result(res))) + }), + Some(WakerHack::Result(_)) => unreachable!(), + None => match self.state.get_or_init(gil, || State::Waken) { + State::Pending(waker) => waker.wake(gil).expect("wake error"), + State::Waken => {} + State::Delegated(_) => panic!("{}", MIXED_AWAITABLE_AND_FUTURE_ERROR), + }, + }) } } -struct LoopAndFuture { - event_loop: PyObject, - future: PyObject, +enum WakerHack { + Argument(PyObject), + Result(Poll>), } -impl LoopAndFuture { - fn new(py: Python<'_>) -> PyResult { - static GET_RUNNING_LOOP: GILOnceCell = GILOnceCell::new(); - let import = || -> PyResult<_> { - let module = py.import_bound("asyncio")?; - Ok(module.getattr("get_running_loop")?.into()) - }; - let event_loop = GET_RUNNING_LOOP.get_or_try_init(py, import)?.call0(py)?; - let future = event_loop.call_method0(py, "create_future")?; - Ok(Self { event_loop, future }) - } - - fn set_result(&self, py: Python<'_>) -> PyResult<()> { - static RELEASE_WAITER: GILOnceCell> = GILOnceCell::new(); - let release_waiter = RELEASE_WAITER.get_or_try_init(py, || { - wrap_pyfunction_bound!(release_waiter, py).map(Bound::unbind) - })?; - // `Future.set_result` must be called in event loop thread, - // so it requires `call_soon_threadsafe` - let call_soon_threadsafe = self.event_loop.call_method1( - py, - intern!(py, "call_soon_threadsafe"), - (release_waiter, self.future.bind(py)), - ); - if let Err(err) = call_soon_threadsafe { - // `call_soon_threadsafe` will raise if the event loop is closed; - // instead of catching an unspecific `RuntimeError`, check directly if it's closed. - let is_closed = self.event_loop.call_method0(py, "is_closed")?; - if !is_closed.extract(py)? { - return Err(err); - } - } - Ok(()) - } +thread_local! { + static WAKER_HACK: Cell> = Cell::new(None); } -/// Call `future.set_result` if the future is not done. -/// -/// Future can be cancelled by the event loop before being waken. -/// See -#[pyfunction(crate = "crate")] -fn release_waiter(future: &Bound<'_, PyAny>) -> PyResult<()> { - let done = future.call_method0(intern!(future.py(), "done"))?; - if !done.extract::()? { - future.call_method1(intern!(future.py(), "set_result"), (future.py().None(),))?; +pub(crate) fn try_delegate( + waker: &Waker, + await_impl: PyObject, +) -> Option>> { + WAKER_HACK.with(|cell| cell.set(Some(WakerHack::Argument(await_impl)))); + waker.wake_by_ref(); + match WAKER_HACK.with(|cell| cell.take()) { + Some(WakerHack::Result(poll)) => Some(poll), + Some(WakerHack::Argument(_)) => None, + None => unreachable!(), } - Ok(()) } diff --git a/src/impl_/coroutine.rs b/src/impl_/coroutine.rs index da2066553df..830bcfa644c 100644 --- a/src/impl_/coroutine.rs +++ b/src/impl_/coroutine.rs @@ -4,7 +4,7 @@ use std::{ }; use crate::{ - coroutine::{cancel::ThrowCallback, Coroutine}, + coroutine::{Coroutine, ThrowCallback}, instance::Bound, marker::Ungil, pycell::impl_::PyClassBorrowChecker, diff --git a/src/lib.rs b/src/lib.rs index b400f143f5a..fa6e1b0f0ff 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -312,27 +312,29 @@ //! [Rust from Python]: https://github.com/PyO3/pyo3#using-rust-from-python //! [Features chapter of the guide]: https://pyo3.rs/latest/features.html#features-reference "Features Reference - PyO3 user guide" //! [`Ungil`]: crate::marker::Ungil -pub use crate::class::*; -pub use crate::conversion::{AsPyPointer, FromPyObject, IntoPy, ToPyObject}; #[allow(deprecated)] pub use crate::conversion::{FromPyPointer, PyTryFrom, PyTryInto}; -pub use crate::err::{ - DowncastError, DowncastIntoError, PyDowncastError, PyErr, PyErrArguments, PyResult, ToPyErr, -}; #[allow(deprecated)] pub use crate::gil::GILPool; #[cfg(not(any(PyPy, GraalPy)))] pub use crate::gil::{prepare_freethreaded_python, with_embedded_python_interpreter}; -pub use crate::instance::{Borrowed, Bound, Py, PyNativeType, PyObject}; -pub use crate::marker::Python; #[allow(deprecated)] pub use crate::pycell::PyCell; -pub use crate::pycell::{PyRef, PyRefMut}; -pub use crate::pyclass::PyClass; -pub use crate::pyclass_init::PyClassInitializer; -pub use crate::type_object::{PyTypeCheck, PyTypeInfo}; -pub use crate::types::PyAny; -pub use crate::version::PythonVersionInfo; +pub use crate::{ + class::*, + conversion::{AsPyPointer, FromPyObject, IntoPy, ToPyObject}, + err::{ + DowncastError, DowncastIntoError, PyDowncastError, PyErr, PyErrArguments, PyResult, ToPyErr, + }, + instance::{Borrowed, Bound, Py, PyNativeType, PyObject}, + marker::Python, + pycell::{PyRef, PyRefMut}, + pyclass::PyClass, + pyclass_init::PyClassInitializer, + type_object::{PyTypeCheck, PyTypeInfo}, + types::PyAny, + version::PythonVersionInfo, +}; pub(crate) mod ffi_ptr_ext; pub(crate) mod py_result_ext; @@ -347,7 +349,6 @@ pub(crate) mod sealed; /// once is resolved. pub mod class { pub use self::gc::{PyTraverseError, PyVisit}; - #[doc(hidden)] pub use self::methods::{ PyClassAttributeDef, PyGetterDef, PyMethodDef, PyMethodDefType, PyMethodType, PySetterDef, @@ -410,16 +411,15 @@ pub mod class { } } +#[cfg(all(feature = "macros", feature = "multiple-pymethods"))] +#[doc(hidden)] +pub use inventory; #[cfg(feature = "macros")] #[doc(hidden)] pub use { indoc, // Re-exported for py_run unindent, // Re-exported for py_run -}; - -#[cfg(all(feature = "macros", feature = "multiple-pymethods"))] -#[doc(hidden)] -pub use inventory; // Re-exported for `#[pyclass]` and `#[pymethods]` with `multiple-pymethods`. +}; // Re-exported for `#[pyclass]` and `#[pymethods]` with `multiple-pymethods`. /// Tests and helpers which reside inside PyO3's main library. Declared first so that macros /// are available in unit tests. @@ -462,14 +462,7 @@ pub mod type_object; pub mod types; mod version; -#[allow(unused_imports)] // with no features enabled this module has no public exports -pub use crate::conversions::*; - -#[cfg(feature = "macros")] -pub use pyo3_macros::{pyfunction, pymethods, pymodule, FromPyObject}; - /// A proc macro used to expose Rust structs and fieldless enums as Python objects. -/// #[doc = include_str!("../guide/pyclass-parameters.md")] /// /// For more on creating Python classes, @@ -478,6 +471,11 @@ pub use pyo3_macros::{pyfunction, pymethods, pymodule, FromPyObject}; /// [1]: https://pyo3.rs/latest/class.html #[cfg(feature = "macros")] pub use pyo3_macros::pyclass; +#[cfg(feature = "macros")] +pub use pyo3_macros::{pyfunction, pymethods, pymodule, FromPyObject}; + +#[allow(unused_imports)] // with no features enabled this module has no public exports +pub use crate::conversions::*; #[cfg(feature = "macros")] #[macro_use] @@ -512,6 +510,7 @@ pub mod doc_test { "README.md" => readme_md, "guide/src/advanced.md" => guide_advanced_md, "guide/src/async-await.md" => guide_async_await_md, + "guide/src/async-await/awaiting_python_awaitables.md" => guide_async_await_awaiting_python_awaitable_md, "guide/src/building-and-distribution.md" => guide_building_and_distribution_md, "guide/src/building-and-distribution/multiple-python-versions.md" => guide_bnd_multiple_python_versions_md, "guide/src/class.md" => guide_class_md, diff --git a/src/tests/common.rs b/src/tests/common.rs index 854d73e4d7b..78962d68cd0 100644 --- a/src/tests/common.rs +++ b/src/tests/common.rs @@ -6,13 +6,14 @@ #[macro_use] mod inner { + use pyo3::{ + prelude::*, + types::{IntoPyDict, PyList}, + }; + #[allow(unused_imports)] // pulls in `use crate as pyo3` in `test_utils.rs` use super::*; - use pyo3::prelude::*; - - use pyo3::types::{IntoPyDict, PyList}; - #[macro_export] macro_rules! py_assert { ($py:expr, $($val:ident)+, $assertion:literal) => { @@ -156,6 +157,17 @@ mod inner { .unwrap(); }}; } + + // see https://stackoverflow.com/questions/60359157/valueerror-set-wakeup-fd-only-works-in-main-thread-on-windows-on-python-3-8-wit + #[cfg(feature = "macros")] + pub fn asyncio_windows(test: &str) -> String { + let set_event_loop_policy = r#" + import asyncio, sys + if sys.platform == "win32": + asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) + "#; + pyo3::unindent::unindent(set_event_loop_policy) + &pyo3::unindent::unindent(test) + } } #[allow(unused_imports)] // some tests use just the macros and none of the other functionality diff --git a/tests/test_await_in_coroutine.rs b/tests/test_await_in_coroutine.rs new file mode 100644 index 00000000000..3f835c27588 --- /dev/null +++ b/tests/test_await_in_coroutine.rs @@ -0,0 +1,177 @@ +#![cfg(feature = "experimental-async")] + +use std::task::Poll; + +use futures::{future::poll_fn, FutureExt}; +use pyo3::{ + coroutine::{await_in_coroutine, CancelHandle}, + exceptions::{PyAttributeError, PyTypeError}, + prelude::*, + py_run, +}; + +#[path = "../src/tests/common.rs"] +mod common; + +#[pyfunction] +async fn wrap_awaitable(awaitable: PyObject) -> PyResult { + let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?; + future.await +} + +#[test] +fn awaitable() { + Python::with_gil(|gil| { + let wrap_awaitable = wrap_pyfunction_bound!(wrap_awaitable, gil).unwrap(); + let test = r#" + import types + import asyncio; + + class BadAwaitable: + def __await__(self): + raise AttributeError("__await__") + + @types.coroutine + def gen_coro(): + yield None + + async def main(): + await wrap_awaitable(...) + asyncio.run(main()) + "#; + let globals = gil.import_bound("__main__").unwrap().dict(); + globals.set_item("wrap_awaitable", wrap_awaitable).unwrap(); + let run = |awaitable| { + gil.run_bound( + &common::asyncio_windows(test).replace("...", awaitable), + Some(&globals), + None, + ) + }; + run("asyncio.sleep(0.001)").unwrap(); + run("gen_coro()").unwrap(); + assert!(run("None").unwrap_err().is_instance_of::(gil)); + assert!(run("BadAwaitable()") + .unwrap_err() + .is_instance_of::(gil)); + }) +} + +#[test] +fn cancel_delegation() { + #[pyfunction] + async fn wrap_cancellable(awaitable: PyObject, #[pyo3(cancel_handle)] cancel: CancelHandle) { + let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil))).unwrap(); + let result = future.await; + Python::with_gil(|gil| { + assert_eq!( + result.unwrap_err().get_type_bound(gil).name().unwrap(), + "CancelledError" + ) + }); + assert!(!cancel.is_cancelled()); + } + Python::with_gil(|gil| { + let wrap_cancellable = wrap_pyfunction_bound!(wrap_cancellable, gil).unwrap(); + let test = r#" + import asyncio; + + async def main(): + task = asyncio.create_task(wrap_cancellable(asyncio.sleep(0.001))) + await asyncio.sleep(0) + task.cancel() + await task + asyncio.run(main()) + "#; + let globals = gil.import_bound("__main__").unwrap().dict(); + globals + .set_item("wrap_cancellable", wrap_cancellable) + .unwrap(); + gil.run_bound(&common::asyncio_windows(test), Some(&globals), None) + .unwrap(); + }) +} + +#[test] +#[should_panic(expected = "PyFuture must be awaited in coroutine context")] +fn pyfuture_without_coroutine() { + #[pyfunction] + fn block_on(awaitable: PyObject) -> PyResult { + let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?; + futures::executor::block_on(future) + } + Python::with_gil(|gil| { + let block_on = wrap_pyfunction_bound!(block_on, gil).unwrap(); + let test = r#" + async def coro(): + ... + block_on(coro()) + "#; + py_run!(gil, block_on, &common::asyncio_windows(test)); + }) +} + +async fn checkpoint() { + let mut ready = false; + poll_fn(|cx| { + if ready { + return Poll::Ready(()); + } + ready = true; + cx.waker().wake_by_ref(); + Poll::Pending + }) + .await +} + +#[test] +#[should_panic(expected = "Python awaitable mixed with Rust future")] +fn pyfuture_in_select() { + #[pyfunction] + async fn select(awaitable: PyObject) -> PyResult { + let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?; + futures::select_biased! { + _ = checkpoint().fuse() => unreachable!(), + res = future.fuse() => res, + } + } + Python::with_gil(|gil| { + let select = wrap_pyfunction_bound!(select, gil).unwrap(); + let test = r#" + import asyncio; + async def main(): + return await select(asyncio.sleep(1)) + asyncio.run(main()) + "#; + let globals = gil.import_bound("__main__").unwrap().dict(); + globals.set_item("select", select).unwrap(); + gil.run_bound(&common::asyncio_windows(test), Some(&globals), None) + .unwrap(); + }) +} + +#[test] +#[should_panic(expected = "Python awaitable mixed with Rust future")] +fn pyfuture_in_select2() { + #[pyfunction] + async fn select2(awaitable: PyObject) -> PyResult { + let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?; + futures::select_biased! { + res = future.fuse() => res, + _ = checkpoint().fuse() => unreachable!(), + } + } + Python::with_gil(|gil| { + let select2 = wrap_pyfunction_bound!(select2, gil).unwrap(); + let test = r#" + import asyncio; + async def main(): + return await select2(asyncio.sleep(1)) + asyncio.run(main()) + "#; + let globals = gil.import_bound("__main__").unwrap().dict(); + globals.set_item("select2", select2).unwrap(); + gil.run_bound(&common::asyncio_windows(test), Some(&globals), None) + .unwrap(); + }) +} diff --git a/tests/test_coroutine.rs b/tests/test_coroutine.rs index 75b524edf78..4ab13733ecf 100644 --- a/tests/test_coroutine.rs +++ b/tests/test_coroutine.rs @@ -17,15 +17,6 @@ use std::sync::atomic::{AtomicBool, Ordering}; #[path = "../src/tests/common.rs"] mod common; -fn handle_windows(test: &str) -> String { - let set_event_loop_policy = r#" - import asyncio, sys - if sys.platform == "win32": - asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) - "#; - pyo3::unindent::unindent(set_event_loop_policy) + &pyo3::unindent::unindent(test) -} - #[test] fn noop_coroutine() { #[pyfunction] @@ -35,7 +26,7 @@ fn noop_coroutine() { Python::with_gil(|gil| { let noop = wrap_pyfunction_bound!(noop, gil).unwrap(); let test = "import asyncio; assert asyncio.run(noop()) == 42"; - py_run!(gil, noop, &handle_windows(test)); + py_run!(gil, noop, &common::asyncio_windows(test)); }) } @@ -79,7 +70,7 @@ fn test_coroutine_qualname() { ("MyClass", gil.get_type_bound::().as_any()), ] .into_py_dict_bound(gil); - py_run!(gil, *locals, &handle_windows(test)); + py_run!(gil, *locals, &common::asyncio_windows(test)); }) } @@ -101,7 +92,7 @@ fn sleep_0_like_coroutine() { Python::with_gil(|gil| { let sleep_0 = wrap_pyfunction_bound!(sleep_0, gil).unwrap(); let test = "import asyncio; assert asyncio.run(sleep_0()) == 42"; - py_run!(gil, sleep_0, &handle_windows(test)); + py_run!(gil, sleep_0, &common::asyncio_windows(test)); }) } @@ -120,7 +111,7 @@ fn sleep_coroutine() { Python::with_gil(|gil| { let sleep = wrap_pyfunction_bound!(sleep, gil).unwrap(); let test = r#"import asyncio; assert asyncio.run(sleep(0.1)) == 42"#; - py_run!(gil, sleep, &handle_windows(test)); + py_run!(gil, sleep, &common::asyncio_windows(test)); }) } @@ -140,11 +131,7 @@ fn cancelled_coroutine() { let globals = gil.import_bound("__main__").unwrap().dict(); globals.set_item("sleep", sleep).unwrap(); let err = gil - .run_bound( - &pyo3::unindent::unindent(&handle_windows(test)), - Some(&globals), - None, - ) + .run_bound(&common::asyncio_windows(test), Some(&globals), None) .unwrap_err(); assert_eq!( err.value_bound(gil).get_type().qualname().unwrap(), @@ -180,12 +167,8 @@ fn coroutine_cancel_handle() { globals .set_item("cancellable_sleep", cancellable_sleep) .unwrap(); - gil.run_bound( - &pyo3::unindent::unindent(&handle_windows(test)), - Some(&globals), - None, - ) - .unwrap(); + gil.run_bound(&common::asyncio_windows(test), Some(&globals), None) + .unwrap(); }) } @@ -210,12 +193,8 @@ fn coroutine_is_cancelled() { "#; let globals = gil.import_bound("__main__").unwrap().dict(); globals.set_item("sleep_loop", sleep_loop).unwrap(); - gil.run_bound( - &pyo3::unindent::unindent(&handle_windows(test)), - Some(&globals), - None, - ) - .unwrap(); + gil.run_bound(&common::asyncio_windows(test), Some(&globals), None) + .unwrap(); }) } @@ -244,7 +223,7 @@ fn coroutine_panic() { else: assert False; "#; - py_run!(gil, panic, &handle_windows(test)); + py_run!(gil, panic, &common::asyncio_windows(test)); }) } @@ -341,6 +320,6 @@ fn test_async_method_receiver_with_other_args() { assert asyncio.run(v.get_value_plus_with(1, 1)) == 12 "#; let locals = [("Value", gil.get_type_bound::())].into_py_dict_bound(gil); - py_run!(gil, *locals, test); + py_run!(gil, *locals, &common::asyncio_windows(test)); }); }