diff --git a/rust/src/federation/queue/native.rs b/rust/src/federation/queue/native.rs index f9c3064373..07bc015bb5 100644 --- a/rust/src/federation/queue/native.rs +++ b/rust/src/federation/queue/native.rs @@ -22,10 +22,10 @@ //! hear about delivered events. Catch-up and the sticky event backlog run //! here too, reading events through the shared event cache. -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use async_trait::async_trait; -use pyo3::{exceptions::PyRuntimeError, intern, prelude::*}; +use pyo3::prelude::*; use super::{ changes::StreamChanges, metrics::QueueMetrics, python::PyHooks, Completion, DestinationQueue, @@ -33,12 +33,13 @@ use super::{ StickyBatch, MAX_PDUS_PER_TRANSACTION, QUEUE_LOG, }; use crate::{ - database::postgres::pool::{ConnectionPool, PyConnectionPool}, + database::postgres::pool::ConnectionPool, events::Event, gil::with_gil, storage::{ events::RedactBehaviour, federation_store::FederationStore, + native_pool::NativePool, sync_store::{SyncReader, SyncStore}, }, visibility::{filter_events_for_server, ServerFilter, ServerFiltered}, @@ -66,6 +67,13 @@ pub struct Native { } impl Native { + /// The connection pool. + fn pool(&self) -> Result { + self.pool + .get() + .map_err(|err| LoopError::Other(err.to_string())) + } + /// The positions of the to-device and device list streams. Takes the /// GIL, but not a trip to the reactor: the ID generators are read where /// they stand. @@ -77,41 +85,6 @@ impl Native { } } -/// The native pool, as the Python `RustConnectionPool` holds it. That -/// replaces its pool when it is closed and started again, so a closed pool -/// is read again from Python. -pub struct NativePool { - /// The Python `RustConnectionPool`. - python: Py, - cached: Mutex>, -} - -impl NativePool { - pub fn new(python: Py) -> Self { - NativePool { - python, - cached: Mutex::new(None), - } - } - - pub fn get(&self) -> Result { - let mut cached = self.cached.lock().expect("poisoned"); - if let Some(pool) = cached.as_ref().filter(|pool| !pool.is_closed()) { - return Ok(pool.clone()); - } - let pool = with_gil("queue_db_pool", |py| -> PyResult { - let pool = self.python.bind(py).getattr(intern!(py, "_pool"))?; - if pool.is_none() { - return Err(PyRuntimeError::new_err("the database pool is closed")); - } - Ok(pool.cast::()?.get().pool.clone()) - }) - .map_err(|err| LoopError::Other(err.to_string()))?; - *cached = Some(pool.clone()); - Ok(pool) - } -} - /// The hooks with the database work done natively. pub(super) struct NativeHooks { pub(super) native: Arc, @@ -232,7 +205,7 @@ impl QueueHooks for NativeHooks { observer: &Observer, ) -> Result<(), LoopError> { let destination = self.destination.as_str(); - let pool = self.native.pool.get()?; + let pool = self.native.pool()?; let store = &self.native.store; let reader = self.native.sync.reader(pool.clone()); @@ -378,7 +351,7 @@ impl QueueHooks for NativeHooks { } let (to_device_token, device_list_token) = self.native.stream_tokens()?; - let pool = self.native.pool.get()?; + let pool = self.native.pool()?; let store = &self.native.store; let (to_device, to_device_stream_id) = if to_device_changed { @@ -442,7 +415,7 @@ impl QueueHooks for NativeHooks { observer: &Observer, ) -> Result, LoopError> { let destination = self.destination.as_str(); - let pool = self.native.pool.get()?; + let pool = self.native.pool()?; let Some((room_id, position, event_ids)) = self .native .store @@ -538,7 +511,7 @@ impl QueueHooks for NativeHooks { || completion.device_list_stream_id.is_some() || completion.last_stream_ordering.is_some(); if writes { - let pool = self.native.pool.get()?; + let pool = self.native.pool()?; let store = &self.native.store; if let Some(stream_id) = completion.to_device_stream_id { store @@ -570,7 +543,7 @@ impl QueueHooks for NativeHooks { return Ok(()); }; if let Some(NativeStickyPosition(position)) = position.downcast_ref() { - let pool = self.native.pool.get()?; + let pool = self.native.pool()?; return self .native .store diff --git a/rust/src/federation/queue/python.rs b/rust/src/federation/queue/python.rs index 01eec2327f..fb8061b1df 100644 --- a/rust/src/federation/queue/python.rs +++ b/rust/src/federation/queue/python.rs @@ -38,7 +38,7 @@ use serde_json::Value; use super::{ changes::StreamChanges, metrics::QueueMetrics, - native::{Native, NativeHooks, NativePool}, + native::{Native, NativeHooks}, Activity, Completion, DestinationQueue, DeviceEdus, DeviceEdusRequest, Edu, LoopError, Observer, PduBody, PresenceState, QueueConfig, QueueHooks, QueuedPdu, Receipt, SendOutcome, StickyBatch, StickyStreamPosition, TransactionSink, @@ -59,6 +59,7 @@ use crate::{ runtime::RustRuntime, storage::{ federation_store::{FederationStore, FederationStoreConfig}, + native_pool::NativePool, sync_store::{SyncStore, SyncStoreConfig}, }, }; @@ -132,7 +133,7 @@ impl RustDestinationQueues { allow_device_name_lookup_over_federation, server_name: server_name.clone(), }), - pool: NativePool::new(db_pool), + pool: NativePool::new(db_pool, "queue_db_pool"), stream_tokens, to_device: StreamChanges::new(to_device_since), device_lists: StreamChanges::new(device_lists_since), diff --git a/rust/src/storage/mod.rs b/rust/src/storage/mod.rs index 8c61cb7331..b14ab715ef 100644 --- a/rust/src/storage/mod.rs +++ b/rust/src/storage/mod.rs @@ -17,6 +17,7 @@ pub mod caches; pub mod db; pub mod events; pub mod federation_store; +pub mod native_pool; pub mod pool_querier; pub mod python; pub mod store; diff --git a/rust/src/storage/native_pool.rs b/rust/src/storage/native_pool.rs new file mode 100644 index 0000000000..c5d1087da0 --- /dev/null +++ b/rust/src/storage/native_pool.rs @@ -0,0 +1,69 @@ +/* + * This file is licensed under the Affero General Public License (AGPL) version 3. + * + * Copyright (C) 2026 Element Creations Ltd + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of the + * License, or (at your option) any later version. + * + * See the GNU Affero General Public License for more details: + * . + * + */ + +//! The native driver's pool, for Rust code that runs on tokio and must not +//! take the GIL for each query. + +use std::sync::Mutex; + +use pyo3::{exceptions::PyRuntimeError, intern, prelude::*}; + +use crate::{ + database::postgres::pool::{ConnectionPool, PyConnectionPool}, + gil::with_gil, +}; + +/// The native pool, as the Python `RustConnectionPool` holds it. That +/// replaces its pool when it is closed and started again, so a closed pool +/// is read again from Python, under the GIL site `site`. +pub struct NativePool { + /// The Python `RustConnectionPool`. + python: Py, + site: &'static str, + cached: Mutex>, +} + +impl NativePool { + pub fn new(python: Py, site: &'static str) -> Self { + NativePool { + python, + site, + cached: Mutex::new(None), + } + } + + pub fn get(&self) -> PyResult { + if let Some(pool) = self + .cached + .lock() + .unwrap_or_else(|e| e.into_inner()) + .as_ref() + .filter(|pool| !pool.is_closed()) + { + return Ok(pool.clone()); + } + // The lock is not held while this takes the GIL. A thread that holds + // the GIL and waits for the lock would otherwise deadlock with us. + let pool = with_gil(self.site, |py| -> PyResult { + let pool = self.python.bind(py).getattr(intern!(py, "_pool"))?; + if pool.is_none() { + return Err(PyRuntimeError::new_err("the database pool is closed")); + } + Ok(pool.cast::()?.get().pool.clone()) + })?; + *self.cached.lock().unwrap_or_else(|e| e.into_inner()) = Some(pool.clone()); + Ok(pool) + } +}