diff --git a/rust/src/database/postgres/helpers.rs b/rust/src/database/postgres/helpers.rs new file mode 100644 index 0000000000..fcc89f5c40 --- /dev/null +++ b/rust/src/database/postgres/helpers.rs @@ -0,0 +1,247 @@ +//! Extension traits for driving [`tokio_postgres`] futures to completion from +//! the synchronous, GIL-holding Python methods. +//! +//! Both [`BlockingPostgres`] and [`BlockingPostgresStream`] release the GIL +//! (`py.detach`) while blocking on the shared tokio runtime, so other Python +//! threads can make progress (and so the runtime's own connection task can run) +//! while we wait. The [`Ungil`] bounds are what let us hand the future across +//! the `detach` boundary. +//! +//! The stream helper is generic over the underlying stream type rather than +//! hard-wired to [`tokio_postgres::RowStream`]. In production it is always used +//! with a `RowStream`, but keeping it generic lets the cursor state machine +//! (which uses these helpers) be unit-tested against an in-memory fake stream +//! with no live database — see [`BlockingPostgresStream`]'s tests. +//! +//! ## Why polling a `RowStream` off the runtime is safe +//! +//! [`BlockingPostgresStream::get_next_if_ready`] polls the stream once on the +//! calling thread *without* entering the runtime. This is sound specifically +//! because a [`tokio_postgres::RowStream`] poll only reads from an in-memory +//! channel that the connection task (running on the runtime) feeds — it never +//! touches the tokio reactor or a timer, so it can't panic with "no reactor +//! running" and a single `Pending` poll genuinely means "nothing buffered yet". +//! A different stream that needs a reactor on the polling thread would *not* be +//! safe to use here, even though the generic bounds would accept it. + +use std::{future::Future, pin::Pin}; + +use futures::{stream::Fuse, FutureExt, StreamExt}; +use pyo3::{marker::Ungil, PyResult, Python}; +use tokio::runtime::Handle; + +use crate::database::postgres::pg_err_to_py; + +/// Block on a future on the shared runtime, releasing the GIL while we wait. +pub trait BlockingPostgres +where + Self: Future + Sized + Send + Ungil, + Self::Output: Ungil + Send, +{ + /// Drive `self` to completion on the shared runtime `handle`, returning its + /// output. Releases the GIL for the duration so the wait doesn't block + /// other Python threads. + /// + /// `handle` must be a handle to the extension's shared runtime (see + /// `crate::tokio_runtime`); the connection carries one from `connect` time. + /// The blocking wait runs on the calling (Python) thread, never on a + /// runtime worker — see this module's docs for why that can't deadlock. + fn block_on(self, py: Python<'_>, handle: &Handle) -> Self::Output { + py.detach(|| handle.block_on(self)) + } +} + +/// Same as [`BlockingPostgres`], but for futures that yield a +/// [`tokio_postgres::Result`], mapping any error into a Python exception. +pub trait BlockingPostgresResult +where + Self: Future> + Sized + Send + Ungil, + Self::Output: Ungil + Send, +{ + /// Block on `self` and convert a Postgres error into a `PyErr`. + fn block_on_result(self, py: Python<'_>, handle: &Handle) -> PyResult { + self.block_on(py, handle).map_err(pg_err_to_py) + } +} + +// Blanket impls: every suitable future automatically gets `block_on` / +// `block_on_result`, so callers can write `fut.block_on(py)` directly. +impl BlockingPostgres for F +where + F: Future + Sized + Send + Ungil, + F::Output: Ungil + Send, +{ +} +impl BlockingPostgresResult for F +where + F: Future> + Sized + Send + Ungil, + F::Output: Ungil + Send, +{ +} + +/// Pull items from a [`Fuse`]d stream from synchronous Python code, blocking on +/// the shared runtime only when the next item isn't already buffered. +/// +/// Implemented for any pinned, fused stream (`Pin<&mut Fuse>`) whose items +/// can cross the GIL-release boundary. In production `S` is +/// [`tokio_postgres::RowStream`]; the generic bound is what lets the cursor +/// logic be tested against an in-memory fake. +/// +/// The [`Fuse`] is *required* by the impl (the trait is implemented only for +/// `Pin<&mut Fuse>`), not merely assumed. This matters because +/// [`Self::get_next_if_ready`] may poll the stream again after it has finished: +/// a bare `Stream` is free to panic if polled past completion, whereas a fused +/// stream simply keeps yielding `None`. So repeated `get_next_if_ready` / +/// `block_on_next` calls after exhaustion are safe by construction. +pub trait BlockingPostgresStream +where + Self: futures::Stream + Sized + Send + Ungil + Unpin, + Self::Item: Ungil + Send, +{ + /// Get the next item from the stream, blocking on the shared runtime if + /// necessary. + /// + /// If the stream is not ready to yield an item, this will release the GIL + /// and block until the next item is available. + /// + /// This method will return `None` if the stream is exhausted. + fn block_on_next(&mut self, py: Python<'_>, handle: &Handle) -> Option { + match self.get_next_if_ready() { + // `Some(Some(item))` (ready) and `Some(None)` (exhausted) are both + // answers we can return immediately — we just hand the inner + // `Option` straight back. + Some(row) => row, + // `None` means "not ready yet": release the GIL and block until the + // next item (or end of stream) arrives. + None => self.next().block_on(py, handle), + } + } + + /// Get the next item from the stream if it's ready, without blocking. + /// + /// Returns `None` if the stream is not ready to yield an item. Returns + /// `Some(None)` if the stream is exhausted. + fn get_next_if_ready(&mut self) -> Option> { + self.next().now_or_never() + } +} + +// Blanket impl over any pinned, fused stream. Requiring the helper bounds here +// (rather than only for `RowStream`) is what makes the cursor logic testable +// with a fake stream. +impl BlockingPostgresStream for Pin<&mut Fuse> +where + Self: futures::Stream + Send + Ungil + Unpin, + ::Item: Ungil + Send, +{ +} + +#[cfg(test)] +mod tests { + //! These tests don't touch Postgres: the future/stream helpers are generic, + //! so we exercise them with plain async blocks and an in-memory stream. + + use std::pin::pin; + + use futures::stream::{self, StreamExt}; + use tokio::runtime::Runtime; + + use super::*; + + /// A throwaway runtime standing in for the shared one. Production code takes + /// its handle from `crate::tokio_runtime`; the helpers only need *a* handle + /// to block on, and (as in production) the blocking wait runs on this test + /// thread rather than on a worker. + fn test_runtime() -> Runtime { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .unwrap() + } + + #[test] + fn block_on_runs_future_and_returns_output() { + Python::initialize(); + let rt = test_runtime(); + Python::attach(|py| { + assert_eq!(async { 1 + 2 }.block_on(py, rt.handle()), 3); + }); + } + + #[test] + fn block_on_result_maps_ok_through() { + Python::initialize(); + let rt = test_runtime(); + Python::attach(|py| { + let ok = async { Ok::(5) }; + assert_eq!(ok.block_on_result(py, rt.handle()).unwrap(), 5); + // The error path (mapping a `tokio_postgres::Error` to a `PyErr`) + // can't be unit-tested here, as that error type can't be + // constructed by hand; it's exercised by the integration tests. + }); + } + + #[test] + fn get_next_if_ready_returns_buffered_rows_then_signals_end() { + Python::initialize(); + let rt = test_runtime(); + Python::attach(|py| { + // `stream::iter` yields each item immediately, so every poll is + // ready: we get the items, then a `Some(None)` end-of-stream once + // it's drained, without ever needing to block. + let stream = stream::iter(vec![Ok::(1), Ok(2)]).fuse(); + let mut stream = pin!(stream); + + assert_eq!(stream.as_mut().get_next_if_ready(), Some(Some(Ok(1)))); + assert_eq!(stream.as_mut().get_next_if_ready(), Some(Some(Ok(2)))); + // Exhausted: the item is "ready" and is `None`. + assert_eq!(stream.as_mut().get_next_if_ready(), Some(None)); + // A fused stream keeps reporting end-of-stream rather than panicking. + assert_eq!(stream.as_mut().get_next_if_ready(), Some(None)); + + // `block_on_next` takes the same already-ready value. + let stream = stream::iter(vec![Ok::(9)]).fuse(); + let mut stream = pin!(stream); + assert_eq!(stream.as_mut().block_on_next(py, rt.handle()), Some(Ok(9))); + assert_eq!(stream.as_mut().block_on_next(py, rt.handle()), None); + }); + } + + #[test] + fn block_on_next_blocks_when_first_poll_is_pending() { + Python::initialize(); + let rt = test_runtime(); + Python::attach(|py| { + // A stream whose first poll is `Pending` (it yields back to the + // runtime before producing the value). `get_next_if_ready` / + // `now_or_never` polls exactly once and so sees `Pending` and gives + // up, forcing `block_on_next` down its blocking path. + let stream = stream::once(async { + tokio::task::yield_now().await; + Ok::(7) + }) + .fuse(); + let mut stream = pin!(stream); + + assert_eq!(stream.as_mut().get_next_if_ready(), None); + // `get_next_if_ready` above polled (and so advanced) the *same* + // pinned stream; `block_on_next` re-polls that same stream via + // `&mut self`, resuming the yielded future rather than restarting + // it, so it still resolves to 7. + assert_eq!(stream.as_mut().block_on_next(py, rt.handle()), Some(Ok(7))); + assert_eq!(stream.as_mut().block_on_next(py, rt.handle()), None); + + // And, on a fresh stream, `block_on_next` handles the pending first + // poll entirely on its own (no preceding `get_next_if_ready`), + // proving it doesn't rely on being "primed" by an earlier call. + let stream = stream::once(async { + tokio::task::yield_now().await; + Ok::(8) + }) + .fuse(); + let mut stream = pin!(stream); + assert_eq!(stream.as_mut().block_on_next(py, rt.handle()), Some(Ok(8))); + }); + } +} diff --git a/rust/src/database/postgres/mod.rs b/rust/src/database/postgres/mod.rs index 5ad02434dd..3710f8b972 100644 --- a/rust/src/database/postgres/mod.rs +++ b/rust/src/database/postgres/mod.rs @@ -7,14 +7,16 @@ //! The driver itself is async; the eventual `Connection` / `Cursor` types will //! drive it from sync Python methods via a shared multi-thread tokio runtime. +use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use pyo3::types::PyModule; -// `pub` (rather than private) so the value-mapping types are reachable from the -// crate root as public API while nothing inside the crate consumes them yet. -// This is what stops clippy's `dead_code` lint from firing on them before the +// `pub` (rather than private) so the not-yet-consumed public items in these +// submodules are reachable from the crate root as public API. This is what +// stops clippy's `dead_code` lint from firing on them before the // cursor/connection code (added in later changes) wires them up; the visibility // is tightened back to private once that happens. +pub mod helpers; pub mod value; /// Register the `postgres` submodule under the parent `database` module. @@ -35,3 +37,8 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> Ok(()) } + +/// Map a [`tokio_postgres`] error into a Python `RuntimeError`. +fn pg_err_to_py(e: tokio_postgres::Error) -> PyErr { + PyRuntimeError::new_err(format!("postgres error: {e}")) +} diff --git a/rust/src/tokio_runtime.rs b/rust/src/tokio_runtime.rs index f239040f30..d1338ed831 100644 --- a/rust/src/tokio_runtime.rs +++ b/rust/src/tokio_runtime.rs @@ -15,7 +15,7 @@ use anyhow::Context; use pyo3::prelude::*; -use tokio::runtime::Runtime; +use tokio::runtime::{Handle, Runtime}; /// This is the name of the attribute where we store the runtime on the reactor static TOKIO_RUNTIME_ATTR: &str = "__synapse_rust_tokio_runtime"; @@ -32,15 +32,7 @@ pub struct PyTokioRuntime { #[pymethods] impl PyTokioRuntime { fn start(&mut self) -> PyResult<()> { - // TODO: allow customization of the runtime like the number of threads - let runtime = tokio::runtime::Builder::new_multi_thread() - .worker_threads(4) - .enable_all() - .build()?; - - self.runtime = Some(runtime); - - Ok(()) + self.ensure_started() } fn shutdown(&mut self) -> PyResult<()> { @@ -57,6 +49,28 @@ impl PyTokioRuntime { } impl PyTokioRuntime { + /// Build the runtime if it hasn't been built yet. + /// + /// Idempotent, so it is safe to call both from the reactor's + /// `callWhenRunning(start)` hook and from a caller that needs the runtime + /// before the reactor has run that hook (see [`runtime_handle`]): whichever + /// runs first builds it, the other is a no-op. + fn ensure_started(&mut self) -> PyResult<()> { + if self.runtime.is_some() { + return Ok(()); + } + + // TODO: allow customization of the runtime like the number of threads + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build()?; + + self.runtime = Some(runtime); + + Ok(()) + } + /// Get the handle to the Tokio runtime, if it is running. pub fn handle(&self) -> PyResult<&tokio::runtime::Handle> { let handle = self @@ -72,11 +86,22 @@ impl PyTokioRuntime { /// Get a handle to the Tokio runtime stored on the reactor instance, or create /// a new one. pub fn runtime<'a>(reactor: &Bound<'a, PyAny>) -> PyResult> { - if !reactor.hasattr(TOKIO_RUNTIME_ATTR)? { - install_runtime(reactor)?; - } + Ok(get_or_install(reactor)?.borrow()) +} - get_runtime(reactor) +/// Get a clonable handle to the shared runtime, starting it on demand. +/// +/// Unlike [`runtime`], this does not require the reactor to have already run +/// its `callWhenRunning(start)` hook: it starts the runtime if necessary. That +/// lets callers that need the runtime before the reactor is up — the database +/// backend's schema setup, `synapse_port_db`, and trial tests — still get a +/// working handle. Once the reactor does run, its `start` hook finds the +/// runtime already built and is a no-op, so there is still only one runtime. +pub fn runtime_handle(reactor: &Bound<'_, PyAny>) -> PyResult { + let runtime = get_or_install(reactor)?; + let mut runtime = runtime.borrow_mut(); + runtime.ensure_started()?; + Ok(runtime.handle()?.clone()) } /// Install a new Tokio runtime on the reactor instance. @@ -97,11 +122,16 @@ fn install_runtime(reactor: &Bound) -> PyResult<()> { Ok(()) } -/// Get a reference to a Tokio runtime handle stored on the reactor instance. -fn get_runtime<'a>(reactor: &Bound<'a, PyAny>) -> PyResult> { +/// Get the [`PyTokioRuntime`] stored on the reactor instance, installing a +/// fresh one (wired to the reactor's start/shutdown) if it isn't there yet. +fn get_or_install<'a>(reactor: &Bound<'a, PyAny>) -> PyResult> { + if !reactor.hasattr(TOKIO_RUNTIME_ATTR)? { + install_runtime(reactor)?; + } + // This will raise if `TOKIO_RUNTIME_ATTR` is not set or if it is // not a `Runtime`. Careful that this could happen if the user sets it // manually, or if multiple versions of `pyo3-twisted` are used! let runtime: Bound = reactor.getattr(TOKIO_RUNTIME_ATTR)?.extract()?; - Ok(runtime.borrow()) + Ok(runtime) }