Add GIL-releasing block_on helpers on the shared tokio runtime

Add the `block_on`/`block_on_result`/`block_on_next` helpers the Postgres
backend uses to drive its async `tokio-postgres` futures to completion from
sync, GIL-holding Python methods, releasing the GIL for the wait. They take
a `tokio::runtime::Handle` and block on it from the calling (Python) thread.

Rather than give the DB backend a runtime of its own, they use the
extension's existing shared runtime (`tokio_runtime::PyTokioRuntime`, stored
on the reactor). `start` is made idempotent and a `runtime_handle` accessor
starts it on demand, so a caller that needs a connection before the reactor
is running still gets a handle; once the reactor runs, its
`callWhenRunning(start)` hook is a no-op.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RPEeXx2fAG67o6u4CnmC8W

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<T>
+where
+    Self: Future<Output = Result<T, tokio_postgres::Error>> + 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<T> {
+        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<F> BlockingPostgres for F
+where
+    F: Future + Sized + Send + Ungil,
+    F::Output: Ungil + Send,
+{
+}
+impl<F, T> BlockingPostgresResult<T> for F
+where
+    F: Future<Output = Result<T, tokio_postgres::Error>> + 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<S>>`) 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<S>>`), 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<Self::Item> {
+        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<Item>` 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<Option<Self::Item>> {
+        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<S> BlockingPostgresStream for Pin<&mut Fuse<S>>
+where
+    Self: futures::Stream + Send + Ungil + Unpin,
+    <Self as futures::Stream>::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::<i32, tokio_postgres::Error>(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::<i32, ()>(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::<i32, ()>(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::<i32, ()>(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::<i32, ()>(8)
+            })
+            .fuse();
+            let mut stream = pin!(stream);
+            assert_eq!(stream.as_mut().block_on_next(py, rt.handle()), Some(Ok(8)));
+        });
+    }
+}

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

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<PyRef<'a, PyTokioRuntime>> {
-    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<Handle> {
+    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<PyAny>) -> PyResult<()> {
     Ok(())
 }
 
-/// Get a reference to a Tokio runtime handle stored on the reactor instance.
-fn get_runtime<'a>(reactor: &Bound<'a, PyAny>) -> PyResult<PyRef<'a, PyTokioRuntime>> {
+/// 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<Bound<'a, PyTokioRuntime>> {
+    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<PyTokioRuntime> = reactor.getattr(TOKIO_RUNTIME_ATTR)?.extract()?;
-    Ok(runtime.borrow())
+    Ok(runtime)
 }