diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cb007b2f..91f88646 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -314,6 +314,14 @@ jobs: - name: Rust unit tests if: matrix.build_type.name == 'standard' && matrix.platform.target == 'x64' run: cargo test --lib --locked + - name: Event loop tests + if: matrix.build_type.name == 'standard' && matrix.platform.target == 'x64' + timeout-minutes: 10 + shell: bash + run: | + python -m pip install --no-index --find-links dist wreq + python -m pip install pytest pytest-asyncio + python -m pytest -v -o faulthandler_timeout=60 tests/aio_test.py tests/cancellation_test.py tests/upload_test.py tests/runtime_test.py - name: Upload wheels uses: actions/upload-artifact@v7 with: @@ -353,6 +361,13 @@ jobs: target: ${{ matrix.platform.target }} args: --release --out dist ${{ matrix.build_type.maturin_args }} --features ${{ matrix.platform.allocator || '' }} sccache: "false" + - name: Event loop tests + if: matrix.build_type.name == 'standard' && matrix.platform.target == 'aarch64' + timeout-minutes: 10 + run: | + python -m pip install --no-index --find-links dist wreq + python -m pip install pytest pytest-asyncio uvloop + python -m pytest -v -o faulthandler_timeout=60 tests/aio_test.py tests/cancellation_test.py tests/upload_test.py tests/runtime_test.py - name: Upload wheels uses: actions/upload-artifact@v7 with: diff --git a/Cargo.lock b/Cargo.lock index dd1596aa..1a219e01 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1175,6 +1175,26 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" +[[package]] +name = "netty" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6d4d60c6b6193ac462e05a1c8ae826a624eff20368e208e59101f5a1fbfc055" +dependencies = [ + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "http2", + "httparse", + "pin-project-lite", + "smallvec", + "tokio", + "tokio-util", + "want", +] + [[package]] name = "nom" version = "7.1.3" @@ -2412,7 +2432,7 @@ checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" [[package]] name = "wreq" version = "0.16.1" -source = "git+https://github.com/0x676e67/wreq?rev=4f2342aba1b33cc3534e8e2a853b21945edf6cef#4f2342aba1b33cc3534e8e2a853b21945edf6cef" +source = "git+https://github.com/0x676e67/wreq?rev=54fdc7c9e0c1290e2f70f38d271ba954e3d86cc2#54fdc7c9e0c1290e2f70f38d271ba954e3d86cc2" dependencies = [ "btls", "btls-sys", @@ -2430,6 +2450,7 @@ dependencies = [ "lru", "mime", "mime_guess", + "netty", "parking_lot", "percent-encoding", "pin-project-lite", @@ -2447,30 +2468,9 @@ dependencies = [ "tower-http", "url", "webpki-root-certs", - "wreq-proto", "wreq-rt", ] -[[package]] -name = "wreq-proto" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb4d829f0427a5064829976fe790641ca21abb1d83d844812baf022dc0bdf2db" -dependencies = [ - "bytes", - "futures-channel", - "futures-util", - "http", - "http-body", - "http2", - "httparse", - "pin-project-lite", - "smallvec", - "tokio", - "tokio-util", - "want", -] - [[package]] name = "wreq-python" version = "0.12.3" @@ -2499,12 +2499,11 @@ dependencies = [ [[package]] name = "wreq-rt" version = "0.2.2-rc.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99e9bce67a3fa3dd3f1503f066d86661c9caf399a763d3bd184da7afaf886c8b" +source = "git+https://github.com/0x676e67/wreq?rev=54fdc7c9e0c1290e2f70f38d271ba954e3d86cc2#54fdc7c9e0c1290e2f70f38d271ba954e3d86cc2" dependencies = [ + "netty", "pin-project-lite", "tokio", - "wreq-proto", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index d8acc3a6..0c9cf0d5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -78,7 +78,8 @@ tikv-jemallocator = { version = "0.7.0", default-features = false, features = [ # Keep the latest emulation profiles and their required TLS support in sync. [patch.crates-io] wreq-util = { git = "https://github.com/0x676e67/wreq-util", rev = "5715529002738f46eea86a7271f918f9374d3582" } -wreq = { git = "https://github.com/0x676e67/wreq", rev = "4f2342aba1b33cc3534e8e2a853b21945edf6cef" } +# demo/netty-migration +wreq = { git = "https://github.com/0x676e67/wreq", rev = "54fdc7c9e0c1290e2f70f38d271ba954e3d86cc2" } # demo/tokio-btls-vectored-write btls = { git = "https://github.com/0x676e67/btls", rev = "ffaf0e8c23d455af4e306986ce81bb57ebd03bc2" } btls-sys = { git = "https://github.com/0x676e67/btls", rev = "ffaf0e8c23d455af4e306986ce81bb57ebd03bc2" } diff --git a/pyproject.toml b/pyproject.toml index e72ad38f..301b556c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ test = [ "pytest", "pytest-asyncio", "pytest-rerunfailures", + "uvloop; sys_platform != 'win32' and platform_python_implementation == 'CPython'", ] docs = [ "zensical", diff --git a/python/wreq/blocking.py b/python/wreq/blocking.py index 391db5cb..c5770812 100644 --- a/python/wreq/blocking.py +++ b/python/wreq/blocking.py @@ -119,7 +119,7 @@ def close(self) -> None: This does not guarantee an immediate socket shutdown or interrupt an active read. Do not close concurrently with a body read. A body transferred to a Streamer is managed separately; previously returned memoryviews remain valid. - Prefer a context manager (`with`) for response cleanup. + `with` instead releases the body and keeps a fully read connection reusable. """ def __enter__(self) -> Any: ... diff --git a/python/wreq/wreq.py b/python/wreq/wreq.py index 43d81741..dc7ab0ce 100644 --- a/python/wreq/wreq.py +++ b/python/wreq/wreq.py @@ -445,7 +445,7 @@ async def close(self) -> None: This does not guarantee an immediate socket shutdown or cancel an active read. Cancel and await any body-read task before closing. A body transferred to a Streamer is managed separately; previously returned memoryviews remain valid. - Prefer an async context manager (`async with`) for response cleanup. + `async with` instead releases the body and keeps a fully read connection reusable. """ async def __aenter__(self) -> Any: ... diff --git a/src/aio.rs b/src/aio.rs new file mode 100644 index 00000000..1b1ef28c --- /dev/null +++ b/src/aio.rs @@ -0,0 +1,99 @@ +//! Awaitables that resume asyncio tasks without attaching Tokio threads to Python. +//! +//! Rust futures complete on Tokio workers, which never take the GIL. A wake marks +//! its [`Coroutine`] ready and queues it on the [`Port`] of the awaiting event loop, +//! ringing a socket the loop polls only when the queue was empty. The loop thread +//! then resolves every queued asyncio future in one batch. + +mod coroutine; +mod port; + +use std::{ + future::{Future, poll_fn}, + task::Poll, +}; + +use pyo3::{IntoPyObjectExt, exceptions::PyRuntimeError, prelude::*}; +use tokio_util::task::AbortOnDropHandle; + +pub use self::coroutine::Coroutine; +use self::port::Port; +use crate::runtime::Runtime; + +/// Run `fut` on the runtime once the coroutine named `qualname` is first awaited. +#[inline] +pub fn spawn<'py, F, T>( + py: Python<'py>, + qualname: &'static str, + runtime: &Runtime, + fut: F, +) -> PyResult> +where + F: Future> + Send + 'static, + T: for<'a> IntoPyObject<'a> + Send + 'static, +{ + local(py, qualname, run(runtime.clone(), fut)) +} + +/// Poll `fut` on the event loop thread, attached to Python, when awaited. +#[inline] +pub fn local<'py, F, T>( + py: Python<'py>, + qualname: &'static str, + fut: F, +) -> PyResult> +where + F: Future> + Send + 'static, + T: for<'a> IntoPyObject<'a>, +{ + Bound::new( + py, + Coroutine::new(qualname, async move { + let value = fut.await?; + Python::attach(|py| value.into_py_any(py)) + }), + ) +} + +/// A coroutine that returns `value` without suspending. +#[inline] +pub fn ready<'py, T>( + qualname: &'static str, + value: Bound<'py, T>, +) -> PyResult> { + let py = value.py(); + let value = value.into_any().unbind(); + Bound::new(py, Coroutine::new(qualname, async move { Ok(value) })) +} + +/// Yield to the event loop once, like `asyncio.sleep(0)`. +#[inline] +pub async fn yield_now() { + let mut yielded = false; + poll_fn(|cx| { + if std::mem::replace(&mut yielded, true) { + return Poll::Ready(()); + } + // A wake during the poll makes the coroutine yield without a future. + cx.waker().wake_by_ref(); + Poll::Pending + }) + .await; +} + +/// Spawn `fut` on the runtime and wait for it; dropping the wait aborts the task. +/// The task keeps the runtime alive until it ends. +#[inline] +pub async fn run(runtime: Runtime, fut: F) -> PyResult +where + F: Future> + Send + 'static, + T: Send + 'static, +{ + let task = runtime.handle().clone().spawn(async move { + let _owner = runtime; + fut.await + }); + AbortOnDropHandle::new(task) + .await + .map_err(|err| PyRuntimeError::new_err(err.to_string()))? +} diff --git a/src/aio/coroutine.rs b/src/aio/coroutine.rs new file mode 100644 index 00000000..cf0ba8e6 --- /dev/null +++ b/src/aio/coroutine.rs @@ -0,0 +1,240 @@ +use std::{ + future::Future, + pin::Pin, + sync::{ + Arc, Mutex, MutexGuard, PoisonError, + atomic::{AtomicU8, Ordering}, + }, + task::{Context, Poll, Wake, Waker}, +}; + +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyStopIteration}, + intern, + prelude::*, +}; + +use super::Port; + +type BoxFuture = Pin>> + Send>>; + +/// An awaitable driving a Rust future on the asyncio event loop thread. +/// +/// Each suspension hands the task a real asyncio future, so task cancellation and +/// the C task fast path work unchanged. A throw or close drops the Rust future, +/// which aborts any spawned work. PyO3's borrow flag rejects reentrant polls. +#[pyclass(module = "wreq")] +pub struct Coroutine { + qualname: &'static str, + /// Polled only through `&mut self`, so `get_mut` reaches it without locking. + future: Mutex>, + slot: Arc, + waker: Waker, +} + +/// Wake state shared with Rust wakers for the coroutine's whole life. +/// +/// Reusing one waker lets Tokio skip re-registering it on every poll. Wakes only +/// change the state and queue the slot; Python objects are touched on the loop thread. +#[derive(Default)] +pub(super) struct Slot { + state: AtomicU8, + /// The port of the loop the coroutine last waited on, set before each wait. + port: Mutex>>, + /// The asyncio future the task awaits while the coroutine waits. + waiter: Mutex>>, +} + +const POLLING: u8 = 0; +const WAITING: u8 = 1; +const NOTIFIED: u8 = 2; +const DONE: u8 = 3; + +// ===== impl Coroutine ===== + +impl Coroutine { + pub(super) fn new(qualname: &'static str, future: F) -> Self + where + F: Future>> + Send + 'static, + { + let slot = Arc::new(Slot::default()); + Coroutine { + qualname, + future: Mutex::new(Some(Box::pin(future))), + waker: Waker::from(slot.clone()), + slot, + } + } + + fn step(&mut self, py: Python<'_>) -> PyResult> { + let future = self + .future + .get_mut() + .unwrap_or_else(PoisonError::into_inner) + .as_mut() + .ok_or_else(|| PyRuntimeError::new_err("cannot reuse already awaited coroutine"))?; + + // A pending waiter means another task is suspended on this coroutine. Once + // it is taken, wakes from now on resume this poll. + if let Some(waiter) = self.slot.take_waiter() { + let done = waiter + .bind(py) + .call_method0(intern!(py, "done")) + .and_then(|done| done.is_truthy()); + if !matches!(done, Ok(true)) { + self.slot.set_waiter(waiter); + done?; + return Err(PyRuntimeError::new_err( + "coroutine is being awaited already", + )); + } + } + + self.slot.state.store(POLLING, Ordering::Release); + if let Poll::Ready(result) = future.as_mut().poll(&mut Context::from_waker(&self.waker)) { + self.finish(); + return Err(PyStopIteration::new_err((result?,))); + } + // A coroutine that cannot wait ends here, releasing any spawned work. + self.slot.suspend(py).inspect_err(|_| self.finish()) + } + + fn finish(&mut self) { + *self + .future + .get_mut() + .unwrap_or_else(PoisonError::into_inner) = None; + self.slot.state.store(DONE, Ordering::Release); + // A finished coroutine must not keep its loop's port open. + let port = self.slot.lock_port().take(); + drop(port); + drop(self.slot.take_waiter()); + } +} + +#[pymethods] +impl Coroutine { + #[getter] + fn __name__(&self) -> &'static str { + self.qualname + .rsplit_once('.') + .map_or(self.qualname, |(_, name)| name) + } + + #[getter] + fn __qualname__(&self) -> &'static str { + self.qualname + } + + fn send(&mut self, py: Python<'_>, _value: &Bound<'_, PyAny>) -> PyResult> { + self.step(py) + } + + fn throw(&mut self, exc: Bound<'_, PyAny>) -> PyResult> { + self.finish(); + Err(PyErr::from_value(exc)) + } + + fn close(&mut self) { + self.finish(); + } + + fn __await__(slf: Py) -> Py { + slf + } + + fn __next__(&mut self, py: Python<'_>) -> PyResult> { + self.step(py) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + // Pending Rust work owns the waiter until its loop closes the port; only then + // can a cycle through the waiter be collected without aborting live requests. + let closed = self + .slot + .port + .try_lock() + .is_ok_and(|port| port.as_ref().is_some_and(|port| !port.is_open())); + if closed && let Ok(waiter) = self.slot.waiter.try_lock() { + visit.call(&*waiter)?; + } + Ok(()) + } + + fn __clear__(&mut self) { + self.finish(); + } +} + +// ===== impl Slot ===== + +impl Slot { + /// Suspend after a pending poll, returning the future the task should await, + /// or `None` for a bare yield when a wake already arrived. + fn suspend(&self, py: Python<'_>) -> PyResult> { + if self.state.load(Ordering::Acquire) == NOTIFIED { + return Ok(py.None()); + } + let (event_loop, port) = Port::current(py)?; + let waiter = event_loop.call_method0(intern!(py, "create_future"))?; + waiter.setattr(intern!(py, "_asyncio_future_blocking"), true)?; + let waiter = waiter.unbind(); + // Publish the port and waiter before waiting, so a wake that sees WAITING + // always reaches the loop that runs this task. + *self.lock_port() = Some(port); + self.set_waiter(waiter.clone_ref(py)); + match self + .state + .compare_exchange(POLLING, WAITING, Ordering::AcqRel, Ordering::Acquire) + { + Ok(_) => Ok(waiter), + // Woken after the poll but before waiting: resume on the next loop turn. + Err(_) => { + drop(self.take_waiter()); + Ok(py.None()) + } + } + } + + #[inline] + fn lock_port(&self) -> MutexGuard<'_, Option>> { + self.port.lock().unwrap_or_else(PoisonError::into_inner) + } + + #[inline] + pub(super) fn set_waiter(&self, waiter: Py) { + *self.waiter.lock().unwrap_or_else(PoisonError::into_inner) = Some(waiter); + } + + #[inline] + pub(super) fn take_waiter(&self) -> Option> { + self.waiter + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + } +} + +impl Wake for Slot { + #[inline] + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + let prev = self + .state + .try_update(Ordering::AcqRel, Ordering::Acquire, |state| { + matches!(state, POLLING | WAITING).then_some(NOTIFIED) + }); + // A wake during a poll is seen by `suspend`; only a waiting task needs the loop. + if prev == Ok(WAITING) { + // Release the lock before pushing, which may attach to Python. + let port = self.lock_port().clone(); + if let Some(port) = port { + port.push(self.clone()); + } + } + } +} diff --git a/src/aio/port.rs b/src/aio/port.rs new file mode 100644 index 00000000..a95ce0b5 --- /dev/null +++ b/src/aio/port.rs @@ -0,0 +1,324 @@ +use std::{ + cell::RefCell, + io::{Read, Write}, + sync::{ + Arc, Mutex, MutexGuard, PoisonError, Weak, + atomic::{AtomicBool, Ordering}, + }, +}; + +use pyo3::{intern, prelude::*, sync::PyOnceLock}; + +use super::coroutine::Slot; + +/// Wakes queued for one event loop, delivered by a single bell per batch. +pub(crate) struct Port { + queue: Mutex>>, + bell: Bell, + /// Cleared once the loop drops its drain; later wakes are discarded. + open: AtomicBool, +} + +enum Bell { + /// A socket pair whose read end the loop watches. The read end is `None` + /// when Python owns it, as with a proactor loop. + Socket { tx: Socket, rx: Option }, + /// A loop that cannot watch the socket: the waking thread schedules the drain. + Loop(Py), +} + +#[cfg(unix)] +type Socket = std::os::unix::net::UnixStream; + +#[cfg(windows)] +type Socket = std::net::TcpStream; + +/// Resolves the asyncio futures of queued coroutines on the loop thread. +#[pyclass(frozen)] +struct Drain(Arc); + +thread_local! { + /// Ports of the loops that ran on this thread, keyed by loop address. + static PORTS: RefCell)>> = const { RefCell::new(Vec::new()) }; +} + +// ===== impl Port ===== + +impl Port { + /// Return the running loop and its port, opening the port on first use. + pub(super) fn current(py: Python<'_>) -> PyResult<(Bound<'_, PyAny>, Arc)> { + static GET_RUNNING_LOOP: PyOnceLock> = PyOnceLock::new(); + let event_loop = GET_RUNNING_LOOP + .get_or_try_init(py, || { + py.import("asyncio")? + .getattr("get_running_loop") + .map(Bound::unbind) + })? + .bind(py) + .call0()?; + + // A loop drops its drain when closed or freed, so a reused address never matches. + let key = event_loop.as_ptr() as usize; + let cached = PORTS.with_borrow_mut(|ports| { + ports.retain(|(_, port)| port.upgrade().is_some_and(|port| port.is_open())); + ports + .iter() + .find(|(id, _)| *id == key) + .and_then(|(_, port)| port.upgrade()) + }); + let port = match cached { + Some(port) => port, + None => { + let port = Self::open(&event_loop)?; + PORTS.with_borrow_mut(|ports| ports.push((key, Arc::downgrade(&port)))); + port + } + }; + Ok((event_loop, port)) + } + + /// Open a port for `event_loop`, preferring a socket the loop watches itself. + fn open(event_loop: &Bound<'_, PyAny>) -> PyResult> { + let py = event_loop.py(); + if let Ok((tx, rx)) = socket_pair() { + #[cfg(windows)] + if event_loop.hasattr(intern!(py, "_proactor"))? { + use std::os::windows::io::IntoRawSocket; + + let port = Self::new(Bell::Socket { tx, rx: None }); + let drain = Drain(port.clone()); + let watched = proactor_watch(py) + .and_then(|watch| watch.call1((event_loop, rx.into_raw_socket(), drain))); + if watched.is_ok() { + return Ok(port); + } + return Ok(Self::new(Bell::Loop(event_loop.clone().unbind()))); + } + + let fd = raw(&rx); + let port = Self::new(Bell::Socket { tx, rx: Some(rx) }); + let drain = Drain(port.clone()); + // Loops without reader support fall back to scheduling drains. + if event_loop + .call_method1(intern!(py, "add_reader"), (fd, drain)) + .is_ok() + { + return Ok(port); + } + } + Ok(Self::new(Bell::Loop(event_loop.clone().unbind()))) + } + + fn new(bell: Bell) -> Arc { + Arc::new(Port { + queue: Mutex::default(), + bell, + open: AtomicBool::new(true), + }) + } + + #[inline] + pub(super) fn is_open(&self) -> bool { + self.open.load(Ordering::Acquire) + } + + #[inline] + fn lock(&self) -> MutexGuard<'_, Vec>> { + self.queue.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// Queue a woken slot, ringing the bell when the queue was empty. + pub(super) fn push(self: &Arc, slot: Arc) { + let first = { + let mut queue = self.lock(); + if !self.is_open() { + return; + } + queue.push(slot); + queue.len() == 1 + }; + if first { + self.ring(); + } + } + + fn ring(self: &Arc) { + match &self.bell { + // A full socket already holds an unread ring. A lost ring would leave the + // queue non-empty, so later pushes never ring: retry interrupted writes. + Bell::Socket { tx, .. } => loop { + match (&*tx).write(&[1]) { + Err(err) if err.kind() == std::io::ErrorKind::Interrupted => continue, + _ => break, + } + }, + Bell::Loop(event_loop) => { + Python::try_attach(|py| { + let drain = Drain(self.clone()); + let scheduled = event_loop + .bind(py) + .call_method1(intern!(py, "call_soon_threadsafe"), (drain,)); + // A closed loop never runs the drain. + if scheduled.is_err() { + self.close(); + } + }); + } + } + } + + /// Discard queued wakes, breaking the cycle between queued slots and the port. + fn close(&self) { + let queue = { + let mut queue = self.lock(); + self.open.store(false, Ordering::Release); + std::mem::take(&mut *queue) + }; + drop(queue); + } +} + +// ===== impl Drain ===== + +#[pymethods] +impl Drain { + fn __call__(&self, py: Python<'_>) -> PyResult<()> { + let port = &self.0; + // Clear the bell before taking the queue so no ring is lost; a short read + // means the bell is empty. + if let Bell::Socket { rx: Some(rx), .. } = &port.bell { + let mut buf = [0; 64]; + while matches!((&*rx).read(&mut buf), Ok(n) if n == buf.len()) {} + } + + let mut slots = std::mem::take(&mut *port.lock()).into_iter(); + while let Some(slot) = slots.next() { + let Some(waiter) = slot.take_waiter() else { + continue; + }; + // Resume the task awaiting the waiter, unless it was already cancelled. + let future = waiter.bind(py); + let done = future + .call_method0(intern!(py, "done")) + .and_then(|done| done.is_truthy()); + let released = match done { + Ok(false) => future + .call_method1(intern!(py, "set_result"), (py.None(),)) + .map(drop), + done => done.map(drop), + }; + if let Err(err) = released { + // A raising callback, such as a signal handler, ends this drain. Requeue + // this wake and the rest; the next drain skips futures already done. + slot.set_waiter(waiter); + port.lock().splice(..0, std::iter::once(slot).chain(slots)); + port.ring(); + return Err(err); + } + } + Ok(()) + } +} + +impl Drop for Drain { + fn drop(&mut self) { + // A socket port's drain is registered with the loop and owns the port's lifetime. + if let Bell::Socket { .. } = self.0.bell { + self.0.close(); + } + } +} + +#[cfg(unix)] +fn socket_pair() -> std::io::Result<(Socket, Socket)> { + let (tx, rx) = Socket::pair()?; + tx.set_nonblocking(true)?; + rx.set_nonblocking(true)?; + Ok((tx, rx)) +} + +/// Connect a loopback pair, accepting only the connection made here. +#[cfg(windows)] +fn socket_pair() -> std::io::Result<(Socket, Socket)> { + use std::net::{Ipv4Addr, TcpListener}; + + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0))?; + let tx = Socket::connect(listener.local_addr()?)?; + let local = tx.local_addr()?; + let rx = loop { + let (rx, peer) = listener.accept()?; + if peer == local { + break rx; + } + }; + tx.set_nodelay(true)?; + tx.set_nonblocking(true)?; + rx.set_nonblocking(true)?; + Ok((tx, rx)) +} + +#[cfg(unix)] +fn raw(socket: &Socket) -> i64 { + use std::os::fd::AsRawFd; + i64::from(socket.as_raw_fd()) +} + +#[cfg(windows)] +fn raw(socket: &Socket) -> std::os::windows::io::RawSocket { + use std::os::windows::io::AsRawSocket; + socket.as_raw_socket() +} + +/// Watch a socket with overlapped receives on a proactor loop, like its own +/// self-pipe; the receive is re-armed before each drain runs. +#[cfg(windows)] +fn proactor_watch(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> { + static WATCH: PyOnceLock> = PyOnceLock::new(); + WATCH + .get_or_try_init(py, || { + PyModule::from_code( + py, + c"import socket + + +class Watch: + # Pending receives reference the watch only through bound methods, so a closed + # loop frees it, and with it the drain, without waiting for the cycle collector. + + def __init__(self, loop, sock, drain): + self.loop = loop + self.sock = sock + self.drain = drain + + def arm(self): + self.loop._proactor.recv(self.sock, 4096).add_done_callback(self.ready) + + def ready(self, fut): + if fut.cancelled() or fut.exception() or not fut.result(): + self.sock.close() + return + try: + self.arm() + except Exception: + self.sock.close() + return + self.drain() + + +def watch(loop, fileno, drain): + sock = socket.socket(fileno=fileno) + try: + # A first receive that fails lets the caller fall back to scheduled drains. + Watch(loop, sock, drain).arm() + except BaseException: + sock.close() + raise +", + c"wreq/_proactor.py", + c"wreq._proactor", + )? + .getattr("watch") + .map(Bound::unbind) + }) + .map(|watch| watch.bind(py)) +} diff --git a/src/client.rs b/src/client.rs index c545dda4..f02b6408 100644 --- a/src/client.rs +++ b/src/client.rs @@ -11,17 +11,17 @@ use std::{ time::Duration, }; -use pyo3::{IntoPyObjectExt, coroutine::CancelHandle, prelude::*, pybacked::PyBackedStr}; +use pyo3::{IntoPyObjectExt, prelude::*, pybacked::PyBackedStr, types::PyDict}; use req::{Request, WebSocketRequest}; use tokio_util::sync::CancellationToken; use wreq::tls::trust::CertStore; use self::{ - nogil::NoGIL, req::{execute_request, execute_websocket_request}, resp::{BlockingResponse, BlockingWebSocket, Response, WebSocket}, }; use crate::{ + aio::{self, Coroutine}, cookie::Jar, dns::{DnsOptions, HickoryResolver, LookupIpStrategy}, emulate::EmulationLike, @@ -251,6 +251,51 @@ pub struct BlockingClient(Client); // ====== Client ===== +impl Client { + /// Return a coroutine named `qualname` that sends the request when awaited. + fn send<'py>( + &self, + py: Python<'py>, + qualname: &'static str, + method: Method, + url: PyBackedStr, + kwds: Option>, + ) -> PyResult> { + let kwds = kwds.map(Bound::unbind); + aio::local(py, qualname, self.clone().execute(method, url, kwds)) + } + + /// Send a request on the client's runtime, extracting options on first await so an + /// async generator body binds to the running loop. + pub(crate) async fn execute( + self, + method: Method, + url: PyBackedStr, + kwds: Option>, + ) -> PyResult { + let kwds = Python::attach(|py| kwds.map(|kwds| kwds.bind(py).extract()).transpose())?; + aio::run( + self.runtime.clone(), + execute_request(self, method, url, kwds), + ) + .await + } + + /// Open a WebSocket on the client's runtime, extracting options on first await. + pub(crate) async fn connect( + self, + url: PyBackedStr, + kwds: Option>, + ) -> PyResult { + let kwds = Python::attach(|py| kwds.map(|kwds| kwds.bind(py).extract()).transpose())?; + aio::run( + self.runtime.clone(), + execute_websocket_request(self, url, kwds), + ) + .await + } +} + impl Default for Client { fn default() -> Self { let runtime = runtime::get(); @@ -496,7 +541,6 @@ impl Client { /// Cancel pending requests and reject new ones with asyncio.CancelledError. /// Existing responses, WebSockets and the shared runtime remain usable. - #[inline] pub fn close(&self) { self.cancel.cancel(); } @@ -508,147 +552,136 @@ impl Client { } /// Make a GET request to the given URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] - pub async fn get( + pub fn get<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - self.request(cancel, Method::GET, url, kwds).await + kwds: Option>, + ) -> PyResult> { + self.send(py, "Client.get", Method::GET, url, kwds) } /// Make a HEAD request to the given URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] - pub async fn head( + pub fn head<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - self.request(cancel, Method::HEAD, url, kwds).await + kwds: Option>, + ) -> PyResult> { + self.send(py, "Client.head", Method::HEAD, url, kwds) } /// Make a POST request to the given URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] - pub async fn post( + pub fn post<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - self.request(cancel, Method::POST, url, kwds).await + kwds: Option>, + ) -> PyResult> { + self.send(py, "Client.post", Method::POST, url, kwds) } /// Make a PUT request to the given URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] - pub async fn put( + pub fn put<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - self.request(cancel, Method::PUT, url, kwds).await + kwds: Option>, + ) -> PyResult> { + self.send(py, "Client.put", Method::PUT, url, kwds) } /// Make a DELETE request to the given URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] - pub async fn delete( + pub fn delete<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - self.request(cancel, Method::DELETE, url, kwds).await + kwds: Option>, + ) -> PyResult> { + self.send(py, "Client.delete", Method::DELETE, url, kwds) } /// Make a PATCH request to the given URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] - pub async fn patch( + pub fn patch<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - self.request(cancel, Method::PATCH, url, kwds).await + kwds: Option>, + ) -> PyResult> { + self.send(py, "Client.patch", Method::PATCH, url, kwds) } /// Make a OPTIONS request to the given URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] - pub async fn options( + pub fn options<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - self.request(cancel, Method::OPTIONS, url, kwds).await + kwds: Option>, + ) -> PyResult> { + self.send(py, "Client.options", Method::OPTIONS, url, kwds) } /// Make a TRACE request to the given URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] - pub async fn trace( + pub fn trace<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - self.request(cancel, Method::TRACE, url, kwds).await + kwds: Option>, + ) -> PyResult> { + self.send(py, "Client.trace", Method::TRACE, url, kwds) } /// Make a request with the given method and URL. - #[inline] #[pyo3(signature = (method, url, **kwds))] - pub async fn request( + pub fn request<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, method: Method, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - NoGIL::with_cancel( - &self.runtime, - execute_request(self.clone(), method, url, kwds), - cancel, - ) - .await + kwds: Option>, + ) -> PyResult> { + self.send(py, "Client.request", method, url, kwds) } /// Make a WebSocket request to the given URL. - #[inline] #[pyo3(signature = (url, **kwds))] - pub async fn websocket( + pub fn websocket<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - NoGIL::with_cancel( - &self.runtime, - execute_websocket_request(self.clone(), url, kwds), - cancel, - ) - .await + kwds: Option>, + ) -> PyResult> { + let kwds = kwds.map(Bound::unbind); + aio::local(py, "Client.websocket", self.clone().connect(url, kwds)) } } #[pymethods] impl Client { - #[inline] - async fn __aenter__(slf: Py) -> PyResult> { - Ok(slf) + fn __aenter__(slf: Bound<'_, Self>) -> PyResult> { + aio::ready("Client.__aenter__", slf) } - #[inline] - async fn __aexit__(&self, _exc_type: Py, _exc_val: Py, _traceback: Py) { - self.close(); + fn __aexit__<'py>( + &self, + py: Python<'py>, + _exc_type: Py, + _exc_val: Py, + _traceback: Py, + ) -> PyResult> { + let cancel = self.cancel.clone(); + aio::local(py, "Client.__aexit__", async move { + cancel.cancel(); + Ok(()) + }) } } @@ -664,14 +697,12 @@ impl BlockingClient { /// Creates a new blocking Client instance. #[new] - #[inline] #[pyo3(signature = (**kwds))] fn new(py: Python, kwds: Option) -> PyResult { Client::new(py, kwds).map(BlockingClient) } /// Get the cookie jar of the client. - #[inline] #[getter] pub fn cookie_jar(&self) -> Option { self.0.cookie_jar.clone() @@ -679,13 +710,11 @@ impl BlockingClient { /// Cancel pending requests and reject new ones with asyncio.CancelledError. /// Existing responses, WebSockets and the shared runtime remain usable. - #[inline] pub fn close(&self) { self.0.close(); } /// Make a GET request to the specified URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] pub fn get( &self, @@ -697,7 +726,6 @@ impl BlockingClient { } /// Make a POST request to the specified URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] pub fn post( &self, @@ -709,7 +737,6 @@ impl BlockingClient { } /// Make a PUT request to the specified URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] pub fn put( &self, @@ -721,7 +748,6 @@ impl BlockingClient { } /// Make a PATCH request to the specified URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] pub fn patch( &self, @@ -733,7 +759,6 @@ impl BlockingClient { } /// Make a DELETE request to the specified URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] pub fn delete( &self, @@ -745,7 +770,6 @@ impl BlockingClient { } /// Make a HEAD request to the specified URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] pub fn head( &self, @@ -757,7 +781,6 @@ impl BlockingClient { } /// Make a OPTIONS request to the specified URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] pub fn options( &self, @@ -769,7 +792,6 @@ impl BlockingClient { } /// Make a TRACE request to the specified URL. - #[inline(always)] #[pyo3(signature = (url, **kwds))] pub fn trace( &self, @@ -818,12 +840,10 @@ impl BlockingClient { #[pymethods] impl BlockingClient { - #[inline] fn __enter__(slf: PyRef) -> PyRef { slf } - #[inline] fn __exit__<'py>( &self, _py: Python<'py>, diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index f6d823f3..76e31c41 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -1,23 +1,29 @@ use std::{ - pin::Pin, - sync::Arc, + future::Future, + pin::{Pin, pin}, + sync::{ + Arc, Mutex, MutexGuard, PoisonError, + atomic::{AtomicUsize, Ordering}, + }, task::{Context, Poll}, }; use bytes::Bytes; use futures_util::{FutureExt, Stream, future::poll_fn}; use http_body_util::BodyExt; -use pyo3::{ - coroutine::CancelHandle, exceptions::PyStopIteration, intern, prelude::*, sync::PyOnceLock, -}; +use pyo3::{exceptions::PyStopIteration, intern, prelude::*, sync::PyOnceLock}; use tokio::{ - sync::{Mutex, mpsc}, + sync::{ + Notify, + mpsc::{self, error::TryRecvError}, + }, task::JoinHandle, }; +use tokio_util::task::AbortOnDropHandle; use crate::{ + aio::{self, Coroutine}, buffer::PyBuffer, - client::nogil::{self, NoGIL}, error::Error, extractor::{BytesInput, StrInput}, header::HeaderMap, @@ -65,40 +71,106 @@ struct Sender(mpsc::Sender>>); /// A response stream yielding read-only memoryviews and any trailing headers. #[derive(Clone)] #[pyclass(subclass, frozen, skip_from_py_object)] -pub struct Streamer(Arc>>, Runtime); +pub struct Streamer(Arc, Runtime); + +/// Frames read ahead by a Tokio task, which starts on the first read so a read +/// timeout does not start before iteration. Readers poll the buffer from Python +/// and wait on `arrived`, which the task signals once per burst of frames. +struct Reader { + state: Mutex, + arrived: Arc, + /// Buffered frames returned since `__anext__` last yielded to the event loop. + ready: AtomicUsize, +} + +enum State { + Idle(Box), + Reading { + rx: mpsc::Receiver>, + _task: AbortOnDropHandle<()>, + }, + Closed, +} + +/// Wakes readers however the read-ahead task ends, including when aborted. +struct NotifyOnDrop(Arc); // ===== impl Streamer ===== impl Streamer { + /// Frames buffered ahead of Python, about 128 KiB of full TLS records. + const READ_AHEAD: usize = 8; + + /// Buffered frames returned before `__anext__` yields to the event loop once. + const YIELD_EVERY: usize = 8; + /// Create a new [`Streamer`] instance. #[inline] pub fn new(resp: wreq::Response, runtime: Runtime) -> Streamer { - Streamer(Arc::new(Mutex::new(Some(resp))), runtime) + let reader = Reader { + state: Mutex::new(State::Idle(Box::new(resp))), + arrived: Arc::new(Notify::new()), + ready: AtomicUsize::new(0), + }; + Streamer(Arc::new(reader), runtime) + } + + async fn read_ahead( + resp: wreq::Response, + tx: mpsc::Sender>, + arrived: Arc, + ) { + // Readers wake only after `pump` drops the sender, so they see the end. + let _end = NotifyOnDrop(arrived.clone()); + Self::pump(resp, tx, &arrived).await; + } + + async fn pump(mut resp: wreq::Response, tx: mpsc::Sender>, arrived: &Notify) { + while let Some(frame) = burst(resp.frame(), arrived).await { + let frame = match frame.map(|frame| frame.into_data()) { + Ok(Ok(bytes)) => Ok(Frame::Bytes(PyBuffer::from(bytes))), + Ok(Err(frame)) => match frame.into_trailers() { + Ok(trailers) => Ok(Frame::Trailers(HeaderMap(trailers))), + Err(_) => continue, + }, + Err(err) => { + // A failed body may leave a stalled HTTP/2 connection behind. + resp.forbid_recycle(); + Err(Error::Library(err).into()) + } + }; + let failed = frame.is_err(); + if burst(tx.send(frame), arrived).await.is_err() || failed { + break; + } + } } - async fn next(self, error: fn() -> Error) -> PyResult { - let frame = self - .0 - .lock() - .await - .as_mut() - .ok_or_else(error)? - .frame() - .await - .ok_or_else(error)? - .map_err(Error::Library)? - .into_data() - .map_err(|frame| frame.into_trailers()); - - match frame { - Ok(bytes) => Ok(Frame::Bytes(PyBuffer::from(bytes))), - Err(Ok(trailers)) => Ok(Frame::Trailers(HeaderMap(trailers))), - Err(Err(frame)) => { - // This branch should be unreachable, as `http_body::Frame` can only be `Data` or - // `Trailers`. The `debug_assert!` will help catch any future - // changes that violate this assumption. - debug_assert!(false, "Unexpected frame type: {:?}", frame); - Err(error().into()) + /// Return a buffered frame, or `None` if the read must wait for `arrived`. + fn try_next(&self, error: fn() -> Error) -> Option> { + let mut state = self.0.lock(); + if let State::Idle(_) = *state + && let State::Idle(resp) = std::mem::replace(&mut *state, State::Closed) + { + let (tx, rx) = mpsc::channel(Self::READ_AHEAD); + let task = self + .1 + .handle() + .spawn(Self::read_ahead(*resp, tx, self.0.arrived.clone())); + *state = State::Reading { + rx, + _task: AbortOnDropHandle::new(task), + }; + } + let State::Reading { rx, .. } = &mut *state else { + return Some(Err(error().into())); + }; + match rx.try_recv() { + Ok(frame) => Some(frame), + Err(TryRecvError::Empty) => None, + Err(TryRecvError::Disconnected) => { + *state = State::Closed; + Some(Err(error().into())) } } } @@ -106,22 +178,30 @@ impl Streamer { #[pymethods] impl Streamer { - #[inline] fn __iter__(slf: PyRef) -> PyRef { slf } - #[inline] fn __next__(&self, py: Python) -> PyResult { - py.detach(|| nogil::block_on(&self.1, self.clone().next(|| Error::StopIteration))) + // Buffered frames are returned without releasing the GIL. + if let Some(frame) = self.try_next(|| Error::StopIteration) { + return frame; + } + py.detach(|| { + loop { + let arrived = pin!(self.0.arrived.notified()); + if let Some(frame) = self.try_next(|| Error::StopIteration) { + return frame; + } + self.1.handle().block_on(arrived); + } + }) } - #[inline] fn __enter__(slf: PyRef) -> PyRef { slf } - #[inline] fn __exit__<'py>( &self, py: Python, @@ -129,66 +209,95 @@ impl Streamer { _exc_value: &Bound<'py, PyAny>, _traceback: &Bound<'py, PyAny>, ) { - py.detach(|| self.0.blocking_lock().take()); + py.detach(|| self.0.close()); } } #[pymethods] impl Streamer { - #[inline] fn __aiter__(slf: PyRef) -> PyRef { slf } - /// Read the next frame when awaited; returns a coroutine, not a Future. - #[inline] - fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { - let this = self.clone(); - let cancel = CancelHandle::new(); - // PyO3 0.29 cannot wrap an async __anext__ slot; use its macro constructor. - // Recheck this internal API when upgrading PyO3. - Bound::new( - py, - pyo3::impl_::coroutine::new_coroutine( - intern!(py, "__anext__"), - Some("Streamer"), - Some(cancel.throw_callback()), - async move { - let runtime = this.1.clone(); - let frame = NoGIL::with_cancel( - &runtime, - this.next(|| Error::StopAsyncIteration), - cancel, - ) - .await?; - // PyO3 polls this coroutine while attached, outside the Tokio task. - Python::attach(|py| frame.into_pyobject(py).map(|obj| obj.unbind())) - }, - ), - ) - .map(Bound::into_any) + /// Read the next frame when awaited. + fn __anext__(slf: Bound<'_, Self>) -> PyResult> { + let py = slf.py(); + let slf = slf.unbind(); + aio::local(py, "Streamer.__anext__", async move { + let this = slf.get(); + // Buffered frames complete without suspending; yield to the event loop + // periodically so timeouts, cancellation and other tasks run. + if this.0.ready.fetch_add(1, Ordering::Relaxed) >= Self::YIELD_EVERY { + this.0.ready.store(0, Ordering::Relaxed); + aio::yield_now().await; + } + loop { + // Readers are woken by `notify_waiters`, which reaches a `Notified` + // from its creation, so it needs no `enable`. + let arrived = pin!(this.0.arrived.notified()); + if let Some(frame) = this.try_next(|| Error::StopAsyncIteration) { + return frame; + } + this.0.ready.store(0, Ordering::Relaxed); + arrived.await; + } + }) } - #[inline] - async fn __aenter__(slf: Py) -> PyResult> { - Ok(slf) + fn __aenter__(slf: Bound<'_, Self>) -> PyResult> { + aio::ready("Streamer.__aenter__", slf) } - #[inline] - async fn __aexit__( + fn __aexit__<'py>( &self, + py: Python<'py>, _exc_type: Py, _exc_val: Py, _traceback: Py, - ) -> PyResult<()> { - let this = self.0.clone(); - NoGIL::new(&self.1, async move { - if let Some(resp) = this.lock().await.take() { - drop(resp) - } + ) -> PyResult> { + let reader = self.0.clone(); + aio::local(py, "Streamer.__aexit__", async move { + reader.close(); Ok(()) }) - .await + } +} + +/// Await `fut`, notifying readers whenever it suspends, so a burst of ready +/// frames costs one wake while a slow stream still delivers each frame at once. +async fn burst(fut: F, arrived: &Notify) -> F::Output { + let mut fut = pin!(fut); + poll_fn(|cx| { + let poll = fut.as_mut().poll(cx); + if poll.is_pending() { + arrived.notify_waiters(); + } + poll + }) + .await +} + +// ===== impl Reader ===== + +impl Reader { + #[inline] + fn lock(&self) -> MutexGuard<'_, State> { + self.state.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// Release the body and end any waiting read. + fn close(&self) { + let state = std::mem::replace(&mut *self.lock(), State::Closed); + drop(state); + self.arrived.notify_waiters(); + } +} + +// ===== impl NotifyOnDrop ===== + +impl Drop for NotifyOnDrop { + fn drop(&mut self) { + self.0.notify_waiters(); } } @@ -372,48 +481,29 @@ impl Drop for PyAsyncStream { #[pymethods] impl Sender { - async fn send( + fn send<'py>( &self, - item: Py, + py: Python<'py>, + item: Bound<'py, PyAny>, error: bool, - #[pyo3(cancel_handle)] cancel: CancelHandle, - ) -> PyResult { - let item = Python::attach(|py| { - if error { - Ok(Err(PyErr::from_value(item.into_bound(py)))) - } else { - item.extract(py).map(Ok) - } - })?; - self.send_item(Some(item), cancel).await + ) -> PyResult> { + let item = if error { + Err(PyErr::from_value(item)) + } else { + Ok(item.extract()?) + }; + let tx = self.0.clone(); + // Channel readiness is runtime-independent, so this waits on the Python loop. + aio::local(py, "Sender.send", async move { + Ok(tx.send(Some(item)).await.is_ok()) + }) } - async fn finish(&self, #[pyo3(cancel_handle)] cancel: CancelHandle) -> PyResult { + fn finish<'py>(&self, py: Python<'py>) -> PyResult> { // Python may retain the sender after completion, especially on PyPy. - self.send_item(None, cancel).await - } -} - -impl Sender { - async fn send_item( - &self, - item: Option>, - mut cancel: CancelHandle, - ) -> PyResult { - let item = match self.0.try_send(item) { - Ok(()) => return Ok(true), - Err(mpsc::error::TrySendError::Closed(_)) => return Ok(false), - Err(mpsc::error::TrySendError::Full(item)) => item, - }; let tx = self.0.clone(); - // Channel readiness is runtime-independent; keep this on the Python loop. - let mut send = std::pin::pin!(tx.send(item)); - tokio::select! { - biased; - exception = poll_fn(|cx| cancel.poll_cancelled(cx)) => { - Err(Python::attach(|py| PyErr::from_value(exception.into_bound(py)))) - } - result = poll_fn(|cx| nogil::poll_with_guard(send.as_mut(), cx)) => Ok(result.is_ok()), - } + aio::local(py, "Sender.finish", async move { + Ok(tx.send(None).await.is_ok()) + }) } } diff --git a/src/client/nogil.rs b/src/client/nogil.rs index 8afe8ff1..e9a6db42 100644 --- a/src/client/nogil.rs +++ b/src/client/nogil.rs @@ -1,90 +1,9 @@ -use std::{ - future::Future, - pin::Pin, - sync::Arc, - task::{Context, Poll, Wake, Waker}, -}; +use std::future::Future; -use pyo3::{coroutine::CancelHandle, exceptions::PyRuntimeError, prelude::*}; -use tokio_util::task::AbortOnDropHandle; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; use crate::runtime::Runtime; -/// A future that allows Python threads to run while it is being polled or executed. -/// It also handles cancellation and spawns the task in tokio runtime. -pub struct NoGIL { - handle: AbortOnDropHandle>, - cancel: CancelHandle, -} - -struct GuardedWaker(Waker); - -// ===== impl NoGIL ===== - -impl NoGIL -where - T: Send + 'static, -{ - /// Spawn internal work without a Python cancellation source. - #[inline] - pub fn new(runtime: &Runtime, fut: Fut) -> Self - where - Fut: Future> + Send + 'static, - { - Self::with_cancel(runtime, fut, CancelHandle::new()) - } - - /// Spawn with Python cancellation, keeping the runtime alive until the task ends. - #[inline] - pub fn with_cancel(runtime: &Runtime, fut: Fut, cancel: CancelHandle) -> Self - where - Fut: Future> + Send + 'static, - { - let owner = runtime.clone(); - Self { - handle: AbortOnDropHandle::new(Python::attach(|py| { - py.detach(|| { - runtime.handle().spawn(async move { - let _owner = owner; - fut.await - }) - }) - })), - cancel, - } - } -} - -impl Future for NoGIL -where - T: Send + 'static, -{ - type Output = PyResult; - - #[inline] - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.get_mut(); - // A Python throw must win even when the Tokio task has already finished. - if let Poll::Ready(exc) = this.cancel.poll_cancelled(cx) { - this.handle.abort(); - return Poll::Ready(Err(Python::attach(|py| { - PyErr::from_value(exc.into_bound(py)) - }))); - } - - let waker = cx.waker(); - Python::attach(|py| { - py.detach(|| { - match poll_with_guard(Pin::new(&mut this.handle), &mut Context::from_waker(waker)) { - Poll::Ready(Ok(result)) => Poll::Ready(result), - Poll::Ready(Err(e)) => Poll::Ready(Err(PyRuntimeError::new_err(e.to_string()))), - Poll::Pending => Poll::Pending, - } - }) - }) - } -} - /// Run network work on its selected worker; only wait for completion on the caller. /// The caller must be detached from Python and outside an async Tokio context. /// Its runtime borrow keeps the workers alive until the join completes. @@ -99,23 +18,3 @@ where .block_on(task) .map_err(|err| PyRuntimeError::new_err(err.to_string()))? } - -/// Protect Python wakers retained by futures that can be woken from Rust threads. -pub fn poll_with_guard(future: Pin<&mut F>, cx: &mut Context<'_>) -> Poll { - let waker = Waker::from(Arc::new(GuardedWaker(cx.waker().clone()))); - future.poll(&mut Context::from_waker(&waker)) -} - -// ===== impl GuardedWaker ===== - -impl Wake for GuardedWaker { - fn wake(self: Arc) { - self.wake_by_ref(); - } - - fn wake_by_ref(self: &Arc) { - // PyO3's nested attach reuses this attachment. If Python is unavailable, - // skip the wake instead of invoking its infallible attachment path. - Python::try_attach(|_| self.0.wake_by_ref()); - } -} diff --git a/src/client/resp/http.rs b/src/client/resp/http.rs index eba23c94..f34e429d 100644 --- a/src/client/resp/http.rs +++ b/src/client/resp/http.rs @@ -1,4 +1,4 @@ -use std::{fmt::Display, sync::Arc}; +use std::{fmt::Display, future::Future, sync::Arc}; use arc_swap::ArcSwapOption; use bytes::Bytes; @@ -8,15 +8,15 @@ use futures_util::{ }; use http::response::{Parts, Response as HttpResponse}; use http_body_util::{BodyExt, Collected}; -use pyo3::{coroutine::CancelHandle, prelude::*, pybacked::PyBackedStr}; +use pyo3::{prelude::*, pybacked::PyBackedStr}; use wreq::{self, Uri}; use crate::{ + aio::{self, Coroutine}, buffer::PyBuffer, client::{ SocketAddr, body::{Json, Streamer}, - nogil::NoGIL, resp::ext::ResponseExt, }, cookie::Cookie, @@ -49,6 +49,9 @@ enum Body { #[pyclass(name = "Response", subclass, frozen, str, skip_from_py_object)] pub struct BlockingResponse(Response); +/// Forbids connection reuse unless disarmed by taking the response parts. +struct RecycleGuard(Option); + // ===== impl Response ===== impl Response { @@ -89,11 +92,15 @@ impl Response { match Arc::into_inner(arc) { Some(Body::Streamable(stream)) => { return Box::pin(async move { + // A failed or cancelled read may leave a stalled HTTP/2 + // connection behind, so only a complete body keeps it reusable. + let mut guard = RecycleGuard(Some(parts)); let bytes = stream .collect() .await .map(Collected::to_bytes) .map_err(Error::Library)?; + let parts = guard.0.take().ok_or(Error::Memory)?; body.store(Some(Arc::new(Body::Reusable(bytes.clone())))); let response = HttpResponse::from_parts(parts, bytes); @@ -122,6 +129,26 @@ impl Response { Err(Error::Memory) } + /// Read the body on the runtime once awaited; the body is taken on first await. + fn read<'py, F, Fut, T>( + slf: Bound<'py, Self>, + qualname: &'static str, + read: F, + ) -> PyResult> + where + F: FnOnce(wreq::Response) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + T: for<'a> IntoPyObject<'a> + Send + 'static, + { + let py = slf.py(); + let slf = slf.unbind(); + aio::local(py, qualname, async move { + let this = slf.get(); + let fut = this.cache_response().and_then(read).map_err(Into::into); + aio::run(this.runtime.clone(), fut).await + }) + } + /// Forcefully destroys the response [`Body`], preventing any further reads. fn destroy(&self) { #[allow(clippy::option_map_unit_fn)] @@ -225,62 +252,64 @@ impl Response { /// Get the text content with the response encoding, defaulting to utf-8 when unspecified. #[pyo3(signature = (encoding = None))] - pub async fn text( - &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn text( + slf: Bound<'_, Self>, encoding: Option, - ) -> PyResult { - let fut = self - .cache_response() - .and_then(|resp| ResponseExt::text(resp, encoding)) - .map_err(Into::into); - NoGIL::with_cancel(&self.runtime, fut, cancel).await + ) -> PyResult> { + Self::read(slf, "Response.text", |resp| { + ResponseExt::text(resp, encoding) + }) } /// Get the JSON content of the response. - pub async fn json(&self, #[pyo3(cancel_handle)] cancel: CancelHandle) -> PyResult { - let fut = self - .cache_response() - .and_then(ResponseExt::json::) - .map_err(Into::into); - NoGIL::with_cancel(&self.runtime, fut, cancel).await + pub fn json(slf: Bound<'_, Self>) -> PyResult> { + Self::read(slf, "Response.json", ResponseExt::json::) } /// Read the body as a read-only memoryview, retaining its data after the response closes. - pub async fn bytes(&self, #[pyo3(cancel_handle)] cancel: CancelHandle) -> PyResult { - let fut = self - .cache_response() - .and_then(ResponseExt::bytes) - .map_ok(PyBuffer::from) - .map_err(Into::into); - NoGIL::with_cancel(&self.runtime, fut, cancel).await + pub fn bytes(slf: Bound<'_, Self>) -> PyResult> { + Self::read(slf, "Response.bytes", |resp| { + ResponseExt::bytes(resp).map_ok(PyBuffer::from) + }) } /// Discard the retained body and mark its connection as non-reusable. /// This does not guarantee an immediate socket shutdown or cancel an active read. /// Cancel and await any body-read task before closing. A body transferred to a /// Streamer is managed separately; previously returned memoryviews remain valid. - /// Prefer an async context manager (`async with`) for response cleanup. - pub async fn close(&self) { - Python::attach(|py| { - py.detach(|| { - self.empty_response().forbid_recycle(); - self.destroy() - }); - }); + /// `async with` instead releases the body and keeps a fully read connection reusable. + pub fn close(slf: Bound<'_, Self>) -> PyResult> { + let py = slf.py(); + let slf = slf.unbind(); + aio::local(py, "Response.close", async move { + let this = slf.get(); + this.empty_response().forbid_recycle(); + this.destroy(); + Ok(()) + }) } } #[pymethods] impl Response { - #[inline] - async fn __aenter__(slf: Py) -> PyResult> { - Ok(slf) - } - - #[inline] - async fn __aexit__(&self, _exc_type: Py, _exc_val: Py, _traceback: Py) { - self.close().await + fn __aenter__(slf: Bound<'_, Self>) -> PyResult> { + aio::ready("Response.__aenter__", slf) + } + + /// Release the body without forbidding reuse: a fully read connection returns + /// to the pool, while an unread HTTP/1 body drains or closes its connection. + fn __aexit__<'py>( + slf: Bound<'py, Self>, + _exc_type: Py, + _exc_val: Py, + _traceback: Py, + ) -> PyResult> { + let py = slf.py(); + let slf = slf.unbind(); + aio::local(py, "Response.__aexit__", async move { + slf.get().destroy(); + Ok(()) + }) } } @@ -303,6 +332,16 @@ impl Drop for Response { } } +// ===== impl RecycleGuard ===== + +impl Drop for RecycleGuard { + fn drop(&mut self) { + if let Some(parts) = self.0.take() { + wreq::Response::from(HttpResponse::from_parts(parts, Bytes::new())).forbid_recycle(); + } + } +} + // ===== impl BlockingResponse ===== #[pymethods] @@ -368,13 +407,11 @@ impl BlockingResponse { } /// Turn a response into an error if the server returned an error. - #[inline] pub fn raise_for_status(&self) -> PyResult<()> { self.0.raise_for_status() } /// Stream read-only memoryviews and any trailing headers from the body. - #[inline] pub fn stream(&self) -> PyResult { self.0.stream() } @@ -421,8 +458,7 @@ impl BlockingResponse { /// This does not guarantee an immediate socket shutdown or interrupt an active read. /// Do not close concurrently with a body read. A body transferred to a Streamer /// is managed separately; previously returned memoryviews remain valid. - /// Prefer a context manager (`with`) for response cleanup. - #[inline] + /// `with` instead releases the body and keeps a fully read connection reusable. pub fn close(&self, py: Python) { py.detach(|| { self.0.empty_response().forbid_recycle(); @@ -433,12 +469,10 @@ impl BlockingResponse { #[pymethods] impl BlockingResponse { - #[inline] fn __enter__(slf: PyRef) -> PyRef { slf } - #[inline] fn __exit__<'py>( &self, py: Python<'py>, @@ -446,7 +480,8 @@ impl BlockingResponse { _exc_value: &Bound<'py, PyAny>, _traceback: &Bound<'py, PyAny>, ) { - self.close(py) + // Like `__aexit__`, release the body and leave reuse to the protocol. + py.detach(|| self.0.destroy()) } } diff --git a/src/client/resp/ws.rs b/src/client/resp/ws.rs index e273928c..bd81a459 100644 --- a/src/client/resp/ws.rs +++ b/src/client/resp/ws.rs @@ -4,7 +4,7 @@ pub mod msg; use std::{fmt::Display, time::Duration}; use msg::Message; -use pyo3::{coroutine::CancelHandle, prelude::*}; +use pyo3::prelude::*; use tokio::sync::mpsc; use wreq::{ header::HeaderValue, @@ -12,7 +12,8 @@ use wreq::{ }; use crate::{ - client::{SocketAddr, nogil::NoGIL}, + aio::{self, Coroutine}, + client::SocketAddr, cookie::Cookie, error::Error, extractor::StrInput, @@ -103,66 +104,81 @@ impl WebSocket { /// Receive a message from the WebSocket. #[pyo3(signature = (timeout=None))] - pub async fn recv( + pub fn recv<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, timeout: Option, - ) -> PyResult> { - let tx = self.cmd.clone(); - NoGIL::with_cancel(&self.runtime, cmd::recv(tx, timeout), cancel).await + ) -> PyResult> { + aio::spawn( + py, + "WebSocket.recv", + &self.runtime, + cmd::recv(self.cmd.clone(), timeout), + ) } /// Send a message to the WebSocket. #[pyo3(signature = (message))] - pub async fn send( - &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, - message: Message, - ) -> PyResult<()> { - let tx = self.cmd.clone(); - NoGIL::with_cancel(&self.runtime, cmd::send(tx, message), cancel).await + pub fn send<'py>(&self, py: Python<'py>, message: Message) -> PyResult> { + aio::spawn( + py, + "WebSocket.send", + &self.runtime, + cmd::send(self.cmd.clone(), message), + ) } /// Send multiple messages to the WebSocket. #[pyo3(signature = (messages))] - pub async fn send_all( + pub fn send_all<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, messages: Vec, - ) -> PyResult<()> { - let tx = self.cmd.clone(); - NoGIL::with_cancel(&self.runtime, cmd::send_all(tx, messages), cancel).await + ) -> PyResult> { + aio::spawn( + py, + "WebSocket.send_all", + &self.runtime, + cmd::send_all(self.cmd.clone(), messages), + ) } /// Close the WebSocket connection. #[pyo3(signature = (code=None, reason=None))] - pub async fn close( + pub fn close<'py>( &self, - #[pyo3(cancel_handle)] cancel: CancelHandle, + py: Python<'py>, code: Option, reason: Option, - ) -> PyResult<()> { - let tx = self.cmd.clone(); - NoGIL::with_cancel(&self.runtime, cmd::close(tx, code, reason), cancel).await + ) -> PyResult> { + aio::spawn( + py, + "WebSocket.close", + &self.runtime, + cmd::close(self.cmd.clone(), code, reason), + ) } } #[pymethods] impl WebSocket { - #[inline] - async fn __aenter__(slf: Py) -> PyResult> { - Ok(slf) + fn __aenter__(slf: Bound<'_, Self>) -> PyResult> { + aio::ready("WebSocket.__aenter__", slf) } - #[inline] - async fn __aexit__( + fn __aexit__<'py>( &self, + py: Python<'py>, _exc_type: Py, _exc_val: Py, _traceback: Py, - ) -> PyResult<()> { - let tx = self.cmd.clone(); - NoGIL::new(&self.runtime, cmd::close(tx, None, None)).await + ) -> PyResult> { + aio::spawn( + py, + "WebSocket.__aexit__", + &self.runtime, + cmd::close(self.cmd.clone(), None, None), + ) } } @@ -259,12 +275,10 @@ impl BlockingWebSocket { #[pymethods] impl BlockingWebSocket { - #[inline] fn __enter__(slf: PyRef) -> PyRef { slf } - #[inline] fn __exit__<'py>( &self, py: Python<'py>, diff --git a/src/header.rs b/src/header.rs index 057aed64..85e68f83 100644 --- a/src/header.rs +++ b/src/header.rs @@ -141,7 +141,6 @@ impl HeaderMap { } /// Returns a list of read-only memoryviews for all keys. - #[inline] fn keys<'py>(&self, py: Python<'py>) -> Vec { py.detach(|| { self.0 @@ -153,7 +152,6 @@ impl HeaderMap { } /// Returns a list of read-only memoryviews for all values. - #[inline] fn values<'py>(&self, py: Python<'py>) -> Vec { py.detach(|| { self.0 @@ -169,7 +167,6 @@ impl HeaderMap { /// This number represents the total number of **values** stored in the map. /// This number can be greater than or equal to the number of **keys** /// stored given that a single key may have more than one associated value. - #[inline] fn len(&self) -> usize { self.0.len() } @@ -178,19 +175,16 @@ impl HeaderMap { /// /// This number will be less than or equal to `len()` as each key may have /// more than one associated value. - #[inline] fn keys_len(&self) -> usize { self.0.keys_len() } /// Returns true if the map contains no elements. - #[inline] fn is_empty(&self) -> bool { self.0.is_empty() } /// Clears the map, removing all key-value pairs. Keeps the allocated memory for reuse. - #[inline] fn clear(&mut self) { self.0.clear(); } @@ -198,27 +192,22 @@ impl HeaderMap { #[pymethods] impl HeaderMap { - #[inline] fn __getitem__<'py>(&self, py: Python<'py>, key: PyBackedStr) -> Option { self.get(py, key, None) } - #[inline] fn __setitem__(&mut self, py: Python, key: PyBackedStr, value: StrInput) { self.insert(py, key, value); } - #[inline] fn __delitem__(&mut self, py: Python, key: PyBackedStr) { self.remove(py, key); } - #[inline] fn __contains__(&self, py: Python, key: PyBackedStr) -> bool { self.contains_key(py, key) } - #[inline] fn __len__(&self) -> usize { self.0.len() } @@ -305,13 +294,11 @@ impl OrigHeaderMap { /// of the list of values currently associated with the key. The key is not /// updated, though; this matters for types that can be `==` without being /// identical. - #[inline] pub fn insert(&mut self, value: StrInput) -> bool { self.0.insert(value.0) } /// Extends the map with all entries from another [`OrigHeaderMap`], preserving order. - #[inline] pub fn extend(&mut self, iter: &Bound<'_, OrigHeaderMap>) { self.0.extend(iter.borrow().0.clone()); } @@ -319,7 +306,6 @@ impl OrigHeaderMap { #[pymethods] impl OrigHeaderMap { - #[inline] fn __len__(&self) -> usize { self.0.len() } diff --git a/src/http.rs b/src/http.rs index 93078582..af1d09db 100644 --- a/src/http.rs +++ b/src/http.rs @@ -42,37 +42,31 @@ impl<'a, 'py> FromPyObject<'a, 'py> for StatusCode { #[pymethods] impl StatusCode { /// Return the status code as an integer. - #[inline] pub const fn as_int(&self) -> u16 { self.0.as_u16() } /// Check if status is within 100-199. - #[inline] pub fn is_informational(&self) -> bool { self.0.is_informational() } /// Check if status is within 200-299. - #[inline] pub fn is_success(&self) -> bool { self.0.is_success() } /// Check if status is within 300-399. - #[inline] pub fn is_redirection(&self) -> bool { self.0.is_redirection() } /// Check if status is within 400-499. - #[inline] pub fn is_client_error(&self) -> bool { self.0.is_client_error() } /// Check if status is within 500-599. - #[inline] pub fn is_server_error(&self) -> bool { self.0.is_server_error() } diff --git a/src/lib.rs b/src/lib.rs index 183a9a1e..d2abbd8d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,6 +5,7 @@ #[macro_use] mod macros; +mod aio; mod buffer; mod client; mod cookie; @@ -27,7 +28,6 @@ use client::{ Streamer, multipart::{Multipart, Part}, }, - req::WebSocketRequest, resp::{BlockingResponse, BlockingWebSocket, Message, Response, WebSocket}, }; use cookie::{Cookie, Jar, SameSite}; @@ -44,10 +44,7 @@ use http2::{ #[cfg(feature = "mimalloc")] use mimalloc as _; use proxy::Proxy; -use pyo3::{ - coroutine::CancelHandle, intern, prelude::*, pybacked::PyBackedStr, types::PyDict, - wrap_pymodule, -}; +use pyo3::{intern, prelude::*, types::PyDict, wrap_pymodule}; use runtime::Runtime; #[cfg(feature = "jemalloc")] use tikv_jemallocator as _; @@ -68,136 +65,141 @@ static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; mod r#async { - use pyo3::{coroutine::CancelHandle, prelude::*, pybacked::PyBackedStr}; + use pyo3::{prelude::*, pybacked::PyBackedStr, types::PyDict}; use crate::{ - client::{ - Client, - req::{Request, WebSocketRequest}, - resp::{Response, WebSocket}, - }, + aio::{self, Coroutine}, + client::Client, http::Method, }; /// Make a GET request with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] - pub async fn get( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn get<'py>( + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - request(cancel, Method::GET, url, kwds).await + kwds: Option>, + ) -> PyResult> { + send(py, "get", Method::GET, url, kwds) } /// Make a POST request with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] - pub async fn post( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn post<'py>( + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - request(cancel, Method::POST, url, kwds).await + kwds: Option>, + ) -> PyResult> { + send(py, "post", Method::POST, url, kwds) } /// Make a PUT request with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] - pub async fn put( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn put<'py>( + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - request(cancel, Method::PUT, url, kwds).await + kwds: Option>, + ) -> PyResult> { + send(py, "put", Method::PUT, url, kwds) } /// Make a PATCH request with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] - pub async fn patch( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn patch<'py>( + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - request(cancel, Method::PATCH, url, kwds).await + kwds: Option>, + ) -> PyResult> { + send(py, "patch", Method::PATCH, url, kwds) } /// Make a DELETE request with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] - pub async fn delete( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn delete<'py>( + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - request(cancel, Method::DELETE, url, kwds).await + kwds: Option>, + ) -> PyResult> { + send(py, "delete", Method::DELETE, url, kwds) } /// Make a HEAD request with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] - pub async fn head( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn head<'py>( + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - request(cancel, Method::HEAD, url, kwds).await + kwds: Option>, + ) -> PyResult> { + send(py, "head", Method::HEAD, url, kwds) } /// Make a OPTIONS request with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] - pub async fn options( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn options<'py>( + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - request(cancel, Method::OPTIONS, url, kwds).await + kwds: Option>, + ) -> PyResult> { + send(py, "options", Method::OPTIONS, url, kwds) } /// Make a TRACE request with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] - pub async fn trace( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn trace<'py>( + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - request(cancel, Method::TRACE, url, kwds).await + kwds: Option>, + ) -> PyResult> { + send(py, "trace", Method::TRACE, url, kwds) } /// Make a request with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (method, url, **kwds))] - pub async fn request( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn request<'py>( + py: Python<'py>, method: Method, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - Client::default().request(cancel, method, url, kwds).await + kwds: Option>, + ) -> PyResult> { + send(py, "request", method, url, kwds) } /// Make a WebSocket connection with the given parameters. - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] - pub async fn websocket( - #[pyo3(cancel_handle)] cancel: CancelHandle, + pub fn websocket<'py>( + py: Python<'py>, url: PyBackedStr, - kwds: Option, - ) -> PyResult { - Client::default().websocket(cancel, url, kwds).await + kwds: Option>, + ) -> PyResult> { + let kwds = kwds.map(Bound::unbind); + aio::local(py, "websocket", async move { + Client::default().connect(url, kwds).await + }) + } + + /// Return a coroutine named `qualname` that sends a request with the default + /// client, created on first await so an unawaited call is inert. + fn send<'py>( + py: Python<'py>, + qualname: &'static str, + method: Method, + url: PyBackedStr, + kwds: Option>, + ) -> PyResult> { + let kwds = kwds.map(Bound::unbind); + aio::local(py, qualname, async move { + Client::default().execute(method, url, kwds).await + }) } } @@ -214,7 +216,6 @@ mod blocking { }; /// Make a GET request with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] pub fn get(py: Python, url: PyBackedStr, kwds: Option) -> PyResult { @@ -222,7 +223,6 @@ mod blocking { } /// Make a POST request with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] pub fn post(py: Python, url: PyBackedStr, kwds: Option) -> PyResult { @@ -230,7 +230,6 @@ mod blocking { } /// Make a PUT request with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] pub fn put(py: Python, url: PyBackedStr, kwds: Option) -> PyResult { @@ -238,7 +237,6 @@ mod blocking { } /// Make a PATCH request with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] pub fn patch( @@ -250,7 +248,6 @@ mod blocking { } /// Make a DELETE request with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] pub fn delete( @@ -262,7 +259,6 @@ mod blocking { } /// Make a HEAD request with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] pub fn head(py: Python, url: PyBackedStr, kwds: Option) -> PyResult { @@ -270,7 +266,6 @@ mod blocking { } /// Make a OPTIONS request with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] pub fn options( @@ -282,7 +277,6 @@ mod blocking { } /// Make a TRACE request with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] pub fn trace( @@ -294,7 +288,6 @@ mod blocking { } /// Make a request with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (method, url, **kwds))] pub fn request( @@ -307,7 +300,6 @@ mod blocking { } /// Make a WebSocket connection with the given parameters (blocking). - #[inline] #[pyfunction] #[pyo3(signature = (url, **kwds))] pub fn websocket( @@ -319,18 +311,6 @@ mod blocking { } } -/// Make a WebSocket connection with the given parameters. -#[inline] -#[pyfunction] -#[pyo3(signature = (url, **kwds))] -pub async fn websocket( - #[pyo3(cancel_handle)] cancel: CancelHandle, - url: PyBackedStr, - kwds: Option, -) -> PyResult { - Client::default().websocket(cancel, url, kwds).await -} - #[pymodule(gil_used = false)] fn wreq(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { use r#async::{delete, get, head, options, patch, post, put, request, trace, websocket}; diff --git a/src/redirect.rs b/src/redirect.rs index 97b0191a..518433ae 100644 --- a/src/redirect.rs +++ b/src/redirect.rs @@ -136,7 +136,6 @@ impl_print_str!(Debug, Policy); #[pymethods] impl Attempt { /// Returns an action meaning the client should follow the next URI. - #[inline] pub fn follow(&self) -> Action { Action { kind: ActionKind::Follow, @@ -146,7 +145,6 @@ impl Attempt { /// Returns an action meaning the client should not follow the next URI. /// /// The 30x response will be returned as the result. - #[inline] pub fn stop(&self) -> Action { Action { kind: ActionKind::Stop, @@ -156,7 +154,6 @@ impl Attempt { /// Returns an action failing the redirect with an error. /// /// The error will be returned for the result of the sent request. - #[inline] pub fn error(&self, message: String) -> Action { Action { kind: ActionKind::Error(message), diff --git a/src/tls.rs b/src/tls.rs index 59e3d2a3..5bcd003d 100644 --- a/src/tls.rs +++ b/src/tls.rs @@ -427,7 +427,6 @@ pub struct TlsInfo(pub wreq::tls::TlsInfo); #[pymethods] impl TlsInfo { /// Get a read-only memoryview of the peer's DER-encoded leaf certificate. - #[inline] pub fn peer_certificate(&self) -> Option { self.0 .peer_certificate() diff --git a/tests/aio_test.py b/tests/aio_test.py new file mode 100644 index 00000000..d09ab0b6 --- /dev/null +++ b/tests/aio_test.py @@ -0,0 +1,143 @@ +import asyncio +import gc +import socket +import sys +import weakref + +import pytest +import wreq + + +class NoReaderLoop(asyncio.SelectorEventLoop): + """A loop that cannot watch the wake socket, so wakes are scheduled thread-safely.""" + + def add_reader(self, *args): + raise NotImplementedError + + +LOOPS = [ + pytest.param(asyncio.new_event_loop, id="default"), + pytest.param(NoReaderLoop, id="no-reader"), +] +if sys.platform == "win32": + LOOPS.append(pytest.param(asyncio.SelectorEventLoop, id="selector")) +try: + import uvloop +except ImportError: + pass +else: + LOOPS.append(pytest.param(uvloop.new_event_loop, id="uvloop")) + + +async def exchange(): + hanging = asyncio.Event() + abandoned = asyncio.Event() + handlers = set() + + async def serve(reader, writer): + handlers.add(asyncio.current_task()) + try: + while True: + head = await reader.readuntil(b"\r\n\r\n") + path = head.split(b" ", 2)[1] + if path == b"/chunks": + writer.write( + b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n" + ) + writer.write(b"1\r\na\r\n1\r\nb\r\n0\r\n\r\n") + await writer.drain() + continue + if path == b"/hang": + hanging.set() + await reader.read() + abandoned.set() + break + body = path.lstrip(b"/") * 4096 + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n" % len(body) + ) + writer.write(body) + await writer.drain() + except (asyncio.IncompleteReadError, ConnectionError): + pass + finally: + writer.close() + + server = await asyncio.start_server(serve, "127.0.0.1", 0) + url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}" + try: + async with wreq.Client(proxies=[]) as client: + + async def fetch(i): + async with await client.get(f"{url}/{i}") as response: + return bytes(await response.bytes()) + + # Concurrent completions resume their tasks through the loop's wake port. + bodies = await asyncio.gather(*(fetch(i) for i in range(64))) + assert bodies == [str(i).encode() * 4096 for i in range(64)] + + response = await client.get(f"{url}/stream") + async with response.stream() as stream: + frames = [bytes(chunk) async for chunk in stream] + assert b"".join(frames) == b"stream" * 4096 + + # Short streams end right after their last frame; readers must see the end. + async def drain_chunks(): + async with await client.get(f"{url}/chunks") as response: + return b"".join([bytes(c) async for c in response.stream()]) + + assert ( + await asyncio.gather(*(drain_chunks() for _ in range(300))) + == [b"ab"] * 300 + ) + + # Cancelling a pending request resumes the task and aborts its Tokio task, + # which drops the connection. + hang = client.get(f"{url}/hang") + pending = asyncio.ensure_future(hang) + await asyncio.wait_for(hanging.wait(), 5) + # A second task awaiting the same coroutine fails without stranding the first. + with pytest.raises(RuntimeError, match="awaited already"): + await hang + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + await asyncio.wait_for(abandoned.wait(), 5) + finally: + server.close() + for handler in handlers: + handler.cancel() + await asyncio.gather(*handlers, return_exceptions=True) + await server.wait_closed() + + +@pytest.mark.parametrize("new_loop", LOOPS) +def test_wakes_resume_tasks_on_event_loops(new_loop): + # Each new loop gets its own wake port. + for _ in range(3): + loop = new_loop() + try: + loop.run_until_complete(asyncio.wait_for(exchange(), 30)) + finally: + loop.close() + + +@pytest.mark.skipif( + sys.implementation.name != "cpython", + reason="PyPy's cpyext does not collect cycles through extension objects", +) +@pytest.mark.parametrize("new_loop", [p for p in LOOPS if p.id != "no-reader"]) +def test_closed_loop_releases_pending_requests(new_loop): + # The listener accepts connections into its backlog but never answers. + with socket.create_server(("127.0.0.1", 0)) as server: + url = f"http://127.0.0.1:{server.getsockname()[1]}/" + client = wreq.Client(proxies=[]) + loop = new_loop() + task = loop.create_task(client.get(url)) + loop.run_until_complete(asyncio.sleep(0.2)) + loop.close() + # Once its loop is closed, a task left pending can be collected with its request. + ref = weakref.ref(task) + del task + gc.collect() + assert ref() is None diff --git a/tests/response_test.py b/tests/response_test.py index df9d0ba9..62a2175c 100644 --- a/tests/response_test.py +++ b/tests/response_test.py @@ -1,8 +1,16 @@ +import asyncio +import socket +import threading +import time +from datetime import timedelta + import pytest import wreq from pathlib import Path from wreq import Version, Multipart, Part +from cancellation_test import local_server + client = wreq.Client(tls_info=True) @@ -120,3 +128,95 @@ async def test_peer_certificate(): certificate = resp.tls_info.peer_certificate() assert type(certificate) is memoryview assert certificate.readonly + + +@pytest.mark.asyncio +async def test_context_exit_keeps_connection_but_close_forbids_reuse(): + reply = b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok" + async with local_server() as (url, connections), wreq.Client(proxies=[]) as client: + task = asyncio.create_task(client.get(url)) + reader, writer = await asyncio.wait_for(connections.get(), 5) + writer.write(reply) + async with await asyncio.wait_for(task, 5) as response: + assert bytes(await response.bytes()) == b"ok" + + # Leaving `async with` after a full read returns the connection to the pool. + task = asyncio.create_task(client.get(url)) + await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), 5) + writer.write(reply) + response = await asyncio.wait_for(task, 5) + assert bytes(await response.bytes()) == b"ok" + await response.close() + + # An explicit close() still forbids reuse, so the next request reconnects. + task = asyncio.create_task(client.get(url)) + _, writer = await asyncio.wait_for(connections.get(), 5) + writer.write(reply) + assert bytes(await (await asyncio.wait_for(task, 5)).bytes()) == b"ok" + assert connections.empty() + + +@pytest.mark.asyncio +async def test_stream_read_ahead_yields_and_closes_waiting_readers(): + stalled = threading.Event() + + def serve(listener): + for size in (64 << 20, None): + conn, _ = listener.accept() + with conn: + while b"\r\n\r\n" not in conn.recv(65536): + pass + length = size or 1 << 20 + conn.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n" % length) + try: + conn.sendall(b"x" * (size or 4096)) + except OSError: + pass + if size is None: + stalled.wait(10) + + listener = socket.create_server(("127.0.0.1", 0)) + url = f"http://127.0.0.1:{listener.getsockname()[1]}/" + thread = threading.Thread(target=serve, args=(listener,), daemon=True) + thread.start() + try: + async with wreq.Client(proxies=[]) as client: + response = await client.get(url, read_timeout=timedelta(seconds=1)) + streamer = response.stream() + # Reading starts on the first iteration, so this wait is not a read timeout. + await asyncio.sleep(1.5) + + ticks = 0 + + async def tick(): + nonlocal ticks + while True: + ticks += 1 + await asyncio.sleep(0) + + ticker = asyncio.create_task(tick()) + async with asyncio.timeout(5), streamer: + frames = 0 + async for _ in streamer: + # A consumer slower than the network must still yield to the loop. + time.sleep(0.002) + frames += 1 + if frames == 64: + break + ticker.cancel() + assert frames == 64 and ticks > 1 + + response = await client.get(url) + streamer = response.stream() + assert len(await anext(streamer)) == 4096 + reader = asyncio.create_task(anext(streamer)) + await asyncio.sleep(0.1) + assert not reader.done() + # A synchronous exit must not wait for the pending read; it ends that read. + streamer.__exit__(None, None, None) + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(reader, 5) + finally: + stalled.set() + listener.close() + thread.join(10) diff --git a/tests/runtime_test.py b/tests/runtime_test.py index 178e685a..5ab1a93d 100644 --- a/tests/runtime_test.py +++ b/tests/runtime_test.py @@ -133,12 +133,16 @@ def request(url): async with local_server() as (url, connections): task = asyncio.create_task(asyncio.to_thread(request, url)) - for _ in range(2): - reader, writer = await asyncio.wait_for(connections.get(), 5) + reader, writer = await asyncio.wait_for(connections.get(), 5) + for index in range(2): + if index: + # Leaving `with` after a full read keeps the connection reusable. + await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), 5) assert await asyncio.wait_for(read_chunked(reader), 5) == b"sync upload" writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}") await writer.drain() assert await asyncio.wait_for(task, 5) == b"{}" + assert connections.empty() del task