From f353108af69958064feb69e36427b51019b99da0 Mon Sep 17 00:00:00 2001 From: Brace Date: Mon, 14 Sep 2026 09:08:15 -0400 Subject: [PATCH] gather-streaming --- Cargo.lock | 1 + Cargo.toml | 1 + src/handle.rs | 6 +- src/lib.rs | 6 + src/runtime.rs | 302 ++++++++++++++++++++++++++++++++++++++++-- src/stream.rs | 135 +++++++++++++++++++ src/task.rs | 23 ++++ tests/http_harness.py | 28 +++- tests/test_fanout.py | 77 +++++++++++ 9 files changed, 565 insertions(+), 14 deletions(-) create mode 100644 src/stream.rs create mode 100644 src/task.rs create mode 100644 tests/test_fanout.py diff --git a/Cargo.lock b/Cargo.lock index 4e8470c..2a10bea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -183,6 +183,7 @@ name = "rust-python-lab" version = "0.1.0" dependencies = [ "futures-channel", + "futures-util", "pyo3", "pyo3-async-runtimes", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 8ef04cb..829bf72 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,3 +13,4 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time"] } tokio-util = { version = "0.7", features = ["rt"] } futures-channel = "0.3" +futures-util = { version = "0.3", default-features = false, features = ["std", "async-await"] } diff --git a/src/handle.rs b/src/handle.rs index 351f35b..ff79fa3 100644 --- a/src/handle.rs +++ b/src/handle.rs @@ -35,6 +35,10 @@ impl Handle { py_task, } } + + pub(crate) fn clone_awaitable(&self, py: Python<'_>) -> Py { + self.awaitable.clone_ref(py) + } } #[pymethods] @@ -44,7 +48,7 @@ impl Handle { } /// Cancel this task: Tokio token and the asyncio Task, if it has started. - fn cancel(&self, py: Python<'_>) -> PyResult<()> { + pub(crate) fn cancel(&self, py: Python<'_>) -> PyResult<()> { self.cancel.cancel(); if let Some(task) = self.py_task.lock().ok().and_then(|guard| { guard.as_ref().map(|task| task.clone_ref(py)) diff --git a/src/lib.rs b/src/lib.rs index 5a9ff8b..abeb1b3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,8 @@ mod handle; mod journal; mod policy; mod runtime; +mod stream; +mod task; use pyo3::prelude::*; @@ -20,4 +22,8 @@ mod rust_python_lab { use crate::policy::RetryPolicy; #[pymodule_export] use crate::runtime::Runtime; + #[pymodule_export] + use crate::stream::Completion; + #[pymodule_export] + use crate::task::Task; } diff --git a/src/runtime.rs b/src/runtime.rs index c49b625..ca92bf5 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -2,9 +2,9 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; -use pyo3::exceptions::PyValueError; +use pyo3::exceptions::{PyStopAsyncIteration, PyValueError}; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyTuple}; +use pyo3::types::{PyDict, PyList, PyTuple}; use pyo3_async_runtimes::tokio::{ future_into_py_with_locals, get_current_locals, get_runtime, scope, }; @@ -17,6 +17,9 @@ use crate::error::QueueFull; use crate::handle::Handle; use crate::journal::{Journal, TaskState}; use crate::policy::{OnFull, RetryConfig, RetryPolicy}; +use crate::stream::{completion_channel, item_channel, CompletionStream, ItemStream}; +use crate::task::Task; +use futures_util::stream::{FuturesUnordered, StreamExt}; struct Inner { journal: Arc>, @@ -159,6 +162,171 @@ impl Runtime { )) } + /// Wait for all tasks, returning results in input order. + #[pyo3(signature = (tasks, return_exceptions=false))] + fn gather<'py>( + slf: &Bound<'py, Self>, + py: Python<'py>, + tasks: Bound<'py, PyAny>, + return_exceptions: bool, + ) -> PyResult> { + let locals = get_current_locals(py)?; + let specs = extract_tasks(&tasks)?; + let mut handles = Vec::with_capacity(specs.len()); + let mut futs = Vec::with_capacity(specs.len()); + for spec in specs { + let handle = submit_spec(slf, py, &spec)?; + let fut = pyo3_async_runtimes::into_future_with_locals( + &locals, + handle.clone_awaitable(py).into_bound(py), + )?; + handles.push(handle); + futs.push(fut); + } + future_into_py_with_locals( + py, + locals, + async move { + let n = futs.len(); + let mut pending = FuturesUnordered::new(); + for (index, fut) in futs.into_iter().enumerate() { + pending.push(async move { (index, fut.await) }); + } + let mut slots: Vec>>> = (0..n).map(|_| None).collect(); + while let Some((index, result)) = pending.next().await { + match result { + Ok(value) => slots[index] = Some(Ok(value)), + Err(err) if return_exceptions => slots[index] = Some(Err(err)), + Err(err) => { + Python::attach(|py| { + for handle in &handles { + let _ = handle.cancel(py); + } + }); + return Err(err); + } + } + } + Python::attach(|py| { + let list = PyList::empty(py); + for slot in slots { + match slot.unwrap_or_else(|| Err(cancelled_error())) { + Ok(value) => list.append(value.bind(py))?, + Err(err) => list.append(err.value(py))?, + } + } + Ok(list.unbind().into_any()) + }) + }, + ) + } + + /// Yield `Completion` items as tasks finish (out of order). + fn as_completed<'py>( + slf: &Bound<'py, Self>, + py: Python<'py>, + tasks: Bound<'py, PyAny>, + ) -> PyResult> { + let locals = get_current_locals(py)?; + let specs = extract_tasks(&tasks)?; + let (tx, rx) = completion_channel(specs.len()); + let mut waiters = Vec::with_capacity(specs.len()); + for (index, spec) in specs.into_iter().enumerate() { + let handle = submit_spec(slf, py, &spec)?; + let fut = pyo3_async_runtimes::into_future_with_locals( + &locals, + handle.clone_awaitable(py).into_bound(py), + )?; + waiters.push((index, fut)); + } + get_runtime().spawn(async move { + let mut set = FuturesUnordered::new(); + for (index, fut) in waiters { + set.push(async move { (index, fut.await) }); + } + while let Some(item) = set.next().await { + if tx.send(item).await.is_err() { + break; + } + } + }); + Bound::new(py, CompletionStream::new(rx)) + } + + /// Stream items from a Python async generator. `buffer` bounds in-flight yields. + #[pyo3(signature = (func, *args, timeout=None, buffer=8))] + fn stream( + &self, + py: Python<'_>, + func: Py, + args: Bound<'_, PyTuple>, + timeout: Option, + buffer: usize, + ) -> PyResult { + if buffer == 0 { + return Err(PyValueError::new_err("buffer must be at least 1")); + } + let locals = get_current_locals(py)?; + let timeout = match timeout { + Some(secs) => Some(parse_timeout_secs(secs)?), + None => self.inner.default_timeout, + }; + let args = args.unbind(); + let cancel = CancellationToken::new(); + let py_task = Arc::new(Mutex::new(None)); + let event_loop = locals.event_loop(py).unbind(); + let (tx, rx) = item_channel(buffer); + + let admission_permit = match self.inner.on_full { + OnFull::REJECT => match Arc::clone(&self.inner.admission).try_acquire_owned() { + Ok(permit) => Some(permit), + Err(_) => return Err(QueueFull::new_err("admission queue is full")), + }, + OnFull::WAIT => None, + }; + + let task_id = self.inner.next_id.fetch_add(1, Ordering::Relaxed); + self.inner + .journal + .lock() + .expect("journal mutex") + .insert_queued(task_id); + + let journal = Arc::clone(&self.inner.journal); + let concurrency = Arc::clone(&self.inner.concurrency); + let admission = Arc::clone(&self.inner.admission); + let cancel_for_task = cancel.clone(); + let event_loop_for_run = event_loop.clone_ref(py); + + get_runtime().spawn(scope(locals.clone(), async move { + let result = run_stream(StreamArgs { + locals, + func, + args, + admission, + admission_permit, + concurrency, + cancel: cancel_for_task, + timeout, + task_id, + py_task, + event_loop: event_loop_for_run, + journal: Arc::clone(&journal), + tx, + }) + .await; + let state = match &result { + Ok(()) => TaskState::Succeeded, + Err(err) => state_for_err(err), + }; + if let Ok(mut journal) = journal.lock() { + journal.finish(task_id, state); + } + })); + + Ok(ItemStream::new(rx, cancel)) + } + fn stats<'py>(&self, py: Python<'py>) -> PyResult> { let stats = self .inner @@ -197,18 +365,22 @@ struct RunArgs { fn state_for(result: &PyResult>) -> TaskState { match result { Ok(_) => TaskState::Succeeded, - Err(err) => Python::attach(|py| { - if err.is_instance_of::(py) { - TaskState::TimedOut - } else if is_cancelled(py, err) { - TaskState::Cancelled - } else { - TaskState::Failed - } - }), + Err(err) => state_for_err(err), } } +fn state_for_err(err: &PyErr) -> TaskState { + Python::attach(|py| { + if err.is_instance_of::(py) { + TaskState::TimedOut + } else if is_cancelled(py, err) { + TaskState::Cancelled + } else { + TaskState::Failed + } + }) +} + fn is_cancelled(py: Python<'_>, err: &PyErr) -> bool { py.import("asyncio") .and_then(|asyncio| asyncio.getattr("CancelledError")) @@ -319,3 +491,111 @@ async fn acquire_permit( } } } + +fn extract_tasks(tasks: &Bound<'_, PyAny>) -> PyResult>> { + let mut out = Vec::new(); + for item in tasks.try_iter()? { + out.push(item?.extract()?); + } + Ok(out) +} + +fn submit_spec(rt: &Bound<'_, Runtime>, py: Python<'_>, spec: &Py) -> PyResult { + let task = spec.bind(py).borrow(); + rt.get().submit( + py, + task.func.clone_ref(py), + task.args.bind(py).clone(), + task.timeout, + ) +} + +struct StreamArgs { + locals: TaskLocals, + func: Py, + args: Py, + admission: Arc, + admission_permit: Option, + concurrency: Arc, + cancel: CancellationToken, + timeout: Option, + task_id: u64, + py_task: Arc>>>, + event_loop: Py, + journal: Arc>, + tx: tokio::sync::mpsc::Sender>>, +} + +async fn run_stream(args: StreamArgs) -> PyResult<()> { + let _admission = match args.admission_permit { + Some(permit) => permit, + None => acquire_permit(args.admission, &args.cancel).await?, + }; + if args.cancel.is_cancelled() { + return Err(cancelled_error()); + } + let _concurrency = acquire_permit(args.concurrency, &args.cancel).await?; + if args.cancel.is_cancelled() { + return Err(cancelled_error()); + } + if let Ok(mut journal) = args.journal.lock() { + journal.mark_running(args.task_id); + } + + let agen = Python::attach(|py| { + args.func + .bind(py) + .call(args.args.bind(py), None) + .map(|value| value.unbind()) + })?; + + loop { + if args.cancel.is_cancelled() { + return Err(cancelled_error()); + } + if args.tx.is_closed() { + return Ok(()); + } + + let (task_rx, result_rx) = Python::attach(|py| { + let next = agen.bind(py).call_method0("__anext__")?; + start_asyncio_task(py, &args.locals, next) + })?; + + let started = tokio::select! { + _ = args.cancel.cancelled() => None, + task = task_rx => task.ok(), + }; + let Some(task) = started else { + return Err(cancelled_error()); + }; + if let Ok(mut slot) = args.py_task.lock() { + *slot = Some(task); + } + + let event_loop = Python::attach(|py| args.event_loop.clone_ref(py)); + match wait_with_policy( + args.cancel.clone(), + args.timeout, + result_rx, + event_loop, + Arc::clone(&args.py_task), + ) + .await + { + Ok(item) => { + if args.tx.send(Ok(item)).await.is_err() { + return Ok(()); + } + } + Err(err) => { + let stop = Python::attach(|py| err.is_instance_of::(py)); + if stop { + return Ok(()); + } + let _ = args.tx.send(Err(err)).await; + return Ok(()); + } + } + } +} diff --git a/src/stream.rs b/src/stream.rs new file mode 100644 index 0000000..7c57a68 --- /dev/null +++ b/src/stream.rs @@ -0,0 +1,135 @@ +use std::sync::Arc; + +use pyo3::exceptions::PyStopAsyncIteration; +use pyo3::prelude::*; +use pyo3_async_runtimes::tokio::{future_into_py_with_locals, get_current_locals}; +use tokio::sync::{Mutex, mpsc}; +use tokio_util::sync::CancellationToken; + +/// One finished gather/as_completed item, in completion order. +#[pyclass(frozen)] +pub struct Completion { + #[pyo3(get)] + index: usize, + #[pyo3(get)] + ok: bool, + #[pyo3(get)] + value: Py, +} + +impl Completion { + pub fn new(py: Python<'_>, index: usize, result: PyResult>) -> Self { + match result { + Ok(value) => Self { + index, + ok: true, + value, + }, + Err(err) => Self { + index, + ok: false, + value: err.value(py).clone().into_any().unbind(), + }, + } + } +} + +#[pyclass] +pub struct CompletionStream { + rx: Arc>)>>>>, +} + +impl CompletionStream { + pub fn new(rx: mpsc::Receiver<(usize, PyResult>)>) -> Self { + Self { + rx: Arc::new(Mutex::new(Some(rx))), + } + } +} + +#[pymethods] +impl CompletionStream { + fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { + let locals = get_current_locals(py)?; + let rx = Arc::clone(&self.rx); + future_into_py_with_locals(py, locals, async move { + let mut guard = rx.lock().await; + let Some(channel) = guard.as_mut() else { + return Err(PyStopAsyncIteration::new_err(())); + }; + match channel.recv().await { + Some((index, result)) => Python::attach(|py| { + Ok(Bound::new(py, Completion::new(py, index, result))?.unbind()) + }), + None => { + *guard = None; + Err(PyStopAsyncIteration::new_err(())) + } + } + }) + } +} + +pub type ItemResult = PyResult>; + +#[pyclass] +pub struct ItemStream { + rx: Arc>>>, + cancel: CancellationToken, +} + +impl ItemStream { + pub fn new(rx: mpsc::Receiver, cancel: CancellationToken) -> Self { + Self { + rx: Arc::new(Mutex::new(Some(rx))), + cancel, + } + } +} + +#[pymethods] +impl ItemStream { + fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn cancel(&self) { + self.cancel.cancel(); + } + + fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { + let locals = get_current_locals(py)?; + let rx = Arc::clone(&self.rx); + future_into_py_with_locals(py, locals, async move { + let mut guard = rx.lock().await; + let Some(channel) = guard.as_mut() else { + return Err(PyStopAsyncIteration::new_err(())); + }; + match channel.recv().await { + Some(Ok(value)) => Ok(value), + Some(Err(err)) => Err(err), + None => { + *guard = None; + Err(PyStopAsyncIteration::new_err(())) + } + } + }) + } +} + +pub fn item_channel(buffer: usize) -> (mpsc::Sender, mpsc::Receiver) { + mpsc::channel(buffer.max(1)) +} + +pub fn completion_channel( + n: usize, +) -> ( + mpsc::Sender<(usize, PyResult>)>, + mpsc::Receiver<(usize, PyResult>)>, +) { + mpsc::channel(n.max(1)) +} diff --git a/src/task.rs b/src/task.rs new file mode 100644 index 0000000..d49b4f1 --- /dev/null +++ b/src/task.rs @@ -0,0 +1,23 @@ +use pyo3::prelude::*; +use pyo3::types::PyTuple; + +/// Spec for work to submit: callable plus args. Re-invoked on retry. +#[pyclass(frozen)] +pub struct Task { + pub func: Py, + pub args: Py, + pub timeout: Option, +} + +#[pymethods] +impl Task { + #[new] + #[pyo3(signature = (func, *args, timeout=None))] + fn new(func: Py, args: Bound<'_, PyTuple>, timeout: Option) -> Self { + Self { + func, + args: args.unbind(), + timeout, + } + } +} diff --git a/tests/http_harness.py b/tests/http_harness.py index 109725c..53b0136 100644 --- a/tests/http_harness.py +++ b/tests/http_harness.py @@ -17,6 +17,7 @@ class Counters: slow_exited: int = 0 cancelled_handlers: int = 0 flaky_hits: int = 0 + sse_sent: int = 0 def _enter(counters: Counters) -> None: @@ -26,11 +27,12 @@ def _enter(counters: Counters) -> None: def create_app(counters: Counters) -> web.Application: - async def work(_request: web.Request) -> web.Response: + async def work(request: web.Request) -> web.Response: + delay = float(request.query.get("delay", "0.15")) counters.work_hits += 1 _enter(counters) try: - await asyncio.sleep(0.15) + await asyncio.sleep(delay) return web.Response(text="ok") finally: counters.in_flight -= 1 @@ -57,10 +59,22 @@ async def flaky(request: web.Request) -> web.Response: return web.Response(status=500, text="fail") return web.Response(text="ok") + async def sse(_request: web.Request) -> web.StreamResponse: + response = web.StreamResponse() + response.headers["Content-Type"] = "text/event-stream" + await response.prepare(_request) + for i in range(8): + await response.write(f"data: token-{i}\n\n".encode()) + counters.sse_sent += 1 + await asyncio.sleep(0.02) + await response.write_eof() + return response + app = web.Application() app.router.add_get("/work", work) app.router.add_get("/slow", slow) app.router.add_get("/flaky", flaky) + app.router.add_get("/sse", sse) return app @@ -78,3 +92,13 @@ async def http_get(url: str) -> str: response = await client.get(url, timeout=None) response.raise_for_status() return response.text + + +async def sse_tokens(url: str): + async with httpx.AsyncClient() as client: + async with client.stream("GET", url, timeout=None) as response: + response.raise_for_status() + async for line in response.aiter_lines(): + if line.startswith("data: "): + yield line[6:] + diff --git a/tests/test_fanout.py b/tests/test_fanout.py new file mode 100644 index 0000000..d657e6d --- /dev/null +++ b/tests/test_fanout.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from rust_python_lab import Runtime, Task +from tests.http_harness import Harness, http_get, sse_tokens + + +async def test_gather_mixed_http(http_harness: Harness): + rt = Runtime(max_concurrency=8) + results = await rt.gather( + [ + Task(http_get, http_harness.url("/work?delay=0.05")), + Task(http_get, http_harness.url("/flaky?fail_times=1")), + Task(http_get, http_harness.url("/work?delay=0.02")), + ], + return_exceptions=True, + ) + assert results[0] == "ok" + assert isinstance(results[1], Exception) + assert results[2] == "ok" + + +async def test_as_completed_out_of_order(http_harness: Harness): + rt = Runtime(max_concurrency=8) + order = [] + async for event in rt.as_completed( + [ + Task(http_get, http_harness.url("/work?delay=0.2")), + Task(http_get, http_harness.url("/work?delay=0.02")), + ] + ): + order.append(event.index) + assert event.ok + assert event.value == "ok" + assert order[0] == 1 + assert sorted(order) == [0, 1] + + +async def test_stream_sse_tokens(http_harness: Harness): + rt = Runtime(max_concurrency=4) + tokens = [] + async for token in rt.stream(sse_tokens, http_harness.url("/sse")): + tokens.append(token) + assert tokens == [f"token-{i}" for i in range(8)] + assert http_harness.counters.sse_sent == 8 + + +async def test_stream_backpressure(): + produced = [] + + async def gen(): + for i in range(10): + produced.append(i) + yield i + + rt = Runtime() + stream = rt.stream(gen, buffer=1) + first = await stream.__anext__() + assert first == 0 + await asyncio.sleep(0.05) + assert produced[-1] <= 2 + rest = [item async for item in stream] + assert [first, *rest] == list(range(10)) + + +async def test_gather_raises_without_return_exceptions(http_harness: Harness): + rt = Runtime(max_concurrency=4) + with pytest.raises(Exception): + await rt.gather( + [ + Task(http_get, http_harness.url("/flaky?fail_times=9")), + Task(http_get, http_harness.url("/work?delay=0.02")), + ] + )