perf(aio): speed up async requests and streams

.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:

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" }

pyproject.toml

@@ -41,6 +41,7 @@ test = [
   "pytest",
   "pytest-asyncio",
   "pytest-rerunfailures",
+  "uvloop; sys_platform != 'win32' and platform_python_implementation == 'CPython'",
 ]
 docs = [
   "zensical",

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: ...

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: ...

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<Bound<'py, Coroutine>>
+where
+    F: Future<Output = PyResult<T>> + 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<Bound<'py, Coroutine>>
+where
+    F: Future<Output = PyResult<T>> + 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<Bound<'py, Coroutine>> {
+    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<F, T>(runtime: Runtime, fut: F) -> PyResult<T>
+where
+    F: Future<Output = PyResult<T>> + 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()))?
+}