inbuxa/inbuxa-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases 5 Wiki Activity

SQL queries time out; readiness follows the data store #45

Merged
jcoffey-dev merged 1 commits from fix/query-timeouts into main 2026-09-25 00:45:10 +00:00
Conversation 0 Commits 1 Files Changed 22
+1525 -787
22 changed files with 1525 additions and 787 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/common/src/config/inner.rs
+2
View File
@@ -94,6 +94,7 @@ impl Data {
span_id_gen: id_generator, span_id_gen: id_generator,
queue_status: true.into(), queue_status: true.into(),
settings_reload: Default::default(), settings_reload: Default::default(),
store_health: Default::default(),
applications, applications,
logos: Default::default(), logos: Default::default(),
smtp_connectors: TlsConnectors::try_new().failed("Failed to build TLS connectors"), smtp_connectors: TlsConnectors::try_new().failed("Failed to build TLS connectors"),
@@ -237,6 +238,7 @@ impl Default for Data {
registry_id_gen: Default::default(), registry_id_gen: Default::default(),
queue_status: true.into(), queue_status: true.into(),
settings_reload: Default::default(), settings_reload: Default::default(),
store_health: Default::default(),
applications: WebApplications::new(), applications: WebApplications::new(),
logos: Default::default(), logos: Default::default(),
smtp_connectors: TlsConnectors::try_new().unwrap(), smtp_connectors: TlsConnectors::try_new().unwrap(),
crates/common/src/lib.rs
+2
View File
@@ -163,6 +163,8 @@ pub struct Data {
pub queue_status: AtomicBool, pub queue_status: AtomicBool,
// inbuxa: coalesces the settings reloads registry writes trigger // inbuxa: coalesces the settings reloads registry writes trigger
pub settings_reload: cache::reload::SettingsReloadGate, pub settings_reload: cache::reload::SettingsReloadGate,
// inbuxa: the readiness probe's cached answer
pub store_health: storage::ready::StoreHealth,
pub applications: WebApplications, pub applications: WebApplications,
pub logos: Mutex<AHashMap<Box<str>, LogoCache>>, pub logos: Mutex<AHashMap<Box<str>, LogoCache>>,
crates/common/src/storage/mod.rs
+1
View File
@@ -26,6 +26,7 @@ pub mod document;
pub mod encryption; pub mod encryption;
pub mod index; pub mod index;
pub mod quota; pub mod quota;
pub mod ready; // inbuxa: readiness follows the data store
pub mod state; pub mod state;
pub mod transaction; pub mod transaction;
crates/common/src/storage/ready.rs
+83
View File
@@ -0,0 +1,83 @@
/*
* SPDX-FileCopyrightText: 2026 Coffey Labs
*
* SPDX-License-Identifier: AGPL-3.0-only
*/
//! Readiness that reflects the data store.
//!
//! /healthz/ready used to answer 200 whenever a data store was configured,
//! so a load balancer kept sending traffic to a node through a database
//! outage. It now reads one key from the data store, with a short time
//! limit, and caches the answer for a couple of seconds so probes can't load
//! the database. Liveness stays 200: restarting a node doesn't bring its
//! database back, and an orchestrator that restarts on failed liveness would
//! otherwise restart every node at once.
use crate::Server;
use parking_lot::Mutex;
use std::{
sync::atomic::{AtomicBool, Ordering},
time::{Duration, Instant},
};
use store::{ValueKey, write::ValueClass};
/// How long a probe's answer is reused.
pub const READY_CACHE: Duration = Duration::from_secs(2);
/// How long a probe waits for the data store.
pub const READY_PROBE_TIMEOUT: Duration = Duration::from_secs(2);
#[derive(Default)]
pub struct StoreHealth {
last: Mutex<Option<(Instant, bool)>>,
probing: AtomicBool,
}
/// Clears the probing flag even when the request is dropped mid-probe.
struct ProbeGuard<'x>(&'x AtomicBool);
impl Drop for ProbeGuard<'_> {
fn drop(&mut self) {
self.0.store(false, Ordering::Release);
}
}
impl Server {
/// Whether the data store answers: a cached result younger than
/// READY_CACHE, or a fresh read bounded by READY_PROBE_TIMEOUT. While
/// one probe is running, other callers get the last answer.
pub async fn is_data_store_ready(&self) -> bool {
let store = &self.core.storage.data;
if store.is_none() {
return false;
}
let health = &self.inner.data.store_health;
let last = *health.last.lock();
if let Some((at, ready)) = last
&& at.elapsed() < READY_CACHE
{
return ready;
}
if health.probing.swap(true, Ordering::AcqRel) {
return last.is_none_or(|(_, ready)| ready);
}
let _guard = ProbeGuard(&health.probing);
let ready = tokio::time::timeout(
READY_PROBE_TIMEOUT,
store.get_value::<u64>(ValueKey::from(ValueClass::Property(0))),
)
.await
.is_ok_and(|result| result.is_ok());
// Say so once per outage, not on every probe
if !ready && last.is_none_or(|(_, ready)| ready) {
trc::event!(
Store(trc::StoreEvent::UnexpectedError),
Details = "Readiness probe: the data store didn't answer",
Limit = READY_PROBE_TIMEOUT,
);
}
*health.last.lock() = Some((Instant::now(), ready));
ready
}
}
crates/http/src/request.rs
+3 -1
View File
@@ -553,8 +553,10 @@ impl ParseHttp for Server {
return Ok(JsonProblemResponse(StatusCode::OK).into_http_response()); return Ok(JsonProblemResponse(StatusCode::OK).into_http_response());
} }
"ready" => { "ready" => {
// inbuxa: ready only while the data store answers
// (a cached, time-limited read); liveness stays 200
return Ok(JsonProblemResponse({ return Ok(JsonProblemResponse({
if !self.core.storage.data.is_none() { if self.is_data_store_ready().await {
StatusCode::OK StatusCode::OK
} else { } else {
StatusCode::SERVICE_UNAVAILABLE StatusCode::SERVICE_UNAVAILABLE
crates/store/src/backend/mod.rs
+3
View File
@@ -30,6 +30,9 @@ pub mod s3;
pub mod sqlite; pub mod sqlite;
// inbuxa: scale-out storage (sharded stores) // inbuxa: scale-out storage (sharded stores)
pub mod scaleout; pub mod scaleout;
// inbuxa: client-side SQL query limits
#[cfg(any(feature = "postgres", feature = "mysql"))]
pub mod query_timeout;
pub const MAX_TOKEN_LENGTH: usize = (u8::MAX >> 1) as usize; pub const MAX_TOKEN_LENGTH: usize = (u8::MAX >> 1) as usize;
crates/store/src/backend/mysql/blob.rs
+50 -35
View File
@@ -10,7 +10,7 @@ use std::ops::Range;
use mysql_async::prelude::Queryable; use mysql_async::prelude::Queryable;
use super::{MysqlStore, into_error}; use super::{MysqlStore, bounded, into_error};
impl MysqlStore { impl MysqlStore {
pub(crate) async fn get_blob( pub(crate) async fn get_blob(
@@ -19,48 +19,63 @@ impl MysqlStore {
range: Range<usize>, range: Range<usize>,
) -> trc::Result<Option<Vec<u8>>> { ) -> trc::Result<Option<Vec<u8>>> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let s = conn let limit = self.timeouts.query;
.prep("SELECT v FROM t WHERE k = ?") let result = tokio::time::timeout(limit, async {
.await let s = conn
.map_err(into_error)?; .prep("SELECT v FROM t WHERE k = ?")
conn.exec_first::<Vec<u8>, _, _>(&s, (key,)) .await
.await .map_err(into_error)?;
.map(|bytes| { conn.exec_first::<Vec<u8>, _, _>(&s, (key,))
if range.start == 0 && range.end == usize::MAX { .await
bytes .map(|bytes| {
} else { if range.start == 0 && range.end == usize::MAX {
bytes.map(|bytes| {
bytes bytes
.get(range.start..std::cmp::min(bytes.len(), range.end)) } else {
.unwrap_or_default() bytes.map(|bytes| {
.to_vec() bytes
}) .get(range.start..std::cmp::min(bytes.len(), range.end))
} .unwrap_or_default()
}) .to_vec()
.map_err(into_error) })
}
})
.map_err(into_error)
})
.await;
bounded(conn, result, limit)
} }
pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> trc::Result<()> { pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> trc::Result<()> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let s = conn let limit = self.timeouts.query;
.prep("INSERT INTO t (k, v) VALUES (?, ?) ON DUPLICATE KEY UPDATE v = VALUES(v)") let result = tokio::time::timeout(limit, async {
.await let s = conn
.map_err(into_error)?; .prep("INSERT INTO t (k, v) VALUES (?, ?) ON DUPLICATE KEY UPDATE v = VALUES(v)")
conn.exec_drop(&s, (key, data)) .await
.await .map_err(into_error)?;
.map_err(into_error) conn.exec_drop(&s, (key, data))
.map(|_| ()) .await
.map_err(into_error)
.map(|_| ())
})
.await;
bounded(conn, result, limit)
} }
pub(crate) async fn delete_blob(&self, key: &[u8]) -> trc::Result<bool> { pub(crate) async fn delete_blob(&self, key: &[u8]) -> trc::Result<bool> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let s = conn let limit = self.timeouts.query;
.prep("DELETE FROM t WHERE k = ?") let result = tokio::time::timeout(limit, async {
.await let s = conn
.map_err(into_error)?; .prep("DELETE FROM t WHERE k = ?")
conn.exec_iter(&s, (key,)) .await
.await .map_err(into_error)?;
.map_err(into_error) conn.exec_iter(&s, (key,))
.map(|hits| hits.affected_rows() > 0) .await
.map_err(into_error)
.map(|hits| hits.affected_rows() > 0)
})
.await;
bounded(conn, result, limit)
} }
} }
crates/store/src/backend/mysql/lookup.rs
+26 -21
View File
@@ -10,7 +10,7 @@ use mysql_async::{Params, Row, prelude::Queryable};
use crate::{IntoRows, QueryResult, QueryType, Value}; use crate::{IntoRows, QueryResult, QueryType, Value};
use super::{MysqlStore, into_error}; use super::{MysqlStore, bounded, into_error};
impl MysqlStore { impl MysqlStore {
pub(crate) async fn sql_query<T: QueryResult>( pub(crate) async fn sql_query<T: QueryResult>(
@@ -19,27 +19,32 @@ impl MysqlStore {
params: &[Value<'_>], params: &[Value<'_>],
) -> trc::Result<T> { ) -> trc::Result<T> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let s = conn.prep(query).await.map_err(into_error)?; let limit = self.timeouts.query;
let params = Params::Positional(params.iter().map(Into::into).collect()); let result = tokio::time::timeout(limit, async {
let s = conn.prep(query).await.map_err(into_error)?;
let params = Params::Positional(params.iter().map(Into::into).collect());
match T::query_type() { match T::query_type() {
QueryType::Execute => conn.exec_drop(s, params).await.map_or_else( QueryType::Execute => conn.exec_drop(s, params).await.map_or_else(
|e| Err(into_error(e)), |e| Err(into_error(e)),
|_| Ok(T::from_exec(conn.affected_rows() as usize)), |_| Ok(T::from_exec(conn.affected_rows() as usize)),
), ),
QueryType::Exists => conn QueryType::Exists => conn
.exec_first::<Row, _, _>(s, params) .exec_first::<Row, _, _>(s, params)
.await .await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_exists(r.is_some()))), .map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_exists(r.is_some()))),
QueryType::QueryOne => conn QueryType::QueryOne => conn
.exec_first::<Row, _, _>(s, params) .exec_first::<Row, _, _>(s, params)
.await .await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_query_one(r))), .map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_query_one(r))),
QueryType::QueryAll => conn QueryType::QueryAll => conn
.exec::<Row, _, _>(s, params) .exec::<Row, _, _>(s, params)
.await .await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_query_all(r))), .map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_query_all(r))),
} }
})
.await;
bounded(conn, result, limit)
} }
} }
crates/store/src/backend/mysql/main.rs
+81 -71
View File
@@ -6,7 +6,7 @@
* Modified by Coffey Labs in 2026 for INBUXA. * Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use super::{MysqlStore, into_error}; use super::{MysqlStore, bounded, into_error};
use crate::{ use crate::{
backend::mysql::MysqlSearchField, backend::mysql::MysqlSearchField,
search::{ search::{
@@ -72,6 +72,7 @@ impl MysqlStore {
.db_name(Some(replica.database.clone())) .db_name(Some(replica.database.clone()))
.tcp_port(replica.port as u16), .tcp_port(replica.port as u16),
), ),
timeouts: Default::default(),
})), })),
replica.host, replica.host,
replica.port as u16, replica.port as u16,
@@ -81,6 +82,7 @@ impl MysqlStore {
let primary = Store::MySQL(Arc::new(MysqlStore { let primary = Store::MySQL(Arc::new(MysqlStore {
conn_pool: Pool::new(opts), conn_pool: Pool::new(opts),
timeouts: Default::default(),
})); }));
// ST-1: no replicas, no change // ST-1: no replicas, no change
@@ -99,88 +101,96 @@ impl MysqlStore {
pub(crate) async fn create_storage_tables(&self) -> trc::Result<()> { pub(crate) async fn create_storage_tables(&self) -> trc::Result<()> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let limit = self.timeouts.maintenance;
let result = tokio::time::timeout(limit, async {
for table in [
SUBSPACE_ACL,
SUBSPACE_TASK_QUEUE,
SUBSPACE_DELETED_ITEMS,
SUBSPACE_SPAM_SAMPLES,
crate::SUBSPACE_INBUXA, // inbuxa: masked email
SUBSPACE_BLOB_LINK,
SUBSPACE_IN_MEMORY_VALUE,
SUBSPACE_PROPERTY,
SUBSPACE_REGISTRY,
SUBSPACE_REGISTRY_PK,
SUBSPACE_DIRECTORY,
SUBSPACE_QUEUE_MESSAGE,
SUBSPACE_QUEUE_EVENT,
SUBSPACE_REPORT_OUT,
SUBSPACE_REPORT_IN,
SUBSPACE_LOGS,
SUBSPACE_TELEMETRY_SPAN,
SUBSPACE_TELEMETRY_METRIC,
] {
let table = char::from(table);
conn.query_drop(format!(
"CREATE TABLE IF NOT EXISTS {table} (
k VARBINARY(255) NOT NULL,
v MEDIUMBLOB NOT NULL,
PRIMARY KEY (k)
) ENGINE=InnoDB"
))
.await
.map_err(into_error)?;
}
for table in [
SUBSPACE_ACL,
SUBSPACE_TASK_QUEUE,
SUBSPACE_DELETED_ITEMS,
SUBSPACE_SPAM_SAMPLES,
crate::SUBSPACE_INBUXA, // inbuxa: masked email
SUBSPACE_BLOB_LINK,
SUBSPACE_IN_MEMORY_VALUE,
SUBSPACE_PROPERTY,
SUBSPACE_REGISTRY,
SUBSPACE_REGISTRY_PK,
SUBSPACE_DIRECTORY,
SUBSPACE_QUEUE_MESSAGE,
SUBSPACE_QUEUE_EVENT,
SUBSPACE_REPORT_OUT,
SUBSPACE_REPORT_IN,
SUBSPACE_LOGS,
SUBSPACE_TELEMETRY_SPAN,
SUBSPACE_TELEMETRY_METRIC,
] {
let table = char::from(table);
conn.query_drop(format!(
"CREATE TABLE IF NOT EXISTS {table} (
k VARBINARY(255) NOT NULL,
v MEDIUMBLOB NOT NULL,
PRIMARY KEY (k)
) ENGINE=InnoDB"
))
.await
.map_err(into_error)?;
}
conn.query_drop(format!(
"CREATE TABLE IF NOT EXISTS {} (
k VARBINARY(255) NOT NULL,
v LONGBLOB NOT NULL,
PRIMARY KEY (k)
) ENGINE=InnoDB",
char::from(SUBSPACE_BLOBS),
))
.await
.map_err(into_error)?;
for table in [SUBSPACE_INDEXES, SUBSPACE_REGISTRY_IDX] {
let table = char::from(table);
conn.query_drop(format!(
"CREATE TABLE IF NOT EXISTS {table} (
k BLOB,
PRIMARY KEY (k(400))
) ENGINE=InnoDB"
))
.await
.map_err(into_error)?;
}
for table in [SUBSPACE_COUNTER, SUBSPACE_QUOTA, SUBSPACE_IN_MEMORY_COUNTER] {
conn.query_drop(format!( conn.query_drop(format!(
"CREATE TABLE IF NOT EXISTS {} ( "CREATE TABLE IF NOT EXISTS {} (
k VARBINARY(255) NOT NULL, k VARBINARY(255) NOT NULL,
v BIGINT NOT NULL DEFAULT 0, v LONGBLOB NOT NULL,
PRIMARY KEY (k) PRIMARY KEY (k)
) ENGINE=InnoDB", ) ENGINE=InnoDB",
char::from(table) char::from(SUBSPACE_BLOBS),
)) ))
.await .await
.map_err(into_error)?; .map_err(into_error)?;
}
Ok(()) for table in [SUBSPACE_INDEXES, SUBSPACE_REGISTRY_IDX] {
let table = char::from(table);
conn.query_drop(format!(
"CREATE TABLE IF NOT EXISTS {table} (
k BLOB,
PRIMARY KEY (k(400))
) ENGINE=InnoDB"
))
.await
.map_err(into_error)?;
}
for table in [SUBSPACE_COUNTER, SUBSPACE_QUOTA, SUBSPACE_IN_MEMORY_COUNTER] {
conn.query_drop(format!(
"CREATE TABLE IF NOT EXISTS {} (
k VARBINARY(255) NOT NULL,
v BIGINT NOT NULL DEFAULT 0,
PRIMARY KEY (k)
) ENGINE=InnoDB",
char::from(table)
))
.await
.map_err(into_error)?;
}
Ok(())
})
.await;
bounded(conn, result, limit)
} }
pub(crate) async fn create_search_tables(&self) -> trc::Result<()> { pub(crate) async fn create_search_tables(&self) -> trc::Result<()> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let limit = self.timeouts.maintenance;
let result = tokio::time::timeout(limit, async {
create_search_tables::<EmailSearchField>(&mut conn).await?;
create_search_tables::<CalendarSearchField>(&mut conn).await?;
create_search_tables::<ContactSearchField>(&mut conn).await?;
//create_search_tables::<FileSearchField>(&mut conn).await?;
create_search_tables::<TracingSearchField>(&mut conn).await?;
create_search_tables::<EmailSearchField>(&mut conn).await?; Ok(())
create_search_tables::<CalendarSearchField>(&mut conn).await?; })
create_search_tables::<ContactSearchField>(&mut conn).await?; .await;
//create_search_tables::<FileSearchField>(&mut conn).await?; bounded(conn, result, limit)
create_search_tables::<TracingSearchField>(&mut conn).await?;
Ok(())
} }
} }
crates/store/src/backend/mysql/mod.rs
+41 -1
View File
@@ -6,6 +6,7 @@
* Modified by Coffey Labs in 2026 for INBUXA. * Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use crate::backend::query_timeout::QueryTimeouts;
use crate::{ use crate::{
search::{ search::{
CalendarSearchField, ContactSearchField, EmailSearchField, FileSearchField, SearchField, CalendarSearchField, ContactSearchField, EmailSearchField, FileSearchField, SearchField,
@@ -14,7 +15,7 @@ use crate::{
write::SearchIndex, write::SearchIndex,
}; };
use mysql_async::Pool; use mysql_async::Pool;
use std::fmt::Display; use std::{fmt::Display, time::Duration};
pub mod blob; pub mod blob;
pub mod lookup; pub mod lookup;
@@ -25,6 +26,8 @@ pub mod write;
pub struct MysqlStore { pub struct MysqlStore {
pub(crate) conn_pool: Pool, pub(crate) conn_pool: Pool,
/// inbuxa: client-side query limits (see backend::query_timeout)
pub(crate) timeouts: QueryTimeouts,
} }
/// inbuxa: how long a request waits for a pooled connection (including /// inbuxa: how long a request waits for a pooled connection (including
@@ -54,6 +57,43 @@ pub(crate) async fn pool_conn(
} }
} }
/// inbuxa: the error for an operation that ran past its time limit.
pub(crate) fn query_timeout_error(limit: Duration) -> trc::Error {
trc::StoreEvent::MysqlError
.reason("Query timed out")
.details(format!(
"No answer from the database within {} s",
limit.as_secs()
))
}
/// inbuxa: ends an operation run on `conn` under `limit`. When it ran out,
/// the connection is closed rather than returned to the pool: a query may
/// still be in flight on it, or a transaction open. Conn::disconnect marks
/// the connection closed before it sends anything, so even when the server
/// doesn't answer and the attempt is dropped, the pool discards it instead
/// of waiting to clean it up.
pub(crate) fn bounded<T>(
conn: mysql_async::Conn,
result: Result<trc::Result<T>, tokio::time::error::Elapsed>,
limit: Duration,
) -> trc::Result<T> {
match result {
Ok(result) => result,
Err(_) => {
discard(conn);
Err(query_timeout_error(limit))
}
}
}
/// inbuxa: closes a connection whose state is unknown (see bounded).
pub(crate) fn discard(conn: mysql_async::Conn) {
tokio::spawn(async move {
let _ = tokio::time::timeout(Duration::from_secs(1), conn.disconnect()).await;
});
}
#[inline(always)] #[inline(always)]
pub(crate) fn into_error(err: impl Display) -> trc::Error { pub(crate) fn into_error(err: impl Display) -> trc::Error {
trc::StoreEvent::MysqlError.reason(err) trc::StoreEvent::MysqlError.reason(err)
crates/store/src/backend/mysql/read.rs
+109 -66
View File
@@ -6,7 +6,7 @@
* Modified by Coffey Labs in 2026 for INBUXA. * Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use super::{MysqlStore, into_error, is_timeout_error}; use super::{MysqlStore, bounded, discard, into_error, is_timeout_error, query_timeout_error};
use crate::{Deserialize, IterateParams, Key, ValueKey, write::ValueClass}; use crate::{Deserialize, IterateParams, Key, ValueKey, write::ValueClass};
use futures::TryStreamExt; use futures::TryStreamExt;
use mysql_async::{Row, prelude::Queryable}; use mysql_async::{Row, prelude::Queryable};
@@ -17,40 +17,50 @@ impl MysqlStore {
U: Deserialize + 'static, U: Deserialize + 'static,
{ {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let s = conn let limit = self.timeouts.query;
.prep(format!( let result = tokio::time::timeout(limit, async {
"SELECT v FROM {} WHERE k = ?", let s = conn
char::from(key.subspace()) .prep(format!(
)) "SELECT v FROM {} WHERE k = ?",
.await char::from(key.subspace())
.map_err(into_error)?; ))
let key = key.serialize(0); .await
conn.exec_first::<Vec<u8>, _, _>(&s, (&key,)) .map_err(into_error)?;
.await let key = key.serialize(0);
.map_err(into_error) conn.exec_first::<Vec<u8>, _, _>(&s, (&key,))
.and_then(|r| { .await
if let Some(r) = r { .map_err(into_error)
Ok(Some(U::deserialize_owned_with_key(&key, r)?)) .and_then(|r| {
} else { if let Some(r) = r {
Ok(None) Ok(Some(U::deserialize_owned_with_key(&key, r)?))
} } else {
}) Ok(None)
}
})
})
.await;
bounded(conn, result, limit)
} }
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> { pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let s = conn let limit = self.timeouts.query;
.prep(format!( let result = tokio::time::timeout(limit, async {
"SELECT 1 FROM {} WHERE k = ?", let s = conn
char::from(key.subspace()) .prep(format!(
)) "SELECT 1 FROM {} WHERE k = ?",
.await char::from(key.subspace())
.map_err(into_error)?; ))
let key = key.serialize(0); .await
conn.exec_first::<u8, _, _>(&s, (&key,)) .map_err(into_error)?;
.await let key = key.serialize(0);
.map_err(into_error) conn.exec_first::<u8, _, _>(&s, (&key,))
.map(|r| r.is_some()) .await
.map_err(into_error)
.map(|r| r.is_some())
})
.await;
bounded(conn, result, limit)
} }
pub(crate) async fn iterate<T: Key>( pub(crate) async fn iterate<T: Key>(
@@ -64,28 +74,36 @@ impl MysqlStore {
let end = params.end.serialize(0); let end = params.end.serialize(0);
let keys = if params.values { "k, v" } else { "k" }; let keys = if params.values { "k, v" } else { "k" };
let s = conn // inbuxa: a scan may run for hours, so the query limit bounds each
.prep(&match (params.first, params.ascending) { // wait for the database (preparing, the query starting, the next
(true, true) => { // row) rather than the scan. A wait that runs out closes the
format!( // connection.
"SELECT {keys} FROM {table} WHERE k >= ? AND k <= ? ORDER BY k ASC LIMIT 1" let limit = self.timeouts.query;
) let query = match (params.first, params.ascending) {
} (true, true) => {
(true, false) => { format!("SELECT {keys} FROM {table} WHERE k >= ? AND k <= ? ORDER BY k ASC LIMIT 1")
format!( }
"SELECT {keys} FROM {table} WHERE k >= ? AND k <= ? ORDER BY k DESC LIMIT 1" (true, false) => {
) format!(
} "SELECT {keys} FROM {table} WHERE k >= ? AND k <= ? ORDER BY k DESC LIMIT 1"
(false, true) => { )
format!("SELECT {keys} FROM {table} WHERE k >= ? AND k <= ? ORDER BY k ASC") }
} (false, true) => {
(false, false) => { format!("SELECT {keys} FROM {table} WHERE k >= ? AND k <= ? ORDER BY k ASC")
format!("SELECT {keys} FROM {table} WHERE k >= ? AND k <= ? ORDER BY k DESC") }
} (false, false) => {
}) format!("SELECT {keys} FROM {table} WHERE k >= ? AND k <= ? ORDER BY k DESC")
.await }
.map_err(into_error)?; };
let s = match tokio::time::timeout(limit, conn.prep(&query)).await {
Ok(s) => s.map_err(into_error)?,
Err(_) => {
discard(conn);
return Err(query_timeout_error(limit));
}
};
let mut from = begin; let mut from = begin;
let mut stalled = false;
let mut to = end; let mut to = end;
let mut resume_key = None; let mut resume_key = None;
@@ -94,13 +112,26 @@ impl MysqlStore {
let mut timed_out = false; let mut timed_out = false;
{ {
let mut rows = conn let mut rows = match tokio::time::timeout(
.exec_stream::<Row, _, _>(&s, (from.clone(), to.clone())) limit,
.await conn.exec_stream::<Row, _, _>(&s, (from.clone(), to.clone())),
.map_err(into_error)?; )
.await
{
Ok(rows) => rows.map_err(into_error)?,
// Leaves the scan loop for the timeout below
Err(_) => break,
};
loop { loop {
match rows.try_next().await { let next = match tokio::time::timeout(limit, rows.try_next()).await {
Ok(next) => next,
Err(_) => {
stalled = true;
break;
}
};
match next {
Ok(Some(mut row)) => { Ok(Some(mut row)) => {
let value = if params.values { let value = if params.values {
row.take_opt::<Vec<u8>, _>(1) row.take_opt::<Vec<u8>, _>(1)
@@ -136,6 +167,10 @@ impl MysqlStore {
} }
} }
if stalled {
break;
}
match last_key { match last_key {
Some(last_key) if timed_out => { Some(last_key) if timed_out => {
if params.ascending { if params.ascending {
@@ -148,6 +183,9 @@ impl MysqlStore {
_ => return Ok(()), _ => return Ok(()),
} }
} }
discard(conn);
Err(query_timeout_error(limit))
} }
pub(crate) async fn get_counter( pub(crate) async fn get_counter(
@@ -158,14 +196,19 @@ impl MysqlStore {
let table = char::from(key.subspace()); let table = char::from(key.subspace());
let key = key.serialize(0); let key = key.serialize(0);
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let s = conn let limit = self.timeouts.query;
.prep(format!("SELECT v FROM {table} WHERE k = ?")) let result = tokio::time::timeout(limit, async {
.await let s = conn
.map_err(into_error)?; .prep(format!("SELECT v FROM {table} WHERE k = ?"))
match conn.exec_first::<i64, _, _>(&s, (key,)).await { .await
Ok(Some(num)) => Ok(num), .map_err(into_error)?;
Ok(None) => Ok(0), match conn.exec_first::<i64, _, _>(&s, (key,)).await {
Err(e) => Err(into_error(e)), Ok(Some(num)) => Ok(num),
} Ok(None) => Ok(0),
Err(e) => Err(into_error(e)),
}
})
.await;
bounded(conn, result, limit)
} }
} }
crates/store/src/backend/mysql/search.rs
+96 -79
View File
@@ -10,8 +10,8 @@ use crate::{
backend::{ backend::{
MAX_TOKEN_LENGTH, MAX_TOKEN_LENGTH,
mysql::{ mysql::{
DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, MysqlSearchField, MysqlStore, into_error, DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, MysqlSearchField, MysqlStore, bounded,
is_timeout_error, into_error, is_timeout_error,
}, },
}, },
search::{ search::{
@@ -27,57 +27,62 @@ use std::fmt::Write;
impl MysqlStore { impl MysqlStore {
pub async fn index(&self, documents: Vec<IndexDocument>) -> trc::Result<()> { pub async fn index(&self, documents: Vec<IndexDocument>) -> trc::Result<()> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let mut tx_opts = TxOpts::default(); let limit = self.timeouts.query;
tx_opts let result = tokio::time::timeout(limit, async {
.with_consistent_snapshot(false) let mut tx_opts = TxOpts::default();
.with_isolation_level(IsolationLevel::ReadCommitted); tx_opts
let mut trx = conn.start_transaction(tx_opts).await.map_err(into_error)?; .with_consistent_snapshot(false)
.with_isolation_level(IsolationLevel::ReadCommitted);
let mut trx = conn.start_transaction(tx_opts).await.map_err(into_error)?;
for document in documents { for document in documents {
let index = document.index; let index = document.index;
let primary_keys = index.primary_keys(); let primary_keys = index.primary_keys();
let all_fields = index.all_fields(); let all_fields = index.all_fields();
let mut fields = document.fields; let mut fields = document.fields;
let mut values = Vec::with_capacity(fields.len() + 2); let mut values = Vec::with_capacity(fields.len() + 2);
let mut query = format!("INSERT INTO {} (", index.mysql_table()); let mut query = format!("INSERT INTO {} (", index.mysql_table());
for (i, field) in primary_keys.iter().chain(all_fields).enumerate() { for (i, field) in primary_keys.iter().chain(all_fields).enumerate() {
if i > 0 { if i > 0 {
query.push(','); query.push(',');
}
query.push_str(field.column());
} }
query.push_str(field.column());
query.push_str(") VALUES (");
for (i, field) in primary_keys.iter().chain(all_fields).enumerate() {
if i > 0 {
query.push(',');
}
if let Some(value) = fields.remove(field) {
query.push('?');
values.push(value);
} else {
query.push_str("NULL");
}
}
query.push_str(") ON DUPLICATE KEY UPDATE ");
for (i, field) in all_fields.iter().enumerate() {
if i > 0 {
query.push(',');
}
let column = field.column();
let _ = write!(&mut query, "{column} = VALUES({column})");
}
let s = trx.prep(&query).await.map_err(into_error)?;
trx.exec_drop(&s, values).await.map_err(into_error)?;
} }
query.push_str(") VALUES ("); trx.commit().await.map_err(into_error)
})
for (i, field) in primary_keys.iter().chain(all_fields).enumerate() { .await;
if i > 0 { bounded(conn, result, limit)
query.push(',');
}
if let Some(value) = fields.remove(field) {
query.push('?');
values.push(value);
} else {
query.push_str("NULL");
}
}
query.push_str(") ON DUPLICATE KEY UPDATE ");
for (i, field) in all_fields.iter().enumerate() {
if i > 0 {
query.push(',');
}
let column = field.column();
let _ = write!(&mut query, "{column} = VALUES({column})");
}
let s = trx.prep(&query).await.map_err(into_error)?;
trx.exec_drop(&s, values).await.map_err(into_error)?;
}
trx.commit().await.map_err(into_error)
} }
pub async fn query<R: SearchDocumentId>( pub async fn query<R: SearchDocumentId>(
@@ -97,12 +102,17 @@ impl MysqlStore {
} }
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let s = conn.prep(query).await.map_err(into_error)?; let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn.prep(query).await.map_err(into_error)?;
conn.exec::<i64, _, _>(s, params) conn.exec::<i64, _, _>(s, params)
.await .await
.map(|r| r.into_iter().map(|r| R::from_u64(r as u64)).collect()) .map(|r| r.into_iter().map(|r| R::from_u64(r as u64)).collect())
.map_err(into_error) .map_err(into_error)
})
.await;
bounded(conn, result, limit)
} }
pub async fn unindex(&self, filter: SearchQuery) -> trc::Result<u64> { pub async fn unindex(&self, filter: SearchQuery) -> trc::Result<u64> {
@@ -111,40 +121,47 @@ impl MysqlStore {
let params = build_filter(&mut query, &filter.filters); let params = build_filter(&mut query, &filter.filters);
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let s = conn.prep(&query).await.map_err(into_error)?; let limit = self.timeouts.maintenance;
let result = tokio::time::timeout(limit, async {
let s = conn.prep(&query).await.map_err(into_error)?;
match conn.exec_drop(s, params.clone()).await { match conn.exec_drop(s, params.clone()).await {
Ok(_) => return Ok(conn.affected_rows()), Ok(_) => return Ok(conn.affected_rows()),
Err(err) if is_timeout_error(&err) => (), Err(err) if is_timeout_error(&err) => (),
Err(err) => return Err(into_error(err)), Err(err) => return Err(into_error(err)),
} }
let mut chunk_size = DELETE_CHUNK_SIZE; let mut chunk_size = DELETE_CHUNK_SIZE;
let mut deleted = 0; let mut deleted = 0;
loop {
let s = conn
.prep(format!("{query} LIMIT {chunk_size}"))
.await
.map_err(into_error)?;
loop { loop {
match conn.exec_drop(&s, params.clone()).await { let s = conn
Ok(_) => { .prep(format!("{query} LIMIT {chunk_size}"))
let affected = conn.affected_rows(); .await
if affected == 0 { .map_err(into_error)?;
return Ok(deleted);
loop {
match conn.exec_drop(&s, params.clone()).await {
Ok(_) => {
let affected = conn.affected_rows();
if affected == 0 {
return Ok(deleted);
}
deleted += affected;
} }
deleted += affected; Err(err)
if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE =>
{
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
break;
}
Err(err) => return Err(into_error(err)),
} }
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
break;
}
Err(err) => return Err(into_error(err)),
} }
} }
} })
.await;
bounded(conn, result, limit)
} }
} }
crates/store/src/backend/mysql/write.rs
+100 -84
View File
@@ -6,7 +6,9 @@
* Modified by Coffey Labs in 2026 for INBUXA. * Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use super::{DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, MysqlStore, into_error, is_timeout_error}; use super::{
DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, MysqlStore, bounded, into_error, is_timeout_error,
};
use crate::{ use crate::{
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
SUBSPACE_REGISTRY_IDX, SUBSPACE_REGISTRY_IDX,
@@ -32,41 +34,45 @@ impl MysqlStore {
let start = Instant::now(); let start = Instant::now();
let mut retry_count = 0; let mut retry_count = 0;
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let limit = self.timeouts.query;
loop { let result = tokio::time::timeout(limit, async {
let err = match self.write_trx(&mut conn, &mut batch).await { loop {
Ok(result) => { let err = match self.write_trx(&mut conn, &mut batch).await {
return Ok(result); Ok(result) => {
} return Ok(result);
Err(err) => err,
};
let _ = conn.query_drop("ROLLBACK;").await;
match err {
CommitError::Mysql(Error::Server(err))
if [1062, 1213].contains(&err.code)
&& retry_count < MAX_COMMIT_ATTEMPTS
&& start.elapsed() < MAX_COMMIT_TIME => {}
/*CommitError::Retry => {
if retry_count > MAX_COMMIT_ATTEMPTS || start.elapsed() > MAX_COMMIT_TIME {
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.caused_by(trc::location!()));
} }
}*/ Err(err) => err,
CommitError::Mysql(err) => { };
return Err(into_error(err));
}
CommitError::Internal(err) => {
return Err(err);
}
}
let backoff = rand::rng().random_range(50..=300); let _ = conn.query_drop("ROLLBACK;").await;
tokio::time::sleep(Duration::from_millis(backoff)).await;
retry_count += 1; match err {
} CommitError::Mysql(Error::Server(err))
if [1062, 1213].contains(&err.code)
&& retry_count < MAX_COMMIT_ATTEMPTS
&& start.elapsed() < MAX_COMMIT_TIME => {}
/*CommitError::Retry => {
if retry_count > MAX_COMMIT_ATTEMPTS || start.elapsed() > MAX_COMMIT_TIME {
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.caused_by(trc::location!()));
}
}*/
CommitError::Mysql(err) => {
return Err(into_error(err));
}
CommitError::Internal(err) => {
return Err(err);
}
}
let backoff = rand::rng().random_range(50..=300);
tokio::time::sleep(Duration::from_millis(backoff)).await;
retry_count += 1;
}
})
.await;
bounded(conn, result, limit)
} }
async fn write_trx( async fn write_trx(
@@ -385,71 +391,81 @@ impl MysqlStore {
pub(crate) async fn purge_store(&self) -> trc::Result<()> { pub(crate) async fn purge_store(&self) -> trc::Result<()> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
for subspace in [SUBSPACE_QUOTA, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER] { let limit = self.timeouts.maintenance;
purge_table(&mut conn, char::from(subspace)).await?; let result = tokio::time::timeout(limit, async {
} for subspace in [SUBSPACE_QUOTA, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER] {
purge_table(&mut conn, char::from(subspace)).await?;
}
Ok(()) Ok(())
})
.await;
bounded(conn, result, limit)
} }
pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> { pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> {
let mut conn = self.conn().await?; let mut conn = self.conn().await?;
let table = char::from(from.subspace()); let limit = self.timeouts.maintenance;
let mut from = from.serialize(0); let result = tokio::time::timeout(limit, async {
let to = to.serialize(0); let table = char::from(from.subspace());
let mut from = from.serialize(0);
let to = to.serialize(0);
let delete = conn let delete = conn
.prep(format!("DELETE FROM {table} WHERE k >= ? AND k < ?")) .prep(format!("DELETE FROM {table} WHERE k >= ? AND k < ?"))
.await
.map_err(into_error)?;
match conn.exec_drop(&delete, (&from, &to)).await {
Ok(_) => return Ok(()),
Err(err) if is_timeout_error(&err) => (),
Err(err) => return Err(into_error(err)),
}
let mut chunk_size = DELETE_CHUNK_SIZE;
loop {
let boundary = conn
.prep(format!(
"SELECT k FROM {table} WHERE k >= ? AND k < ? ORDER BY k ASC LIMIT 1 OFFSET {chunk_size}"
))
.await .await
.map_err(into_error)?; .map_err(into_error)?;
match conn.exec_drop(&delete, (&from, &to)).await {
Ok(_) => return Ok(()),
Err(err) if is_timeout_error(&err) => (),
Err(err) => return Err(into_error(err)),
}
let mut chunk_size = DELETE_CHUNK_SIZE;
loop { loop {
let next = match conn let boundary = conn
.exec_first::<Vec<u8>, _, _>(&boundary, (&from, &to)) .prep(format!(
"SELECT k FROM {table} WHERE k >= ? AND k < ? ORDER BY k ASC LIMIT 1 OFFSET {chunk_size}"
))
.await .await
{ .map_err(into_error)?;
Ok(next) => next,
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
break;
}
Err(err) => return Err(into_error(err)),
};
match conn loop {
.exec_drop(&delete, (&from, next.as_ref().unwrap_or(&to))) let next = match conn
.await .exec_first::<Vec<u8>, _, _>(&boundary, (&from, &to))
{ .await
Ok(_) => (), {
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => { Ok(next) => next,
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE); Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
break; chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
} break;
Err(err) => return Err(into_error(err)), }
} Err(err) => return Err(into_error(err)),
};
match next { match conn
Some(next) => from = next, .exec_drop(&delete, (&from, next.as_ref().unwrap_or(&to)))
None => return Ok(()), .await
{
Ok(_) => (),
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
break;
}
Err(err) => return Err(into_error(err)),
}
match next {
Some(next) => from = next,
None => return Ok(()),
}
} }
} }
} })
.await;
bounded(conn, result, limit)
} }
} }
crates/store/src/backend/postgres/blob.rs
+57 -40
View File
@@ -2,13 +2,15 @@
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]> * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
* *
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*
* Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use std::ops::Range; use std::ops::Range;
use crate::backend::postgres::into_pool_error; use crate::backend::postgres::into_pool_error;
use super::{PostgresStore, into_error}; use super::{PostgresStore, bounded, into_error};
impl PostgresStore { impl PostgresStore {
pub(crate) async fn get_blob( pub(crate) async fn get_blob(
@@ -17,53 +19,68 @@ impl PostgresStore {
range: Range<usize>, range: Range<usize>,
) -> trc::Result<Option<Vec<u8>>> { ) -> trc::Result<Option<Vec<u8>>> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn let limit = self.timeouts.query;
.prepare_cached("SELECT v FROM t WHERE k = $1") let result = tokio::time::timeout(limit, async {
.await let s = conn
.map_err(into_error)?; .prepare_cached("SELECT v FROM t WHERE k = $1")
conn.query_opt(&s, &[&key]) .await
.await .map_err(into_error)?;
.and_then(|row| { conn.query_opt(&s, &[&key])
if let Some(row) = row { .await
Ok(Some(if range.start == 0 && range.end == usize::MAX { .and_then(|row| {
row.try_get::<_, Vec<u8>>(0)? if let Some(row) = row {
Ok(Some(if range.start == 0 && range.end == usize::MAX {
row.try_get::<_, Vec<u8>>(0)?
} else {
let bytes = row.try_get::<_, &[u8]>(0)?;
bytes
.get(range.start..std::cmp::min(bytes.len(), range.end))
.unwrap_or_default()
.to_vec()
}))
} else { } else {
let bytes = row.try_get::<_, &[u8]>(0)?; Ok(None)
bytes }
.get(range.start..std::cmp::min(bytes.len(), range.end)) })
.unwrap_or_default() .map_err(into_error)
.to_vec() })
})) .await;
} else { bounded(conn, result, limit)
Ok(None)
}
})
.map_err(into_error)
} }
pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> trc::Result<()> { pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> trc::Result<()> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn let limit = self.timeouts.query;
.prepare_cached( let result = tokio::time::timeout(limit, async {
"INSERT INTO t (k, v) VALUES ($1, $2) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v", let s = conn
) .prepare_cached(
.await "INSERT INTO t (k, v) VALUES ($1, $2) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v",
.map_err(into_error)?; )
conn.execute(&s, &[&key, &data]) .await
.await .map_err(into_error)?;
.map_err(into_error) conn.execute(&s, &[&key, &data])
.map(|_| ()) .await
.map_err(into_error)
.map(|_| ())
})
.await;
bounded(conn, result, limit)
} }
pub(crate) async fn delete_blob(&self, key: &[u8]) -> trc::Result<bool> { pub(crate) async fn delete_blob(&self, key: &[u8]) -> trc::Result<bool> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn let limit = self.timeouts.query;
.prepare_cached("DELETE FROM t WHERE k = $1") let result = tokio::time::timeout(limit, async {
.await let s = conn
.map_err(into_error)?; .prepare_cached("DELETE FROM t WHERE k = $1")
conn.execute(&s, &[&key]) .await
.await .map_err(into_error)?;
.map_err(into_error) conn.execute(&s, &[&key])
.map(|hits| hits > 0) .await
.map_err(into_error)
.map(|hits| hits > 0)
})
.await;
bounded(conn, result, limit)
} }
} }
crates/store/src/backend/postgres/lookup.rs
+32 -25
View File
@@ -2,6 +2,8 @@
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]> * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
* *
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*
* Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use crate::{QueryResult, QueryType, backend::postgres::into_pool_error}; use crate::{QueryResult, QueryType, backend::postgres::into_pool_error};
@@ -12,7 +14,7 @@ use tokio_postgres::types::{FromSql, ToSql, Type};
use crate::IntoRows; use crate::IntoRows;
use super::{PostgresStore, into_error}; use super::{PostgresStore, bounded, into_error};
impl PostgresStore { impl PostgresStore {
pub(crate) async fn sql_query<T: QueryResult>( pub(crate) async fn sql_query<T: QueryResult>(
@@ -21,33 +23,38 @@ impl PostgresStore {
params_: &[crate::Value<'_>], params_: &[crate::Value<'_>],
) -> trc::Result<T> { ) -> trc::Result<T> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn.prepare_cached(query).await.map_err(into_error)?; let limit = self.timeouts.query;
let params = params_ let result = tokio::time::timeout(limit, async {
.iter() let s = conn.prepare_cached(query).await.map_err(into_error)?;
.map(|v| v as &(dyn tokio_postgres::types::ToSql + Sync)) let params = params_
.collect::<Vec<_>>(); .iter()
.map(|v| v as &(dyn tokio_postgres::types::ToSql + Sync))
.collect::<Vec<_>>();
match T::query_type() { match T::query_type() {
QueryType::Execute => conn QueryType::Execute => conn
.execute(&s, params.as_slice()) .execute(&s, params.as_slice())
.await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_exec(r as usize))),
QueryType::Exists => {
let rows = conn.query_raw(&s, params).await.map_err(into_error)?;
pin_mut!(rows);
rows.try_next()
.await .await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_exists(r.is_some()))) .map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_exec(r as usize))),
QueryType::Exists => {
let rows = conn.query_raw(&s, params).await.map_err(into_error)?;
pin_mut!(rows);
rows.try_next()
.await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_exists(r.is_some())))
}
QueryType::QueryOne => conn
.query_opt(&s, params.as_slice())
.await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_query_one(r))),
QueryType::QueryAll => conn
.query(&s, params.as_slice())
.await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_query_all(r))),
} }
QueryType::QueryOne => conn })
.query_opt(&s, params.as_slice()) .await;
.await bounded(conn, result, limit)
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_query_one(r))),
QueryType::QueryAll => conn
.query(&s, params.as_slice())
.await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_query_all(r))),
}
} }
} }
crates/store/src/backend/postgres/main.rs
+81 -71
View File
@@ -6,7 +6,7 @@
* Modified by Coffey Labs in 2026 for INBUXA. * Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use super::{PostgresStore, into_error}; use super::{PostgresStore, bounded, into_error};
use crate::{ use crate::{
backend::postgres::{ backend::postgres::{
PsqlSearchField, into_pool_error, PsqlSearchField, into_pool_error,
@@ -119,6 +119,7 @@ impl PostgresStore {
Store::PostgreSQL(Arc::new(PostgresStore { Store::PostgreSQL(Arc::new(PostgresStore {
conn_pool: pool, conn_pool: pool,
ts_configs: ts_configs.clone(), ts_configs: ts_configs.clone(),
timeouts: Default::default(),
})), })),
replica.host, replica.host,
replica.port as u16, replica.port as u16,
@@ -129,6 +130,7 @@ impl PostgresStore {
let primary = Store::PostgreSQL(Arc::new(PostgresStore { let primary = Store::PostgreSQL(Arc::new(PostgresStore {
conn_pool: primary_pool, conn_pool: primary_pool,
ts_configs, ts_configs,
timeouts: Default::default(),
})); }));
// ST-1: no replicas, no change // ST-1: no replicas, no change
@@ -147,84 +149,92 @@ impl PostgresStore {
pub(crate) async fn create_storage_tables(&self) -> trc::Result<()> { pub(crate) async fn create_storage_tables(&self) -> trc::Result<()> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let limit = self.timeouts.maintenance;
let result = tokio::time::timeout(limit, async {
for table in [
SUBSPACE_ACL,
SUBSPACE_TASK_QUEUE,
SUBSPACE_DELETED_ITEMS,
SUBSPACE_SPAM_SAMPLES,
crate::SUBSPACE_INBUXA, // inbuxa: masked email
SUBSPACE_BLOB_LINK,
SUBSPACE_IN_MEMORY_VALUE,
SUBSPACE_PROPERTY,
SUBSPACE_REGISTRY,
SUBSPACE_REGISTRY_PK,
SUBSPACE_QUEUE_MESSAGE,
SUBSPACE_QUEUE_EVENT,
SUBSPACE_REPORT_OUT,
SUBSPACE_REPORT_IN,
SUBSPACE_LOGS,
SUBSPACE_BLOBS,
SUBSPACE_DIRECTORY,
SUBSPACE_TELEMETRY_SPAN,
SUBSPACE_TELEMETRY_METRIC,
] {
let table = char::from(table);
conn.execute(
&format!(
"CREATE TABLE IF NOT EXISTS {table} (
k BYTEA PRIMARY KEY,
v BYTEA NOT NULL
)"
),
&[],
)
.await
.map_err(into_error)?;
}
for table in [ for table in [SUBSPACE_INDEXES, SUBSPACE_REGISTRY_IDX] {
SUBSPACE_ACL, let table = char::from(table);
SUBSPACE_TASK_QUEUE, conn.execute(
SUBSPACE_DELETED_ITEMS, &format!(
SUBSPACE_SPAM_SAMPLES, "CREATE TABLE IF NOT EXISTS {table} (
crate::SUBSPACE_INBUXA, // inbuxa: masked email k BYTEA PRIMARY KEY
SUBSPACE_BLOB_LINK, )"
SUBSPACE_IN_MEMORY_VALUE, ),
SUBSPACE_PROPERTY, &[],
SUBSPACE_REGISTRY, )
SUBSPACE_REGISTRY_PK, .await
SUBSPACE_QUEUE_MESSAGE, .map_err(into_error)?;
SUBSPACE_QUEUE_EVENT, }
SUBSPACE_REPORT_OUT,
SUBSPACE_REPORT_IN, for table in [SUBSPACE_COUNTER, SUBSPACE_QUOTA, SUBSPACE_IN_MEMORY_COUNTER] {
SUBSPACE_LOGS, conn.execute(
SUBSPACE_BLOBS, &format!(
SUBSPACE_DIRECTORY, "CREATE TABLE IF NOT EXISTS {} (
SUBSPACE_TELEMETRY_SPAN,
SUBSPACE_TELEMETRY_METRIC,
] {
let table = char::from(table);
conn.execute(
&format!(
"CREATE TABLE IF NOT EXISTS {table} (
k BYTEA PRIMARY KEY, k BYTEA PRIMARY KEY,
v BYTEA NOT NULL v BIGINT NOT NULL DEFAULT 0
)" )",
), char::from(table)
&[], ),
) &[],
.await )
.map_err(into_error)?; .await
} .map_err(into_error)?;
}
for table in [SUBSPACE_INDEXES, SUBSPACE_REGISTRY_IDX] { Ok(())
let table = char::from(table); })
conn.execute( .await;
&format!( bounded(conn, result, limit)
"CREATE TABLE IF NOT EXISTS {table} (
k BYTEA PRIMARY KEY
)"
),
&[],
)
.await
.map_err(into_error)?;
}
for table in [SUBSPACE_COUNTER, SUBSPACE_QUOTA, SUBSPACE_IN_MEMORY_COUNTER] {
conn.execute(
&format!(
"CREATE TABLE IF NOT EXISTS {} (
k BYTEA PRIMARY KEY,
v BIGINT NOT NULL DEFAULT 0
)",
char::from(table)
),
&[],
)
.await
.map_err(into_error)?;
}
Ok(())
} }
pub(crate) async fn create_search_tables(&self) -> trc::Result<()> { pub(crate) async fn create_search_tables(&self) -> trc::Result<()> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let limit = self.timeouts.maintenance;
let result = tokio::time::timeout(limit, async {
create_search_tables::<EmailSearchField>(&conn).await?;
create_search_tables::<CalendarSearchField>(&conn).await?;
create_search_tables::<ContactSearchField>(&conn).await?;
//create_search_tables::<FileSearchField>(&conn).await?;
create_search_tables::<TracingSearchField>(&conn).await?;
create_search_tables::<EmailSearchField>(&conn).await?; Ok(())
create_search_tables::<CalendarSearchField>(&conn).await?; })
create_search_tables::<ContactSearchField>(&conn).await?; .await;
//create_search_tables::<FileSearchField>(&conn).await?; bounded(conn, result, limit)
create_search_tables::<TracingSearchField>(&conn).await?;
Ok(())
} }
} }
crates/store/src/backend/postgres/mod.rs
+33 -1
View File
@@ -6,6 +6,7 @@
* Modified by Coffey Labs in 2026 for INBUXA. * Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use crate::backend::query_timeout::QueryTimeouts;
use crate::{ use crate::{
search::{ search::{
CalendarSearchField, ContactSearchField, EmailSearchField, FileSearchField, SearchField, CalendarSearchField, ContactSearchField, EmailSearchField, FileSearchField, SearchField,
@@ -14,7 +15,8 @@ use crate::{
write::SearchIndex, write::SearchIndex,
}; };
use ahash::AHashSet; use ahash::AHashSet;
use deadpool_postgres::Pool; use deadpool_postgres::{Object, Pool};
use std::time::Duration;
use tokio_postgres::error::SqlState; use tokio_postgres::error::SqlState;
pub mod blob; pub mod blob;
@@ -28,6 +30,8 @@ pub mod write;
pub struct PostgresStore { pub struct PostgresStore {
pub(crate) conn_pool: Pool, pub(crate) conn_pool: Pool,
pub(crate) ts_configs: AHashSet<&'static str>, pub(crate) ts_configs: AHashSet<&'static str>,
/// inbuxa: client-side query limits (see backend::query_timeout)
pub(crate) timeouts: QueryTimeouts,
} }
#[inline(always)] #[inline(always)]
@@ -72,6 +76,34 @@ pub(crate) fn is_timeout_error(err: &tokio_postgres::Error) -> bool {
}) })
} }
/// inbuxa: the error for an operation that ran past its time limit.
pub(crate) fn query_timeout_error(limit: Duration) -> trc::Error {
trc::StoreEvent::PostgresqlError
.reason("Query timed out")
.details(format!(
"No answer from the database within {} s",
limit.as_secs()
))
}
/// inbuxa: ends an operation run on `conn` under `limit`. When it ran out,
/// the connection is taken out of the pool and closed: a query may still be
/// in flight on it, or a transaction open, so it can't be handed to the
/// next caller.
pub(crate) fn bounded<T>(
conn: Object,
result: Result<trc::Result<T>, tokio::time::error::Elapsed>,
limit: Duration,
) -> trc::Result<T> {
match result {
Ok(result) => result,
Err(_) => {
drop(Object::take(conn));
Err(query_timeout_error(limit))
}
}
}
#[inline(always)] #[inline(always)]
pub(crate) fn into_pool_error(err: deadpool_postgres::PoolError) -> trc::Error { pub(crate) fn into_pool_error(err: deadpool_postgres::PoolError) -> trc::Error {
match err { match err {
crates/store/src/backend/postgres/read.rs
+108 -63
View File
@@ -2,9 +2,11 @@
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]> * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
* *
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*
* Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use super::{PostgresStore, into_error, is_timeout_error}; use super::{PostgresStore, bounded, into_error, is_timeout_error, query_timeout_error};
use crate::{ use crate::{
Deserialize, IterateParams, Key, ValueKey, backend::postgres::into_pool_error, Deserialize, IterateParams, Key, ValueKey, backend::postgres::into_pool_error,
write::ValueClass, write::ValueClass,
@@ -17,40 +19,50 @@ impl PostgresStore {
U: Deserialize + 'static, U: Deserialize + 'static,
{ {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn let limit = self.timeouts.query;
.prepare_cached(&format!( let result = tokio::time::timeout(limit, async {
"SELECT v FROM {} WHERE k = $1", let s = conn
char::from(key.subspace()) .prepare_cached(&format!(
)) "SELECT v FROM {} WHERE k = $1",
.await char::from(key.subspace())
.map_err(into_error)?; ))
let key = key.serialize(0); .await
conn.query_opt(&s, &[&key]) .map_err(into_error)?;
.await let key = key.serialize(0);
.map_err(into_error) conn.query_opt(&s, &[&key])
.and_then(|r| { .await
if let Some(r) = r { .map_err(into_error)
Ok(Some(U::deserialize_with_key(&key, r.get(0))?)) .and_then(|r| {
} else { if let Some(r) = r {
Ok(None) Ok(Some(U::deserialize_with_key(&key, r.get(0))?))
} } else {
}) Ok(None)
}
})
})
.await;
bounded(conn, result, limit)
} }
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> { pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn let limit = self.timeouts.query;
.prepare_cached(&format!( let result = tokio::time::timeout(limit, async {
"SELECT 1 FROM {} WHERE k = $1", let s = conn
char::from(key.subspace()) .prepare_cached(&format!(
)) "SELECT 1 FROM {} WHERE k = $1",
.await char::from(key.subspace())
.map_err(into_error)?; ))
let key = key.serialize(0); .await
conn.query_opt(&s, &[&key]) .map_err(into_error)?;
.await let key = key.serialize(0);
.map_err(into_error) conn.query_opt(&s, &[&key])
.map(|r| r.is_some()) .await
.map_err(into_error)
.map(|r| r.is_some())
})
.await;
bounded(conn, result, limit)
} }
pub(crate) async fn iterate<T: Key>( pub(crate) async fn iterate<T: Key>(
@@ -64,44 +76,65 @@ impl PostgresStore {
let end = params.end.serialize(0); let end = params.end.serialize(0);
let keys = if params.values { "k, v" } else { "k" }; let keys = if params.values { "k, v" } else { "k" };
let s = conn // inbuxa: a scan may run for hours, so the query limit bounds each
.prepare_cached(&match (params.first, params.ascending) { // wait for the database (preparing, the query starting, the next
(true, true) => { // row) rather than the scan. A wait that runs out closes the
format!( // connection.
"SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC LIMIT 1" let limit = self.timeouts.query;
) let query = match (params.first, params.ascending) {
} (true, true) => {
(true, false) => { format!(
format!( "SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC LIMIT 1"
)
}
(true, false) => {
format!(
"SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC LIMIT 1" "SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC LIMIT 1"
) )
} }
(false, true) => { (false, true) => {
format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC") format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC")
} }
(false, false) => { (false, false) => {
format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC") format!("SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k DESC")
} }
}) };
.await.map_err(into_error)?; let s = match tokio::time::timeout(limit, conn.prepare_cached(&query)).await {
Ok(s) => s.map_err(into_error)?,
Err(_) => {
drop(deadpool_postgres::Object::take(conn));
return Err(query_timeout_error(limit));
}
};
let mut from = begin; let mut from = begin;
let mut to = end; let mut to = end;
let mut resume_key: Option<Vec<u8>> = None; let mut resume_key: Option<Vec<u8>> = None;
let mut stalled = false;
loop { loop {
let mut last_key = None; let mut last_key = None;
let mut timed_out = false; let mut timed_out = false;
{ {
let rows = conn let rows =
.query_raw(&s, &[&from, &to]) match tokio::time::timeout(limit, conn.query_raw(&s, &[&from, &to])).await {
.await Ok(rows) => rows.map_err(into_error)?,
.map_err(into_error)?; // Leaves the scan loop for the timeout below
Err(_) => break,
};
pin_mut!(rows); pin_mut!(rows);
loop { loop {
match rows.try_next().await { let next = match tokio::time::timeout(limit, rows.try_next()).await {
Ok(next) => next,
Err(_) => {
stalled = true;
break;
}
};
match next {
Ok(Some(row)) => { Ok(Some(row)) => {
let key = row.try_get::<_, &[u8]>(0).map_err(into_error)?; let key = row.try_get::<_, &[u8]>(0).map_err(into_error)?;
let value = if params.values { let value = if params.values {
@@ -132,6 +165,10 @@ impl PostgresStore {
} }
} }
if stalled {
break;
}
match last_key { match last_key {
Some(last_key) if timed_out => { Some(last_key) if timed_out => {
if params.ascending { if params.ascending {
@@ -144,6 +181,9 @@ impl PostgresStore {
_ => return Ok(()), _ => return Ok(()),
} }
} }
drop(deadpool_postgres::Object::take(conn));
Err(query_timeout_error(limit))
} }
pub(crate) async fn get_counter( pub(crate) async fn get_counter(
@@ -155,14 +195,19 @@ impl PostgresStore {
let key = key.serialize(0); let key = key.serialize(0);
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn let limit = self.timeouts.query;
.prepare_cached(&format!("SELECT v FROM {table} WHERE k = $1")) let result = tokio::time::timeout(limit, async {
.await let s = conn
.map_err(into_error)?; .prepare_cached(&format!("SELECT v FROM {table} WHERE k = $1"))
match conn.query_opt(&s, &[&key]).await { .await
Ok(Some(row)) => row.try_get(0).map_err(into_error), .map_err(into_error)?;
Ok(None) => Ok(0), match conn.query_opt(&s, &[&key]).await {
Err(e) => Err(into_error(e)), Ok(Some(row)) => row.try_get(0).map_err(into_error),
} Ok(None) => Ok(0),
Err(e) => Err(into_error(e)),
}
})
.await;
bounded(conn, result, limit)
} }
} }
crates/store/src/backend/postgres/search.rs
+151 -136
View File
@@ -10,8 +10,8 @@ use crate::{
backend::{ backend::{
MAX_TOKEN_LENGTH, MAX_TOKEN_LENGTH,
postgres::{ postgres::{
DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, PostgresStore, PsqlSearchField, into_error, DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, PostgresStore, PsqlSearchField, bounded,
into_pool_error, is_timeout_error, into_error, into_pool_error, is_timeout_error,
}, },
}, },
search::{ search::{
@@ -36,125 +36,130 @@ impl PostgresStore {
pub async fn index(&self, documents: Vec<IndexDocument>) -> trc::Result<()> { pub async fn index(&self, documents: Vec<IndexDocument>) -> trc::Result<()> {
let mut conn = self.conn_pool.get().await.map_err(into_pool_error)?; let mut conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let trx = conn let limit = self.timeouts.query;
.build_transaction() let result = tokio::time::timeout(limit, async {
.isolation_level(IsolationLevel::ReadCommitted) let trx = conn
.start() .build_transaction()
.await .isolation_level(IsolationLevel::ReadCommitted)
.map_err(into_error)?; .start()
.await
.map_err(into_error)?;
for document in documents { for document in documents {
let index = document.index; let index = document.index;
let primary_keys = index.primary_keys(); let primary_keys = index.primary_keys();
let all_fields = index.all_fields(); let all_fields = index.all_fields();
let fields = document.fields; let fields = document.fields;
// inbuxa: keyword text (addresses, contact fields, ...) is split into // inbuxa: keyword text (addresses, contact fields, ...) is split into
// words before it reaches the text parser, see keyword_terms(). // words before it reaches the text parser, see keyword_terms().
let keywords = primary_keys let keywords = primary_keys
.iter() .iter()
.chain(all_fields) .chain(all_fields)
.map(|field| match fields.get(field) { .map(|field| match fields.get(field) {
Some(SearchValue::Text { Some(SearchValue::Text {
value, value,
language: Language::None, language: Language::None,
}) if field.is_text() => Some(keyword_terms(value)), }) if field.is_text() => Some(keyword_terms(value)),
_ => None, _ => None,
}) })
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let mut values = Vec::with_capacity(fields.len() + 2); let mut values = Vec::with_capacity(fields.len() + 2);
let mut query = format!("INSERT INTO {} (", index.psql_table()); let mut query = format!("INSERT INTO {} (", index.psql_table());
for (i, field) in primary_keys.iter().chain(all_fields).enumerate() { for (i, field) in primary_keys.iter().chain(all_fields).enumerate() {
if i > 0 { if i > 0 {
query.push(','); query.push(',');
} }
query.push_str(field.column()); query.push_str(field.column());
if let Some(sort_column) = field.sort_column() { if let Some(sort_column) = field.sort_column() {
query.push(','); query.push(',');
query.push_str(sort_column); query.push_str(sort_column);
} }
}
query.push_str(") VALUES (");
for (i, field) in primary_keys.iter().chain(all_fields).enumerate() {
if i > 0 {
query.push(',');
} }
if let Some(value) = fields.get(field) { query.push_str(") VALUES (");
let value_ref = format!("${}", values.len() + 1);
let (text_len, language) = if let SearchValue::Text { value, language } = value
{
(value.len(), self.ts_config(language))
} else {
(0, PG_UNSTEMMED_LANG)
};
if let Some(keywords) = &keywords[i] { for (i, field) in primary_keys.iter().chain(all_fields).enumerate() {
let _ = write!(&mut query, "to_tsvector('{language}',{value_ref})"); if i > 0 {
values.push(keywords as &(dyn ToSql + Sync)); query.push(',');
if field.sort_column().is_some() { }
let value_ref = format!("${}", values.len() + 1);
if text_len > 255 { if let Some(value) = fields.get(field) {
let _ = write!(&mut query, ",left({value_ref},255)"); let value_ref = format!("${}", values.len() + 1);
let (text_len, language) =
if let SearchValue::Text { value, language } = value {
(value.len(), self.ts_config(language))
} else { } else {
let _ = write!(&mut query, ",{value_ref}"); (0, PG_UNSTEMMED_LANG)
};
if let Some(keywords) = &keywords[i] {
let _ = write!(&mut query, "to_tsvector('{language}',{value_ref})");
values.push(keywords as &(dyn ToSql + Sync));
if field.sort_column().is_some() {
let value_ref = format!("${}", values.len() + 1);
if text_len > 255 {
let _ = write!(&mut query, ",left({value_ref},255)");
} else {
let _ = write!(&mut query, ",{value_ref}");
}
values.push(value as &(dyn ToSql + Sync));
} }
values.push(value as &(dyn ToSql + Sync)); continue;
} } else if field.is_text() {
continue; let _ = write!(&mut query, "to_tsvector('{language}',{value_ref})");
} else if field.is_text() { } else if text_len > 512 {
let _ = write!(&mut query, "to_tsvector('{language}',{value_ref})"); query.push_str("left(");
} else if text_len > 512 {
query.push_str("left(");
query.push_str(&value_ref);
query.push_str(",512)");
} else {
query.push_str(&value_ref);
}
if field.sort_column().is_some() {
if text_len > 255 {
query.push_str(",left(");
query.push_str(&value_ref); query.push_str(&value_ref);
query.push_str(",255)"); query.push_str(",512)");
} else { } else {
query.push(',');
query.push_str(&value_ref); query.push_str(&value_ref);
} }
}
values.push(value as &(dyn ToSql + Sync)); if field.sort_column().is_some() {
} else { if text_len > 255 {
query.push_str("NULL"); query.push_str(",left(");
if field.sort_column().is_some() { query.push_str(&value_ref);
query.push_str(",NULL"); query.push_str(",255)");
} else {
query.push(',');
query.push_str(&value_ref);
}
}
values.push(value as &(dyn ToSql + Sync));
} else {
query.push_str("NULL");
if field.sort_column().is_some() {
query.push_str(",NULL");
}
} }
} }
}
query.push_str(") ON CONFLICT ("); query.push_str(") ON CONFLICT (");
for (i, pkey) in primary_keys.iter().enumerate() { for (i, pkey) in primary_keys.iter().enumerate() {
if i > 0 { if i > 0 {
query.push(','); query.push(',');
}
query.push_str(pkey.column());
} }
query.push_str(pkey.column()); query.push_str(") DO UPDATE SET ");
} for (i, field) in all_fields.iter().enumerate() {
query.push_str(") DO UPDATE SET "); if i > 0 {
for (i, field) in all_fields.iter().enumerate() { query.push(',');
if i > 0 { }
query.push(','); let column = field.column();
let _ = write!(&mut query, "{column} = EXCLUDED.{column}");
} }
let column = field.column();
let _ = write!(&mut query, "{column} = EXCLUDED.{column}"); trx.execute(&query, &values).await.map_err(into_error)?;
} }
trx.execute(&query, &values).await.map_err(into_error)?; trx.commit().await.map_err(into_error)
} })
.await;
trx.commit().await.map_err(into_error) bounded(conn, result, limit)
} }
pub async fn query<R: SearchDocumentId>( pub async fn query<R: SearchDocumentId>(
@@ -170,16 +175,21 @@ impl PostgresStore {
build_sort(&mut query, sort); build_sort(&mut query, sort);
} }
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn.prepare_cached(&query).await.map_err(into_error)?; let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn.prepare_cached(&query).await.map_err(into_error)?;
conn.query(&s, params.as_slice()) conn.query(&s, params.as_slice())
.await .await
.and_then(|rows| { .and_then(|rows| {
rows.into_iter() rows.into_iter()
.map(|row| row.try_get::<_, DocId>(0).map(|v| R::from_u64(v.0))) .map(|row| row.try_get::<_, DocId>(0).map(|v| R::from_u64(v.0)))
.collect::<Result<Vec<R>, _>>() .collect::<Result<Vec<R>, _>>()
}) })
.map_err(into_error) .map_err(into_error)
})
.await;
bounded(conn, result, limit)
} }
pub async fn unindex(&self, filter: SearchQuery) -> trc::Result<u64> { pub async fn unindex(&self, filter: SearchQuery) -> trc::Result<u64> {
@@ -189,40 +199,45 @@ impl PostgresStore {
let params = self.build_filter(&mut where_clause, &filter.filters); let params = self.build_filter(&mut where_clause, &filter.filters);
let params = params.iter().map(SqlParam::as_sql).collect::<Vec<_>>(); let params = params.iter().map(SqlParam::as_sql).collect::<Vec<_>>();
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let s = conn let limit = self.timeouts.maintenance;
.prepare_cached(&format!("DELETE FROM {table}{where_clause}")) let result = tokio::time::timeout(limit, async {
.await
.map_err(into_error)?;
match conn.execute(&s, params.as_slice()).await {
Ok(deleted) => return Ok(deleted),
Err(err) if is_timeout_error(&err) => (),
Err(err) => return Err(into_error(err)),
}
let mut chunk_size = DELETE_CHUNK_SIZE;
let mut deleted = 0;
loop {
let s = conn let s = conn
.prepare_cached(&format!( .prepare_cached(&format!("DELETE FROM {table}{where_clause}"))
"DELETE FROM {table} WHERE ctid IN (SELECT ctid FROM {table}{where_clause} LIMIT {chunk_size})"
))
.await .await
.map_err(into_error)?; .map_err(into_error)?;
match conn.execute(&s, params.as_slice()).await {
Ok(deleted) => return Ok(deleted),
Err(err) if is_timeout_error(&err) => (),
Err(err) => return Err(into_error(err)),
}
let mut chunk_size = DELETE_CHUNK_SIZE;
let mut deleted = 0;
loop { loop {
match conn.execute(&s, params.as_slice()).await { let s = conn
Ok(0) => return Ok(deleted), .prepare_cached(&format!(
Ok(affected) => deleted += affected, "DELETE FROM {table} WHERE ctid IN (SELECT ctid FROM {table}{where_clause} LIMIT {chunk_size})"
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => { ))
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE); .await
break; .map_err(into_error)?;
loop {
match conn.execute(&s, params.as_slice()).await {
Ok(0) => return Ok(deleted),
Ok(affected) => deleted += affected,
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
break;
}
Err(err) => return Err(into_error(err)),
} }
Err(err) => return Err(into_error(err)),
} }
} }
} })
.await;
bounded(conn, result, limit)
} }
fn build_filter<'x>( fn build_filter<'x>(
crates/store/src/backend/postgres/write.rs
+106 -90
View File
@@ -2,9 +2,11 @@
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]> * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
* *
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*
* Modified by Coffey Labs in 2026 for INBUXA.
*/ */
use super::{PostgresStore, into_error, is_timeout_error}; use super::{PostgresStore, bounded, into_error, is_timeout_error};
use crate::{ use crate::{
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA, IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
SUBSPACE_REGISTRY_IDX, SUBSPACE_REGISTRY_IDX,
@@ -30,48 +32,53 @@ enum CommitError {
impl PostgresStore { impl PostgresStore {
pub(crate) async fn write(&self, mut batch: Batch<'_>) -> trc::Result<AssignedIds> { pub(crate) async fn write(&self, mut batch: Batch<'_>) -> trc::Result<AssignedIds> {
let mut conn = self.conn_pool.get().await.map_err(into_pool_error)?; let mut conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let start = Instant::now(); let limit = self.timeouts.query;
let mut retry_count = 0; let result = tokio::time::timeout(limit, async {
let start = Instant::now();
let mut retry_count = 0;
loop { loop {
match self.write_trx(&mut conn, &mut batch).await { match self.write_trx(&mut conn, &mut batch).await {
Ok(result) => { Ok(result) => {
return Ok(result); return Ok(result);
}
Err(err) => {
match err {
CommitError::Postgres(err) => match err.code() {
Some(
&SqlState::T_R_SERIALIZATION_FAILURE
| &SqlState::T_R_DEADLOCK_DETECTED,
) if retry_count < MAX_COMMIT_ATTEMPTS
&& start.elapsed() < MAX_COMMIT_TIME => {}
Some(&SqlState::UNIQUE_VIOLATION) => {
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.reason("Unique violation")
.caused_by(trc::location!()));
}
_ => return Err(into_error(err)),
},
CommitError::Internal(err) => return Err(err),
/*CommitError::Retry => {
if retry_count > MAX_COMMIT_ATTEMPTS
|| start.elapsed() > MAX_COMMIT_TIME
{
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.caused_by(trc::location!()));
}
}*/
} }
Err(err) => {
match err {
CommitError::Postgres(err) => match err.code() {
Some(
&SqlState::T_R_SERIALIZATION_FAILURE
| &SqlState::T_R_DEADLOCK_DETECTED,
) if retry_count < MAX_COMMIT_ATTEMPTS
&& start.elapsed() < MAX_COMMIT_TIME => {}
Some(&SqlState::UNIQUE_VIOLATION) => {
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.reason("Unique violation")
.caused_by(trc::location!()));
}
_ => return Err(into_error(err)),
},
CommitError::Internal(err) => return Err(err),
/*CommitError::Retry => {
if retry_count > MAX_COMMIT_ATTEMPTS
|| start.elapsed() > MAX_COMMIT_TIME
{
return Err(trc::StoreEvent::AssertValueFailed
.into_err()
.caused_by(trc::location!()));
}
}*/
}
let backoff = rand::rng().random_range(50..=300); let backoff = rand::rng().random_range(50..=300);
tokio::time::sleep(Duration::from_millis(backoff)).await; tokio::time::sleep(Duration::from_millis(backoff)).await;
retry_count += 1; retry_count += 1;
}
} }
} }
} })
.await;
bounded(conn, result, limit)
} }
async fn write_trx( async fn write_trx(
@@ -393,72 +400,81 @@ impl PostgresStore {
pub(crate) async fn purge_store(&self) -> trc::Result<()> { pub(crate) async fn purge_store(&self) -> trc::Result<()> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let limit = self.timeouts.maintenance;
let result = tokio::time::timeout(limit, async {
for subspace in [SUBSPACE_QUOTA, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER] {
purge_table(&conn, char::from(subspace)).await?;
}
for subspace in [SUBSPACE_QUOTA, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER] { Ok(())
purge_table(&conn, char::from(subspace)).await?; })
} .await;
bounded(conn, result, limit)
Ok(())
} }
pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> { pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?; let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let table = char::from(from.subspace()); let limit = self.timeouts.maintenance;
let mut from = from.serialize(0); let result = tokio::time::timeout(limit, async {
let to = to.serialize(0); let table = char::from(from.subspace());
let mut from = from.serialize(0);
let to = to.serialize(0);
let delete = conn let delete = conn
.prepare_cached(&format!("DELETE FROM {table} WHERE k >= $1 AND k < $2")) .prepare_cached(&format!("DELETE FROM {table} WHERE k >= $1 AND k < $2"))
.await
.map_err(into_error)?;
match conn.execute(&delete, &[&from, &to]).await {
Ok(_) => return Ok(()),
Err(err) if is_timeout_error(&err) => (),
Err(err) => return Err(into_error(err)),
}
let mut chunk_size = DELETE_CHUNK_SIZE;
loop {
let boundary = conn
.prepare_cached(&format!(
"SELECT k FROM {table} WHERE k >= $1 AND k < $2 ORDER BY k ASC LIMIT 1 OFFSET {chunk_size}"
))
.await .await
.map_err(into_error)?; .map_err(into_error)?;
match conn.execute(&delete, &[&from, &to]).await {
Ok(_) => return Ok(()),
Err(err) if is_timeout_error(&err) => (),
Err(err) => return Err(into_error(err)),
}
let mut chunk_size = DELETE_CHUNK_SIZE;
loop { loop {
let next = match conn.query_opt(&boundary, &[&from, &to]).await { let boundary = conn
Ok(next) => match next { .prepare_cached(&format!(
Some(row) => Some(row.try_get::<_, Vec<u8>>(0).map_err(into_error)?), "SELECT k FROM {table} WHERE k >= $1 AND k < $2 ORDER BY k ASC LIMIT 1 OFFSET {chunk_size}"
None => None, ))
},
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
break;
}
Err(err) => return Err(into_error(err)),
};
match conn
.execute(&delete, &[&from, next.as_ref().unwrap_or(&to)])
.await .await
{ .map_err(into_error)?;
Ok(_) => (),
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
break;
}
Err(err) => return Err(into_error(err)),
}
match next { loop {
Some(next) => from = next, let next = match conn.query_opt(&boundary, &[&from, &to]).await {
None => return Ok(()), Ok(next) => match next {
Some(row) => Some(row.try_get::<_, Vec<u8>>(0).map_err(into_error)?),
None => None,
},
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
break;
}
Err(err) => return Err(into_error(err)),
};
match conn
.execute(&delete, &[&from, next.as_ref().unwrap_or(&to)])
.await
{
Ok(_) => (),
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
break;
}
Err(err) => return Err(into_error(err)),
}
match next {
Some(next) => from = next,
None => return Ok(()),
}
} }
} }
} })
.await;
bounded(conn, result, limit)
} }
} }
crates/store/src/backend/query_timeout.rs
+77
View File
@@ -0,0 +1,77 @@
/*
* SPDX-FileCopyrightText: 2026 Coffey Labs
*
* SPDX-License-Identifier: AGPL-3.0-only
*/
//! Client-side limits on SQL queries.
//!
//! The pool timeouts bound getting a connection, not using one. A database
//! that stops answering while the TCP connection stays up (a paused
//! container, a hung server whose kernel still acknowledges keepalives)
//! left a query on a checked-out connection waiting for as long as it took.
//! A server-side statement_timeout can't help there: the server that would
//! enforce it is the one not answering. So each operation on a PostgreSQL
//! or MySQL connection runs under a time limit here, and a connection whose
//! operation ran out is closed rather than put back in the pool, since its
//! protocol state is unknown.
//!
//! Two limits:
//! - `query`, two minutes, for request-path work: reads, writes, blob
//! transfers, search queries and document indexing. Those take
//! milliseconds; two minutes leaves room for a large blob over a slow
//! link and still ends a hang.
//! - `maintenance`, thirty minutes, for work that legitimately runs long in
//! one statement: range deletes (account removal, purges), unindexing,
//! and creating tables and indexes at startup.
//!
//! Iterating over a range (exports, reindexing, maintenance scans) can run
//! for hours, so there the `query` limit applies to each wait for the next
//! row instead of the whole scan.
use std::time::Duration;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct QueryTimeouts {
pub query: Duration,
pub maintenance: Duration,
}
impl QueryTimeouts {
pub const QUERY: Duration = Duration::from_secs(120);
pub const MAINTENANCE: Duration = Duration::from_secs(30 * 60);
}
impl Default for QueryTimeouts {
fn default() -> Self {
Self {
query: Self::QUERY,
maintenance: Self::MAINTENANCE,
}
}
}
#[cfg(feature = "test_mode")]
impl crate::Store {
/// Sets the query limits of a SQL store that was just built (tests only:
/// the limits aren't configurable).
pub fn with_query_timeouts(self, timeouts: QueryTimeouts) -> Self {
match self {
#[cfg(feature = "postgres")]
crate::Store::PostgreSQL(mut store) => {
std::sync::Arc::get_mut(&mut store)
.expect("store already shared")
.timeouts = timeouts;
crate::Store::PostgreSQL(store)
}
#[cfg(feature = "mysql")]
crate::Store::MySQL(mut store) => {
std::sync::Arc::get_mut(&mut store)
.expect("store already shared")
.timeouts = timeouts;
crate::Store::MySQL(store)
}
store => store,
}
}
}
tests/src/store/pool_timeout.rs
+283 -3
View File
@@ -9,11 +9,31 @@
//! the pool's timeouts. Upstream's pools had none, so the worker waited for //! the pool's timeouts. Upstream's pools had none, so the worker waited for
//! good. No database is needed: a local listener that never answers plays //! good. No database is needed: a local listener that never answers plays
//! the server. //! the server.
//!
//! inbuxa: the same for a database that stops answering while connections
//! are already open (a paused container): a query on a checked-out
//! connection ends within the query limit, the store works again once the
//! database is back, and /healthz/ready says 503 in between while
//! /healthz/live stays 200. These need the local test databases; a proxy
//! that can stop forwarding plays the pause.
use registry::schema::structs::DataStore; use registry::schema::structs::DataStore;
use std::time::{Duration, Instant}; use std::{
use store::{Store, ValueKey, write::ValueClass}; sync::{
use tokio::net::TcpListener; Arc,
atomic::{AtomicBool, Ordering},
},
time::{Duration, Instant},
};
use store::{
IterateParams, Store, ValueKey,
backend::query_timeout::QueryTimeouts,
write::{BatchBuilder, ValueClass},
};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
};
/// Accepts connections on a local port and never sends a byte. /// Accepts connections on a local port and never sends a byte.
async fn silent_server() -> u16 { async fn silent_server() -> u16 {
@@ -94,3 +114,263 @@ pub async fn mysql_pool_timeout() {
) )
.await; .await;
} }
/// A TCP proxy to a local port that can stop forwarding, in both
/// directions, while keeping every connection open: a paused server whose
/// kernel still keeps the connections up.
struct PausableProxy {
port: u16,
paused: Arc<AtomicBool>,
}
impl PausableProxy {
async fn start(upstream: u16) -> Self {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let paused = Arc::new(AtomicBool::new(false));
let paused_ = paused.clone();
tokio::spawn(async move {
while let Ok((client, _)) = listener.accept().await {
let Ok(server) = TcpStream::connect(("127.0.0.1", upstream)).await else {
continue;
};
let (client_rx, client_tx) = client.into_split();
let (server_rx, server_tx) = server.into_split();
tokio::spawn(forward(client_rx, server_tx, paused_.clone()));
tokio::spawn(forward(server_rx, client_tx, paused_.clone()));
}
});
PausableProxy { port, paused }
}
fn pause(&self, paused: bool) {
self.paused.store(paused, Ordering::SeqCst);
}
}
async fn forward(
mut from: tokio::net::tcp::OwnedReadHalf,
mut to: tokio::net::tcp::OwnedWriteHalf,
paused: Arc<AtomicBool>,
) {
let mut buf = vec![0u8; 16384];
loop {
while paused.load(Ordering::SeqCst) {
tokio::time::sleep(Duration::from_millis(20)).await;
}
let n = match from.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => n,
};
// Hold what arrived while paused until the pause ends
while paused.load(Ordering::SeqCst) {
tokio::time::sleep(Duration::from_millis(20)).await;
}
if to.write_all(&buf[..n]).await.is_err() {
return;
}
}
}
const TEST_LIMITS: QueryTimeouts = QueryTimeouts {
query: Duration::from_secs(2),
maintenance: Duration::from_secs(3),
};
/// Opens `connections` pooled connections at once, so the operations that
/// follow find one idle and check it out.
async fn warm(store: &Store, connections: usize) {
let reads = (0..connections).map(|_| async {
store
.get_value::<u64>(ValueKey::from(ValueClass::Property(0)))
.await
.unwrap();
});
futures::future::join_all(reads).await;
}
/// With the database paused, reads, scans and writes on connections the
/// pool already holds end in an error within the query limit; once it is
/// back, the store works again.
async fn assert_queries_time_out(store: Store, proxy: &PausableProxy) {
store.create_tables().await.unwrap();
warm(&store, 4).await;
// mysql_async resets a connection on its way back to the pool; let
// those finish, or the connections are stuck in the reset when the
// pause starts and the pool's own wait timeout answers instead
tokio::time::sleep(Duration::from_secs(1)).await;
proxy.pause(true);
let key = || ValueKey::from(ValueClass::Property(0));
let limit = TEST_LIMITS.query;
for (what, op) in [("read", 0), ("scan", 1), ("write", 2)] {
let started = Instant::now();
let result = tokio::time::timeout(Duration::from_secs(20), async {
match op {
0 => store.get_value::<u64>(key()).await.map(|_| ()),
1 => {
store
.iterate(
IterateParams::new(
ValueKey::from(ValueClass::Property(0)),
ValueKey::from(ValueClass::Property(u8::MAX)),
),
|_, _| Ok(true),
)
.await
}
_ => {
let mut batch = BatchBuilder::new();
batch
.with_account_id(u32::MAX - 7)
.with_collection(types::collection::Collection::Email)
.with_document(0)
.set(ValueClass::Property(0), 1u64.to_be_bytes().to_vec());
store.write(batch.build_all()).await.map(|_| ())
}
}
})
.await;
let elapsed = started.elapsed();
match result {
Ok(Err(err)) => {
let err = format!("{err:?}");
println!("Paused database, {what}: {err} after {elapsed:?}");
assert!(err.contains("Query timed out"), "{what}: {err}");
assert!(
elapsed >= limit && elapsed < limit * 3,
"{what} ended after {elapsed:?}"
);
}
Ok(Ok(())) => panic!("{what} succeeded against a paused database"),
Err(_) => panic!("{what} still waiting after {elapsed:?}"),
}
}
proxy.pause(false);
tokio::time::timeout(Duration::from_secs(20), store.get_value::<u64>(key()))
.await
.expect("still waiting after the database came back")
.expect("the store didn't recover");
}
#[cfg(feature = "postgres")]
#[tokio::test(flavor = "multi_thread")]
pub async fn postgres_query_timeout() {
println!("Running PostgreSQL query timeout test...");
let DataStore::PostgreSql(mut config) =
crate::utils::storage::build_data_store("PostgreSql", "").await
else {
unreachable!()
};
let proxy = PausableProxy::start(config.port as u16).await;
config.host = "127.0.0.1".into();
config.port = proxy.port as u64;
// New connections through the paused proxy give up as quickly
config.timeout = Some(TEST_LIMITS.query.into());
let store = Store::build(DataStore::PostgreSql(config))
.await
.unwrap()
.with_query_timeouts(TEST_LIMITS);
assert_queries_time_out(store, &proxy).await;
}
#[cfg(feature = "mysql")]
#[tokio::test(flavor = "multi_thread")]
pub async fn mysql_query_timeout() {
println!("Running MySQL query timeout test...");
let DataStore::MySql(mut config) = crate::utils::storage::build_data_store("MySql", "").await
else {
unreachable!()
};
let proxy = PausableProxy::start(config.port as u16).await;
config.host = "127.0.0.1".into();
config.port = proxy.port as u64;
let store = Store::build(DataStore::MySql(config))
.await
.unwrap()
.with_query_timeouts(TEST_LIMITS);
assert_queries_time_out(store, &proxy).await;
}
/// /healthz/ready follows the data store; /healthz/live doesn't.
#[cfg(feature = "postgres")]
#[tokio::test(flavor = "multi_thread")]
pub async fn postgres_readiness() {
use crate::utils::server::TestServerBuilder;
use registry::schema::enums::NetworkListenerProtocol;
const HTTP_PORT: u16 = 11_320;
if std::env::var("STORE").as_deref() != Ok("PostgreSql") {
println!("Skipping the readiness test: it runs with STORE=PostgreSql.");
return;
}
println!("Running readiness test...");
let test = TestServerBuilder::new("postgres_readiness")
.await
.with_listener(NetworkListenerProtocol::Http, "http", HTTP_PORT, true)
.await
.build()
.await;
// Point the running node's data store at the database through the proxy
let DataStore::PostgreSql(mut config) =
crate::utils::storage::build_data_store("PostgreSql", "").await
else {
unreachable!()
};
let proxy = PausableProxy::start(config.port as u16).await;
config.host = "127.0.0.1".into();
config.port = proxy.port as u64;
config.timeout = Some(TEST_LIMITS.query.into());
let store = Store::build(DataStore::PostgreSql(config))
.await
.unwrap()
.with_query_timeouts(TEST_LIMITS);
let inner = &test.server.inner;
let mut core = inner.shared_core.load_full().as_ref().clone();
core.storage.data = store;
inner.shared_core.store(Arc::new(core));
let health = |path: &'static str| async move {
reqwest::Client::builder()
.danger_accept_invalid_certs(true)
.timeout(Duration::from_secs(10))
.build()
.unwrap()
.get(format!("https://127.0.0.1:{HTTP_PORT}/healthz/{path}"))
.send()
.await
.unwrap()
.status()
.as_u16()
};
let wait_for = |path: &'static str, status: u16| async move {
let started = Instant::now();
loop {
let got = health(path).await;
if got == status {
println!("/healthz/{path}: {got} after {:?}", started.elapsed());
return;
}
assert!(
started.elapsed() < Duration::from_secs(20),
"/healthz/{path} still {got}, expected {status}"
);
tokio::time::sleep(Duration::from_millis(250)).await;
}
};
wait_for("ready", 200).await;
proxy.pause(true);
wait_for("ready", 503).await;
assert_eq!(health("live").await, 200);
proxy.pause(false);
wait_for("ready", 200).await;
assert_eq!(health("live").await, 200);
if test.is_reset() {
test.temp_dir.delete();
}
}