From 3918c148af9032f830fba7312628e89feb112a90 Mon Sep 17 00:00:00 2001 From: Brace Date: Tue, 15 Sep 2026 08:29:42 -0400 Subject: [PATCH] The README now leads with what you actually get: a Python asyncio surface with a Rust scheduler, real cancel/timeout, bounded admission, retries that re-call the callable --- README.md | 64 +++++++++++++--------- src/bridge.rs | 67 ++++++++++++++--------- src/handle.rs | 29 +++++----- src/journal.rs | 9 +--- src/policy.rs | 6 ++- src/runtime.rs | 144 ++++++++++++++++++++++++------------------------- src/stream.rs | 42 ++++++++------- 7 files changed, 201 insertions(+), 160 deletions(-) diff --git a/README.md b/README.md index 1aad19b..3c3a5da 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,21 @@ # rust-python-lab -A Tokio-backed **execution controller** for Python AI/agent apps. Python stays the authoring surface (`await handle`, `async with`, `async for`). Rust owns admission, concurrency, timeouts, cancellation, retries, streaming, and shutdown. Python's asyncio event loop still executes coroutine bodies — LLM SDKs, HTTP clients, DB drivers, and tools. +Keep writing agents in Python. Let **Rust run the scheduler**. -This is not a LangChain/LangGraph rewrite and not a Rust HTTP/LLM client. Rust waits on bridged asyncio futures; it does not step Python bytecode. +`rust_python_lab` is a Tokio-backed execution controller for asyncio apps. You `submit`, `gather`, `stream`, and `await` as usual. Rust owns concurrency, backpressure, timeouts, cancellation, retries, and shutdown — without rewriting your HTTP, DB, or LLM calls, and without locking you into LangChain or LangGraph. + +Python still steps coroutine bodies. Rust does not execute Python bytecode. That is the point: the control plane stays off the asyncio thread, so a slow or greedy task cannot stall admission, timers, or cancel. + +## Why this exists + +- **Python-native API.** `await handle`, `async with Runtime`, `async for` completions and tokens. `submit` is synchronous so you can `cancel()` before the first `await`. +- **Real cancel and timeout.** Dropping a Tokio waiter is not enough. The runtime cancels the **asyncio Task** the bridge created, so httpx connections and tool calls actually stop. +- **Limits that hold under load.** `max_concurrency` is a Tokio semaphore. The admission queue is bounded. `OnFull.REJECT` raises `QueueFull`; `OnFull.WAIT` blocks until a slot exists (and aborts on shutdown). +- **Retries that work.** Rust re-invokes the Python callable for a **new** coroutine. A consumed coroutine object is never replayed. +- **Fan-out without a Python control-plane.** `gather`, `as_completed`, and bounded `stream()` keep in-flight state in Rust. `stats()` is O(1) counters, not a walk of Python objects. +- **Clean shutdown.** `await rt.shutdown(timeout=…)` stops new work, drains, then cancels the rest. `async with` does this on exit. Optional `idle_ttl` closes admission after quiet. Further `submit` raises `RuntimeClosed`. +- **Use the libraries you already have.** httpx, OpenAI/Anthropic SDKs, databases, tools — they stay asyncio. Several `Runtime` instances can apply different policies (for example an LLM pool of 8 and a DB pool of 64). +- **Proven on sockets, not `asyncio.sleep`.** CI drives a local HTTP/SSE server through httpx: concurrency caps, handler abort on timeout/cancel, retries, backpressure, streams, shutdown, and a short soak. ## Execution split @@ -12,11 +25,9 @@ This is not a LangChain/LangGraph rewrite and not a Rust HTTP/LLM client. Rust w | **Python asyncio** | coroutine bodies, sockets, SDK calls, object lifetimes, GIL | | **Bridge** | `pyo3-async-runtimes`: capture `TaskLocals`, convert awaitables both ways, never hold the GIL across `.await` | -Cancel and timeout must cancel the **asyncio Task**, not only drop a Tokio waiter. Retries re-invoke the Python callable to get a new coroutine. - ## Setup -Requires Rust (stable) and Python 3.12. Local development uses 3.12.3. +Rust (stable) and Python 3.12+. ```bash python -m venv .venv @@ -40,7 +51,9 @@ async with Runtime( ) as rt: handle = rt.submit(call_llm, "hello", timeout=10.0) result = await handle - handle.cancel() + + slow = rt.submit(call_llm, "long job") + slow.cancel() results = await rt.gather( [Task(call_llm, prompt) for prompt in prompts], @@ -51,41 +64,42 @@ async with Runtime( async for token in rt.stream(llm_tokens, "hello"): ... - print(rt.stats()) - await rt.shutdown(timeout=5.0) + print(rt.stats()) # queued, in_flight, succeeded, failed, cancelled, timed_out, completed ``` -After close, `submit` raises `RuntimeClosed`. A full queue with `OnFull.REJECT` raises `QueueFull`. +`Handle.status()` is `queued`, `running`, `succeeded`, `failed`, `cancelled`, or `timed_out`. After close, `submit` raises `RuntimeClosed`. A full queue with `OnFull.REJECT` raises `QueueFull`. Without `async with`, call `await rt.shutdown(timeout=5.0)` yourself. ## Test -CI and local default: **localhost HTTP only**, no API keys, no public internet. +CI and the default local run use **localhost HTTP only** — no API keys, no public internet. ```bash pytest ``` -The suite drives a local `aiohttp` server (`tests/http_harness.py`) through **httpx** so policy is proven against real sockets: +The suite uses `tests/http_harness.py` (aiohttp on `127.0.0.1`) and **httpx** task bodies: -- concurrency cap, timeout, and cancel against `/work` and `/slow` -- retries against `/flaky` -- `gather` / `as_completed` / SSE `stream` against mixed latencies and `/sse` -- shutdown cancels in-flight `/slow` and rejects later submits +- `/work` — peak in-flight ≤ `max_concurrency` +- `/slow` — timeout and `handle.cancel()` abort the server handler +- `/flaky` — retries re-call the callable (extra server hits) +- `/sse` — `rt.stream` plus bounded-buffer backpressure +- mixed latencies in `gather` / `as_completed` +- shutdown cancels in-flight `/slow`; later `submit` is `RuntimeClosed` -Short soak (2k tiny GETs; `SOAK=1` runs 20k): +Soak (2k tiny GETs; `SOAK=1` runs 20k): ```bash pytest tests/test_soak.py SOAK=1 pytest tests/test_soak.py ``` -Agent-shaped example (tools + SSE “LLM”) against the same harness: +Agent-shaped example (parallel tools + SSE “LLM”) against the same harness: ```bash python examples/agent_tools.py ``` -Control-plane baseline (not a CI gate): `Runtime.gather` vs `asyncio.gather` + `asyncio.Semaphore` on local HTTP. +Scheduler baseline (not a CI gate): `Runtime.gather` vs `asyncio.gather` + `asyncio.Semaphore` on local HTTP. ```bash python scripts/bench_fanout.py @@ -93,17 +107,19 @@ python scripts/bench_fanout.py ### Optional live LLM -Skipped unless a key is present. CI never sets keys. Cheap smoke: 2–3 short prompts, `as_completed`, one token stream, one cancelled handle. +Skipped unless a key is set. CI never sets keys. Cheap smoke: two completions via `as_completed`, one token stream, one cancelled handle. Uses **httpx** only (no official SDK, no LangChain). ```bash python examples/live_llm.py ``` -Uses `OPENAI_API_KEY` (preferred) or `ANTHROPIC_API_KEY`. Optional `OPENAI_MODEL` / `ANTHROPIC_MODEL`. To also run it under pytest: `LIVE_LLM=1 pytest tests/test_examples.py`. Do not log secrets. +Set `OPENAI_API_KEY` (preferred) or `ANTHROPIC_API_KEY`. Optional `OPENAI_MODEL` / `ANTHROPIC_MODEL`. To run it under pytest: `LIVE_LLM=1 pytest tests/test_examples.py`. Do not log secrets. + +## Honest limits -## What this does not prove +These are deliberate, not missing checkboxes: -- Rust is not faster at token generation. The model and Python SDK still dominate latency. Benchmarks measure **scheduler** cost on local HTTP, not provider SLA. +- Rust does not make the model faster. Token generation and the Python SDK still dominate latency. `scripts/bench_fanout.py` measures **scheduler** cost on local HTTP, not provider SLA. - No crash-durable journal replay, no multi-process supervisor, no sandbox for tool code. -- Live tests need a key, cost a few tokens, and will be rate-limited. They are a smoke, not a soak. -- No built-in Rust HTTP or LLM adapters in this lab. Coroutine bodies stay on asyncio on purpose. +- Live LLM is an opt-in smoke. It needs a key, spends a few tokens, and will be rate-limited. +- No built-in Rust HTTP or LLM client. Coroutine bodies stay on asyncio so existing Python libraries keep working. diff --git a/src/bridge.rs b/src/bridge.rs index 662df04..28afa64 100644 --- a/src/bridge.rs +++ b/src/bridge.rs @@ -9,9 +9,31 @@ use pyo3::types::PyDict; use pyo3_async_runtimes::TaskLocals; use tokio_util::sync::CancellationToken; -static CANCELLED_EXC: OnceLock> = OnceLock::new(); +pub type PyValue = Py; +pub type PyValueResult = PyResult; -fn cancelled_type(py: Python<'_>) -> PyResult> { +type PyBound<'py> = Bound<'py, PyAny>; +type PyTaskCell = Option; +type PyTaskLock = Mutex; +pub type SharedPyTask = Arc; + +type TaskRx = oneshot::Receiver; +type ResultRx = oneshot::Receiver; +type TaskTx = oneshot::Sender; +type ResultTx = oneshot::Sender; + +pub struct BridgedTask { + pub task_rx: TaskRx, + pub result_rx: ResultRx, +} + +pub fn shared_py_task() -> SharedPyTask { + Arc::new(Mutex::new(None)) +} + +static CANCELLED_EXC: OnceLock = OnceLock::new(); + +fn cancelled_type(py: Python<'_>) -> PyResult> { if let Some(cls) = CANCELLED_EXC.get() { return Ok(cls.bind(py).clone()); } @@ -25,8 +47,8 @@ fn cancelled_type(py: Python<'_>) -> PyResult> { pub fn start_asyncio_task( py: Python<'_>, locals: &TaskLocals, - awaitable: Bound<'_, PyAny>, -) -> PyResult<(oneshot::Receiver>, oneshot::Receiver>>)> { + awaitable: PyBound<'_>, +) -> PyResult { let (task_tx, task_rx) = oneshot::channel(); let (result_tx, result_rx) = oneshot::channel(); let event_loop = locals.event_loop(py); @@ -45,13 +67,10 @@ pub fn start_asyncio_task( kwargs.set_item("context", locals.context(py))?; event_loop.call_method("call_soon_threadsafe", (callback,), Some(&kwargs))?; - Ok((task_rx, result_rx)) + Ok(BridgedTask { task_rx, result_rx }) } -pub fn cancel_asyncio_task( - event_loop: &Bound<'_, PyAny>, - task: &Bound<'_, PyAny>, -) -> PyResult<()> { +pub fn cancel_asyncio_task(event_loop: &PyBound<'_>, task: &PyBound<'_>) -> PyResult<()> { let cancel = task.getattr("cancel")?; event_loop.call_method1("call_soon_threadsafe", (cancel,))?; Ok(()) @@ -85,13 +104,13 @@ pub fn parse_timeout_secs(secs: f64) -> PyResult { /// Cancels the asyncio task if the Tokio waiter is dropped mid-flight. pub struct CancelOnDrop { - event_loop: Py, - py_task: Arc>>>, + event_loop: PyValue, + py_task: SharedPyTask, armed: bool, } impl CancelOnDrop { - pub fn new(event_loop: Py, py_task: Arc>>>) -> Self { + pub fn new(event_loop: PyValue, py_task: SharedPyTask) -> Self { Self { event_loop, py_task, @@ -113,7 +132,7 @@ impl Drop for CancelOnDrop { return; }; Python::attach(|py| { - let _ = cancel_asyncio_task(&self.event_loop.bind(py), &task.bind(py)); + let _ = cancel_asyncio_task(self.event_loop.bind(py), task.bind(py)); }); } } @@ -121,10 +140,10 @@ impl Drop for CancelOnDrop { pub async fn wait_with_policy( cancel: CancellationToken, timeout: Option, - result_rx: oneshot::Receiver>>, - event_loop: Py, - py_task: Arc>>>, -) -> PyResult> { + result_rx: ResultRx, + event_loop: PyValue, + py_task: SharedPyTask, +) -> PyValueResult { let mut guard = CancelOnDrop::new(event_loop, py_task); let result = match timeout { @@ -151,7 +170,7 @@ pub async fn wait_with_policy( result } -fn map_oneshot(result: Result>, oneshot::Canceled>) -> PyResult> { +fn map_oneshot(result: Result) -> PyValueResult { match result { Ok(inner) => inner, Err(_) => Err(cancelled_error()), @@ -160,10 +179,10 @@ fn map_oneshot(result: Result>, oneshot::Canceled>) -> PyResu #[pyclass] struct CaptureTask { - awaitable: Py, - event_loop: Py, - task_tx: Option>>, - result_tx: Option>>>, + awaitable: PyValue, + event_loop: PyValue, + task_tx: Option, + result_tx: Option, } #[pymethods] @@ -189,12 +208,12 @@ impl CaptureTask { #[pyclass] struct ResultCompleter { - tx: Option>>>, + tx: Option, } #[pymethods] impl ResultCompleter { - fn __call__(&mut self, task: Bound<'_, PyAny>) -> PyResult<()> { + fn __call__(&mut self, task: PyBound<'_>) -> PyResult<()> { let result = match task.call_method0("result") { Ok(value) => Ok(value.unbind()), Err(err) => Err(err), diff --git a/src/handle.rs b/src/handle.rs index 8f73b36..7d8a66e 100644 --- a/src/handle.rs +++ b/src/handle.rs @@ -1,28 +1,28 @@ -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use pyo3::prelude::*; use tokio_util::sync::CancellationToken; -use crate::bridge::cancel_asyncio_task; +use crate::bridge::{PyValue, SharedPyTask, cancel_asyncio_task}; use crate::journal::TaskSlot; /// Submitted work. Await this object to get the Python coroutine's result. #[pyclass(frozen)] pub struct Handle { - awaitable: Py, + awaitable: PyValue, slot: Arc, cancel: CancellationToken, - event_loop: Py, - py_task: Arc>>>, + event_loop: PyValue, + py_task: SharedPyTask, } impl Handle { pub fn new( - awaitable: Py, + awaitable: PyValue, slot: Arc, cancel: CancellationToken, - event_loop: Py, - py_task: Arc>>>, + event_loop: PyValue, + py_task: SharedPyTask, ) -> Self { Self { awaitable, @@ -33,7 +33,7 @@ impl Handle { } } - pub(crate) fn clone_awaitable(&self, py: Python<'_>) -> Py { + pub(crate) fn clone_awaitable(&self, py: Python<'_>) -> PyValue { self.awaitable.clone_ref(py) } } @@ -47,10 +47,13 @@ impl Handle { /// Cancel this task: Tokio token and the asyncio Task, if it has started. 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)) - }) { - cancel_asyncio_task(&self.event_loop.bind(py), &task.bind(py))?; + if let Some(task) = self + .py_task + .lock() + .ok() + .and_then(|guard| guard.as_ref().map(|task| task.clone_ref(py))) + { + cancel_asyncio_task(self.event_loop.bind(py), task.bind(py))?; } Ok(()) } diff --git a/src/journal.rs b/src/journal.rs index ec086db..4ea181b 100644 --- a/src/journal.rs +++ b/src/journal.rs @@ -1,5 +1,5 @@ use std::sync::Arc; -use std::sync::atomic::{AtomicU64, AtomicU8, Ordering}; +use std::sync::atomic::{AtomicU8, AtomicU64, Ordering}; #[repr(u8)] #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -64,12 +64,7 @@ impl TaskSlot { fn transition(&self, from: TaskState, to: TaskState) -> bool { self.state - .compare_exchange( - from as u8, - to as u8, - Ordering::AcqRel, - Ordering::Acquire, - ) + .compare_exchange(from as u8, to as u8, Ordering::AcqRel, Ordering::Acquire) .is_ok() } diff --git a/src/policy.rs b/src/policy.rs index d4b7839..25b317b 100644 --- a/src/policy.rs +++ b/src/policy.rs @@ -6,8 +6,10 @@ use pyo3::prelude::*; #[pyclass(eq, eq_int, frozen, from_py_object, name = "OnFull")] #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum OnFull { - REJECT, - WAIT, + #[pyo3(name = "REJECT")] + Reject, + #[pyo3(name = "WAIT")] + Wait, } #[pyclass(frozen, from_py_object)] diff --git a/src/runtime.rs b/src/runtime.rs index d3a4937..0259dd5 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1,25 +1,26 @@ +use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; use std::time::Duration; use pyo3::exceptions::{PyStopAsyncIteration, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyDict, PyList, PyTuple}; +use pyo3_async_runtimes::TaskLocals; use pyo3_async_runtimes::tokio::{ future_into_py_with_locals, get_current_locals, get_runtime, scope, }; -use pyo3_async_runtimes::TaskLocals; use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore}; use tokio_util::sync::CancellationToken; use crate::bridge::{ - cancelled_error, is_cancelled, parse_timeout_secs, start_asyncio_task, wait_with_policy, + PyValue, PyValueResult, SharedPyTask, cancelled_error, is_cancelled, parse_timeout_secs, + shared_py_task, start_asyncio_task, wait_with_policy, }; use crate::error::{QueueFull, RuntimeClosed}; use crate::handle::Handle; use crate::journal::{Journal, TaskSlot, TaskState}; use crate::policy::{OnFull, RetryConfig, RetryPolicy}; -use crate::stream::{completion_channel, item_channel, CompletionStream, ItemStream}; +use crate::stream::{CompletionStream, ItemResult, ItemStream, completion_channel, item_channel}; use crate::task::Task; use futures_util::stream::{FuturesUnordered, StreamExt}; @@ -101,7 +102,7 @@ impl Runtime { #[pyo3(signature = ( max_concurrency=32, queue_capacity=128, - on_full=OnFull::REJECT, + on_full=OnFull::Reject, default_timeout=None, retry=None, idle_ttl=None @@ -168,17 +169,17 @@ impl Runtime { }; let args = args.unbind(); let cancel = self.inner.shutdown.child_token(); - let py_task = Arc::new(Mutex::new(None)); + let py_task = shared_py_task(); let event_loop = locals.event_loop(py).unbind(); let admission_permit = match self.inner.on_full { - OnFull::REJECT => match Arc::clone(&self.inner.admission).try_acquire_owned() { + 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, + OnFull::Wait => None, }; let work = WorkGuard::try_begin(Arc::clone(&self.inner))?; @@ -245,42 +246,38 @@ impl Runtime { 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); - } + 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))?, - } + } + 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()) - }) - }, - ) + } + Ok(list.unbind().into_any()) + }) + }) } /// Yield `Completion` items as tasks finish (out of order). @@ -337,16 +334,16 @@ impl Runtime { }; let args = args.unbind(); let cancel = self.inner.shutdown.child_token(); - let py_task = Arc::new(Mutex::new(None)); + let py_task = shared_py_task(); 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() { + 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, + OnFull::Wait => None, }; let work = WorkGuard::try_begin(Arc::clone(&self.inner))?; @@ -408,10 +405,7 @@ impl Runtime { self.inner.is_closed() } - fn __aenter__<'py>( - slf: Bound<'py, Self>, - py: Python<'py>, - ) -> PyResult> { + fn __aenter__<'py>(slf: Bound<'py, Self>, py: Python<'py>) -> PyResult> { slf.get().inner.ensure_open()?; let slf = slf.unbind(); let locals = get_current_locals(py)?; @@ -434,18 +428,18 @@ impl Runtime { struct RunArgs { inner: Arc, locals: TaskLocals, - func: Py, + func: PyValue, args: Py, admission_permit: Option, cancel: CancellationToken, timeout: Option, task_id: u64, - py_task: Arc>>>, - event_loop: Py, + py_task: SharedPyTask, + event_loop: PyValue, slot: Arc, } -fn state_for(result: &PyResult>) -> TaskState { +fn state_for(result: &PyValueResult) -> TaskState { match result { Ok(_) => TaskState::Succeeded, Err(err) => state_for_err(err), @@ -464,7 +458,7 @@ fn state_for_err(err: &PyErr) -> TaskState { }) } -async fn run_task(args: RunArgs) -> PyResult> { +async fn run_task(args: RunArgs) -> PyValueResult { let slot = Arc::clone(&args.slot); let inner = Arc::clone(&args.inner); let result = run_task_inner(args).await; @@ -472,7 +466,7 @@ async fn run_task(args: RunArgs) -> PyResult> { result } -async fn run_task_inner(args: RunArgs) -> PyResult> { +async fn run_task_inner(args: RunArgs) -> PyValueResult { let retry = args.inner.retry; let _admission = match args.admission_permit { Some(permit) => permit, @@ -513,8 +507,7 @@ async fn run_task_inner(args: RunArgs) -> PyResult> { { Ok(value) => return Ok(value), Err(err) => { - let more = attempt + 1 < attempts - && Python::attach(|py| !is_cancelled(py, &err)); + let more = attempt + 1 < attempts && Python::attach(|py| !is_cancelled(py, &err)); last_err = Some(err); if !more { break; @@ -527,17 +520,17 @@ async fn run_task_inner(args: RunArgs) -> PyResult> { async fn run_once( locals: &TaskLocals, - func: &Py, + func: &PyValue, args: &Py, cancel: CancellationToken, timeout: Option, - py_task: &Arc>>>, - event_loop: &Py, -) -> PyResult> { + py_task: &SharedPyTask, + event_loop: &PyValue, +) -> PyValueResult { let (task_rx, result_rx, event_loop) = Python::attach(|py| -> PyResult<_> { let coro = func.bind(py).call(args.bind(py), None)?; - let (task_rx, result_rx) = start_asyncio_task(py, locals, coro)?; - Ok((task_rx, result_rx, event_loop.clone_ref(py))) + let bridged = start_asyncio_task(py, locals, coro)?; + Ok((bridged.task_rx, bridged.result_rx, event_loop.clone_ref(py))) })?; let started = tokio::select! { @@ -593,25 +586,28 @@ fn submit_spec(rt: &Bound<'_, Runtime>, py: Python<'_>, spec: &Py) -> PyRe struct StreamArgs { inner: Arc, locals: TaskLocals, - func: Py, + func: PyValue, args: Py, admission_permit: Option, cancel: CancellationToken, timeout: Option, - py_task: Arc>>>, - event_loop: Py, + py_task: SharedPyTask, + event_loop: PyValue, slot: Arc, - tx: tokio::sync::mpsc::Sender>>, + tx: tokio::sync::mpsc::Sender, } async fn run_stream(args: StreamArgs) -> PyResult<()> { let slot = Arc::clone(&args.slot); let inner = Arc::clone(&args.inner); let result = run_stream_inner(args).await; - inner.journal.finish(&slot, match &result { - Ok(()) => TaskState::Succeeded, - Err(err) => state_for_err(err), - }); + inner.journal.finish( + &slot, + match &result { + Ok(()) => TaskState::Succeeded, + Err(err) => state_for_err(err), + }, + ); result } @@ -640,8 +636,12 @@ async fn run_stream_inner(args: StreamArgs) -> PyResult<()> { let started = Python::attach(|py| -> PyResult<_> { let next = agen.bind(py).call_method0("__anext__")?; - let (task_rx, result_rx) = start_asyncio_task(py, &args.locals, next)?; - Ok((task_rx, result_rx, args.event_loop.clone_ref(py))) + let bridged = start_asyncio_task(py, &args.locals, next)?; + Ok(( + bridged.task_rx, + bridged.result_rx, + args.event_loop.clone_ref(py), + )) }); let (task_rx, result_rx, event_loop) = match started { Ok(parts) => parts, diff --git a/src/stream.rs b/src/stream.rs index 7c57a68..8f48606 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -3,9 +3,22 @@ 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::sync::{Mutex as AsyncMutex, mpsc}; use tokio_util::sync::CancellationToken; +use crate::bridge::{PyValue, PyValueResult}; + +pub type ItemResult = PyValueResult; +pub type CompletionItem = (usize, ItemResult); + +type ItemTx = mpsc::Sender; +type ItemRx = mpsc::Receiver; +type CompletionTx = mpsc::Sender; +type CompletionRx = mpsc::Receiver; +type ChannelRx = mpsc::Receiver; +type LockedRx = AsyncMutex>>; +type SharedRx = Arc>; + /// One finished gather/as_completed item, in completion order. #[pyclass(frozen)] pub struct Completion { @@ -14,11 +27,11 @@ pub struct Completion { #[pyo3(get)] ok: bool, #[pyo3(get)] - value: Py, + value: PyValue, } impl Completion { - pub fn new(py: Python<'_>, index: usize, result: PyResult>) -> Self { + pub fn new(py: Python<'_>, index: usize, result: ItemResult) -> Self { match result { Ok(value) => Self { index, @@ -36,13 +49,13 @@ impl Completion { #[pyclass] pub struct CompletionStream { - rx: Arc>)>>>>, + rx: SharedRx, } impl CompletionStream { - pub fn new(rx: mpsc::Receiver<(usize, PyResult>)>) -> Self { + pub fn new(rx: CompletionRx) -> Self { Self { - rx: Arc::new(Mutex::new(Some(rx))), + rx: Arc::new(AsyncMutex::new(Some(rx))), } } } @@ -74,18 +87,16 @@ impl CompletionStream { } } -pub type ItemResult = PyResult>; - #[pyclass] pub struct ItemStream { - rx: Arc>>>, + rx: SharedRx, cancel: CancellationToken, } impl ItemStream { - pub fn new(rx: mpsc::Receiver, cancel: CancellationToken) -> Self { + pub fn new(rx: ItemRx, cancel: CancellationToken) -> Self { Self { - rx: Arc::new(Mutex::new(Some(rx))), + rx: Arc::new(AsyncMutex::new(Some(rx))), cancel, } } @@ -121,15 +132,10 @@ impl ItemStream { } } -pub fn item_channel(buffer: usize) -> (mpsc::Sender, mpsc::Receiver) { +pub fn item_channel(buffer: usize) -> (ItemTx, ItemRx) { mpsc::channel(buffer.max(1)) } -pub fn completion_channel( - n: usize, -) -> ( - mpsc::Sender<(usize, PyResult>)>, - mpsc::Receiver<(usize, PyResult>)>, -) { +pub fn completion_channel(n: usize) -> (CompletionTx, CompletionRx) { mpsc::channel(n.max(1)) }