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

Merge pull request 'SQL queries time out; readiness follows the data store' (#45) from fix/query-timeouts into main
ci / fork-checks (push) Successful in 55s
Details
ci / build (push) Canceled after 11m21s
Details

Browse Source
This commit was merged in pull request #45.
This commit is contained in:
jcoffey-dev
2026-09-25 00:45:09 +00:00
parent 181ab1c140 08f29926d4
commit ad58c35f39
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,
queue_status: true.into(),
settings_reload: Default::default(),
store_health: Default::default(),
applications,
logos: Default::default(),
smtp_connectors: TlsConnectors::try_new().failed("Failed to build TLS connectors"),
@@ -237,6 +238,7 @@ impl Default for Data {
registry_id_gen: Default::default(),
queue_status: true.into(),
settings_reload: Default::default(),
store_health: Default::default(),
applications: WebApplications::new(),
logos: Default::default(),
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,
// inbuxa: coalesces the settings reloads registry writes trigger
pub settings_reload: cache::reload::SettingsReloadGate,
// inbuxa: the readiness probe's cached answer
pub store_health: storage::ready::StoreHealth,
pub applications: WebApplications,
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 index;
pub mod quota;
pub mod ready; // inbuxa: readiness follows the data store
pub mod state;
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());
}
"ready" => {
// inbuxa: ready only while the data store answers
// (a cached, time-limited read); liveness stays 200
return Ok(JsonProblemResponse({
if !self.core.storage.data.is_none() {
if self.is_data_store_ready().await {
StatusCode::OK
} else {
StatusCode::SERVICE_UNAVAILABLE
crates/store/src/backend/mod.rs
+3
View File
@@ -30,6 +30,9 @@ pub mod s3;
pub mod sqlite;
// inbuxa: scale-out storage (sharded stores)
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;
crates/store/src/backend/mysql/blob.rs
+16 -1
View File
@@ -10,7 +10,7 @@ use std::ops::Range;
use mysql_async::prelude::Queryable;
use super::{MysqlStore, into_error};
use super::{MysqlStore, bounded, into_error};
impl MysqlStore {
pub(crate) async fn get_blob(
@@ -19,6 +19,8 @@ impl MysqlStore {
range: Range<usize>,
) -> trc::Result<Option<Vec<u8>>> {
let mut conn = self.conn().await?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prep("SELECT v FROM t WHERE k = ?")
.await
@@ -38,10 +40,15 @@ impl MysqlStore {
}
})
.map_err(into_error)
})
.await;
bounded(conn, result, limit)
}
pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> trc::Result<()> {
let mut conn = self.conn().await?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prep("INSERT INTO t (k, v) VALUES (?, ?) ON DUPLICATE KEY UPDATE v = VALUES(v)")
.await
@@ -50,10 +57,15 @@ impl MysqlStore {
.await
.map_err(into_error)
.map(|_| ())
})
.await;
bounded(conn, result, limit)
}
pub(crate) async fn delete_blob(&self, key: &[u8]) -> trc::Result<bool> {
let mut conn = self.conn().await?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prep("DELETE FROM t WHERE k = ?")
.await
@@ -62,5 +74,8 @@ impl MysqlStore {
.await
.map_err(into_error)
.map(|hits| hits.affected_rows() > 0)
})
.await;
bounded(conn, result, limit)
}
}
crates/store/src/backend/mysql/lookup.rs
+6 -1
View File
@@ -10,7 +10,7 @@ use mysql_async::{Params, Row, prelude::Queryable};
use crate::{IntoRows, QueryResult, QueryType, Value};
use super::{MysqlStore, into_error};
use super::{MysqlStore, bounded, into_error};
impl MysqlStore {
pub(crate) async fn sql_query<T: QueryResult>(
@@ -19,6 +19,8 @@ impl MysqlStore {
params: &[Value<'_>],
) -> trc::Result<T> {
let mut conn = self.conn().await?;
let limit = self.timeouts.query;
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());
@@ -40,6 +42,9 @@ impl MysqlStore {
.await
.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
+13 -3
View File
@@ -6,7 +6,7 @@
* Modified by Coffey Labs in 2026 for INBUXA.
*/
use super::{MysqlStore, into_error};
use super::{MysqlStore, bounded, into_error};
use crate::{
backend::mysql::MysqlSearchField,
search::{
@@ -72,6 +72,7 @@ impl MysqlStore {
.db_name(Some(replica.database.clone()))
.tcp_port(replica.port as u16),
),
timeouts: Default::default(),
})),
replica.host,
replica.port as u16,
@@ -81,6 +82,7 @@ impl MysqlStore {
let primary = Store::MySQL(Arc::new(MysqlStore {
conn_pool: Pool::new(opts),
timeouts: Default::default(),
}));
// ST-1: no replicas, no change
@@ -99,7 +101,8 @@ impl MysqlStore {
pub(crate) async fn create_storage_tables(&self) -> trc::Result<()> {
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,
@@ -169,11 +172,15 @@ impl MysqlStore {
}
Ok(())
})
.await;
bounded(conn, result, limit)
}
pub(crate) async fn create_search_tables(&self) -> trc::Result<()> {
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?;
@@ -181,6 +188,9 @@ impl MysqlStore {
create_search_tables::<TracingSearchField>(&mut conn).await?;
Ok(())
})
.await;
bounded(conn, result, limit)
}
}
crates/store/src/backend/mysql/mod.rs
+41 -1
View File
@@ -6,6 +6,7 @@
* Modified by Coffey Labs in 2026 for INBUXA.
*/
use crate::backend::query_timeout::QueryTimeouts;
use crate::{
search::{
CalendarSearchField, ContactSearchField, EmailSearchField, FileSearchField, SearchField,
@@ -14,7 +15,7 @@ use crate::{
write::SearchIndex,
};
use mysql_async::Pool;
use std::fmt::Display;
use std::{fmt::Display, time::Duration};
pub mod blob;
pub mod lookup;
@@ -25,6 +26,8 @@ pub mod write;
pub struct MysqlStore {
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
@@ -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)]
pub(crate) fn into_error(err: impl Display) -> trc::Error {
trc::StoreEvent::MysqlError.reason(err)
crates/store/src/backend/mysql/read.rs
+56 -13
View File
@@ -6,7 +6,7 @@
* 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 futures::TryStreamExt;
use mysql_async::{Row, prelude::Queryable};
@@ -17,6 +17,8 @@ impl MysqlStore {
U: Deserialize + 'static,
{
let mut conn = self.conn().await?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prep(format!(
"SELECT v FROM {} WHERE k = ?",
@@ -35,10 +37,15 @@ impl MysqlStore {
Ok(None)
}
})
})
.await;
bounded(conn, result, limit)
}
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
let mut conn = self.conn().await?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prep(format!(
"SELECT 1 FROM {} WHERE k = ?",
@@ -51,6 +58,9 @@ impl MysqlStore {
.await
.map_err(into_error)
.map(|r| r.is_some())
})
.await;
bounded(conn, result, limit)
}
pub(crate) async fn iterate<T: Key>(
@@ -64,12 +74,14 @@ impl MysqlStore {
let end = params.end.serialize(0);
let keys = if params.values { "k, v" } else { "k" };
let s = conn
.prep(&match (params.first, params.ascending) {
// inbuxa: a scan may run for hours, so the query limit bounds each
// wait for the database (preparing, the query starting, the next
// row) rather than the scan. A wait that runs out closes the
// connection.
let limit = self.timeouts.query;
let query = match (params.first, params.ascending) {
(true, true) => {
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 ASC LIMIT 1")
}
(true, false) => {
format!(
@@ -82,10 +94,16 @@ impl MysqlStore {
(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 stalled = false;
let mut to = end;
let mut resume_key = None;
@@ -94,13 +112,26 @@ impl MysqlStore {
let mut timed_out = false;
{
let mut rows = conn
.exec_stream::<Row, _, _>(&s, (from.clone(), to.clone()))
let mut rows = match tokio::time::timeout(
limit,
conn.exec_stream::<Row, _, _>(&s, (from.clone(), to.clone())),
)
.await
.map_err(into_error)?;
{
Ok(rows) => rows.map_err(into_error)?,
// Leaves the scan loop for the timeout below
Err(_) => break,
};
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)) => {
let value = if params.values {
row.take_opt::<Vec<u8>, _>(1)
@@ -136,6 +167,10 @@ impl MysqlStore {
}
}
if stalled {
break;
}
match last_key {
Some(last_key) if timed_out => {
if params.ascending {
@@ -148,6 +183,9 @@ impl MysqlStore {
_ => return Ok(()),
}
}
discard(conn);
Err(query_timeout_error(limit))
}
pub(crate) async fn get_counter(
@@ -158,6 +196,8 @@ impl MysqlStore {
let table = char::from(key.subspace());
let key = key.serialize(0);
let mut conn = self.conn().await?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prep(format!("SELECT v FROM {table} WHERE k = ?"))
.await
@@ -167,5 +207,8 @@ impl MysqlStore {
Ok(None) => Ok(0),
Err(e) => Err(into_error(e)),
}
})
.await;
bounded(conn, result, limit)
}
}
crates/store/src/backend/mysql/search.rs
+20 -3
View File
@@ -10,8 +10,8 @@ use crate::{
backend::{
MAX_TOKEN_LENGTH,
mysql::{
DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, MysqlSearchField, MysqlStore, into_error,
is_timeout_error,
DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, MysqlSearchField, MysqlStore, bounded,
into_error, is_timeout_error,
},
},
search::{
@@ -27,6 +27,8 @@ use std::fmt::Write;
impl MysqlStore {
pub async fn index(&self, documents: Vec<IndexDocument>) -> trc::Result<()> {
let mut conn = self.conn().await?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let mut tx_opts = TxOpts::default();
tx_opts
.with_consistent_snapshot(false)
@@ -78,6 +80,9 @@ impl MysqlStore {
}
trx.commit().await.map_err(into_error)
})
.await;
bounded(conn, result, limit)
}
pub async fn query<R: SearchDocumentId>(
@@ -97,12 +102,17 @@ impl MysqlStore {
}
let mut conn = self.conn().await?;
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)
.await
.map(|r| r.into_iter().map(|r| R::from_u64(r as u64)).collect())
.map_err(into_error)
})
.await;
bounded(conn, result, limit)
}
pub async fn unindex(&self, filter: SearchQuery) -> trc::Result<u64> {
@@ -111,6 +121,8 @@ impl MysqlStore {
let params = build_filter(&mut query, &filter.filters);
let mut conn = self.conn().await?;
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 {
@@ -137,7 +149,9 @@ impl MysqlStore {
}
deleted += affected;
}
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_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);
break;
}
@@ -145,6 +159,9 @@ impl MysqlStore {
}
}
}
})
.await;
bounded(conn, result, limit)
}
}
crates/store/src/backend/mysql/write.rs
+18 -2
View File
@@ -6,7 +6,9 @@
* 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::{
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
SUBSPACE_REGISTRY_IDX,
@@ -32,7 +34,8 @@ impl MysqlStore {
let start = Instant::now();
let mut retry_count = 0;
let mut conn = self.conn().await?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
loop {
let err = match self.write_trx(&mut conn, &mut batch).await {
Ok(result) => {
@@ -67,6 +70,9 @@ impl MysqlStore {
tokio::time::sleep(Duration::from_millis(backoff)).await;
retry_count += 1;
}
})
.await;
bounded(conn, result, limit)
}
async fn write_trx(
@@ -385,15 +391,22 @@ impl MysqlStore {
pub(crate) async fn purge_store(&self) -> trc::Result<()> {
let mut conn = self.conn().await?;
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(&mut conn, char::from(subspace)).await?;
}
Ok(())
})
.await;
bounded(conn, result, limit)
}
pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> {
let mut conn = self.conn().await?;
let limit = self.timeouts.maintenance;
let result = tokio::time::timeout(limit, async {
let table = char::from(from.subspace());
let mut from = from.serialize(0);
let to = to.serialize(0);
@@ -450,6 +463,9 @@ impl MysqlStore {
}
}
}
})
.await;
bounded(conn, result, limit)
}
}
crates/store/src/backend/postgres/blob.rs
+18 -1
View File
@@ -2,13 +2,15 @@
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
*
* Modified by Coffey Labs in 2026 for INBUXA.
*/
use std::ops::Range;
use crate::backend::postgres::into_pool_error;
use super::{PostgresStore, into_error};
use super::{PostgresStore, bounded, into_error};
impl PostgresStore {
pub(crate) async fn get_blob(
@@ -17,6 +19,8 @@ impl PostgresStore {
range: Range<usize>,
) -> trc::Result<Option<Vec<u8>>> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prepare_cached("SELECT v FROM t WHERE k = $1")
.await
@@ -39,10 +43,15 @@ impl PostgresStore {
}
})
.map_err(into_error)
})
.await;
bounded(conn, result, limit)
}
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 limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prepare_cached(
"INSERT INTO t (k, v) VALUES ($1, $2) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v",
@@ -53,10 +62,15 @@ impl PostgresStore {
.await
.map_err(into_error)
.map(|_| ())
})
.await;
bounded(conn, result, limit)
}
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 limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prepare_cached("DELETE FROM t WHERE k = $1")
.await
@@ -65,5 +79,8 @@ impl PostgresStore {
.await
.map_err(into_error)
.map(|hits| hits > 0)
})
.await;
bounded(conn, result, limit)
}
}
crates/store/src/backend/postgres/lookup.rs
+8 -1
View File
@@ -2,6 +2,8 @@
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* 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};
@@ -12,7 +14,7 @@ use tokio_postgres::types::{FromSql, ToSql, Type};
use crate::IntoRows;
use super::{PostgresStore, into_error};
use super::{PostgresStore, bounded, into_error};
impl PostgresStore {
pub(crate) async fn sql_query<T: QueryResult>(
@@ -21,6 +23,8 @@ impl PostgresStore {
params_: &[crate::Value<'_>],
) -> trc::Result<T> {
let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn.prepare_cached(query).await.map_err(into_error)?;
let params = params_
.iter()
@@ -48,6 +52,9 @@ impl PostgresStore {
.await
.map_or_else(|e| Err(into_error(e)), |r| Ok(T::from_query_all(r))),
}
})
.await;
bounded(conn, result, limit)
}
}
crates/store/src/backend/postgres/main.rs
+13 -3
View File
@@ -6,7 +6,7 @@
* Modified by Coffey Labs in 2026 for INBUXA.
*/
use super::{PostgresStore, into_error};
use super::{PostgresStore, bounded, into_error};
use crate::{
backend::postgres::{
PsqlSearchField, into_pool_error,
@@ -119,6 +119,7 @@ impl PostgresStore {
Store::PostgreSQL(Arc::new(PostgresStore {
conn_pool: pool,
ts_configs: ts_configs.clone(),
timeouts: Default::default(),
})),
replica.host,
replica.port as u16,
@@ -129,6 +130,7 @@ impl PostgresStore {
let primary = Store::PostgreSQL(Arc::new(PostgresStore {
conn_pool: primary_pool,
ts_configs,
timeouts: Default::default(),
}));
// ST-1: no replicas, no change
@@ -147,7 +149,8 @@ impl PostgresStore {
pub(crate) async fn create_storage_tables(&self) -> trc::Result<()> {
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,
@@ -213,11 +216,15 @@ impl PostgresStore {
}
Ok(())
})
.await;
bounded(conn, result, limit)
}
pub(crate) async fn create_search_tables(&self) -> trc::Result<()> {
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?;
@@ -225,6 +232,9 @@ impl PostgresStore {
create_search_tables::<TracingSearchField>(&conn).await?;
Ok(())
})
.await;
bounded(conn, result, limit)
}
}
crates/store/src/backend/postgres/mod.rs
+33 -1
View File
@@ -6,6 +6,7 @@
* Modified by Coffey Labs in 2026 for INBUXA.
*/
use crate::backend::query_timeout::QueryTimeouts;
use crate::{
search::{
CalendarSearchField, ContactSearchField, EmailSearchField, FileSearchField, SearchField,
@@ -14,7 +15,8 @@ use crate::{
write::SearchIndex,
};
use ahash::AHashSet;
use deadpool_postgres::Pool;
use deadpool_postgres::{Object, Pool};
use std::time::Duration;
use tokio_postgres::error::SqlState;
pub mod blob;
@@ -28,6 +30,8 @@ pub mod write;
pub struct PostgresStore {
pub(crate) conn_pool: Pool,
pub(crate) ts_configs: AHashSet<&'static str>,
/// inbuxa: client-side query limits (see backend::query_timeout)
pub(crate) timeouts: QueryTimeouts,
}
#[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)]
pub(crate) fn into_pool_error(err: deadpool_postgres::PoolError) -> trc::Error {
match err {
crates/store/src/backend/postgres/read.rs
+55 -10
View File
@@ -2,9 +2,11 @@
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* 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::{
Deserialize, IterateParams, Key, ValueKey, backend::postgres::into_pool_error,
write::ValueClass,
@@ -17,6 +19,8 @@ impl PostgresStore {
U: Deserialize + 'static,
{
let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prepare_cached(&format!(
"SELECT v FROM {} WHERE k = $1",
@@ -35,10 +39,15 @@ impl PostgresStore {
Ok(None)
}
})
})
.await;
bounded(conn, result, limit)
}
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 limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prepare_cached(&format!(
"SELECT 1 FROM {} WHERE k = $1",
@@ -51,6 +60,9 @@ impl PostgresStore {
.await
.map_err(into_error)
.map(|r| r.is_some())
})
.await;
bounded(conn, result, limit)
}
pub(crate) async fn iterate<T: Key>(
@@ -64,8 +76,12 @@ impl PostgresStore {
let end = params.end.serialize(0);
let keys = if params.values { "k, v" } else { "k" };
let s = conn
.prepare_cached(&match (params.first, params.ascending) {
// inbuxa: a scan may run for hours, so the query limit bounds each
// wait for the database (preparing, the query starting, the next
// row) rather than the scan. A wait that runs out closes the
// connection.
let limit = self.timeouts.query;
let query = match (params.first, params.ascending) {
(true, true) => {
format!(
"SELECT {keys} FROM {table} WHERE k >= $1 AND k <= $2 ORDER BY k ASC LIMIT 1"
@@ -82,26 +98,43 @@ impl PostgresStore {
(false, false) => {
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 to = end;
let mut resume_key: Option<Vec<u8>> = None;
let mut stalled = false;
loop {
let mut last_key = None;
let mut timed_out = false;
{
let rows = conn
.query_raw(&s, &[&from, &to])
.await
.map_err(into_error)?;
let rows =
match tokio::time::timeout(limit, conn.query_raw(&s, &[&from, &to])).await {
Ok(rows) => rows.map_err(into_error)?,
// Leaves the scan loop for the timeout below
Err(_) => break,
};
pin_mut!(rows);
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)) => {
let key = row.try_get::<_, &[u8]>(0).map_err(into_error)?;
let value = if params.values {
@@ -132,6 +165,10 @@ impl PostgresStore {
}
}
if stalled {
break;
}
match last_key {
Some(last_key) if timed_out => {
if params.ascending {
@@ -144,6 +181,9 @@ impl PostgresStore {
_ => return Ok(()),
}
}
drop(deadpool_postgres::Object::take(conn));
Err(query_timeout_error(limit))
}
pub(crate) async fn get_counter(
@@ -155,6 +195,8 @@ impl PostgresStore {
let key = key.serialize(0);
let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let s = conn
.prepare_cached(&format!("SELECT v FROM {table} WHERE k = $1"))
.await
@@ -164,5 +206,8 @@ impl PostgresStore {
Ok(None) => Ok(0),
Err(e) => Err(into_error(e)),
}
})
.await;
bounded(conn, result, limit)
}
}
crates/store/src/backend/postgres/search.rs
+19 -4
View File
@@ -10,8 +10,8 @@ use crate::{
backend::{
MAX_TOKEN_LENGTH,
postgres::{
DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, PostgresStore, PsqlSearchField, into_error,
into_pool_error, is_timeout_error,
DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, PostgresStore, PsqlSearchField, bounded,
into_error, into_pool_error, is_timeout_error,
},
},
search::{
@@ -36,6 +36,8 @@ impl PostgresStore {
pub async fn index(&self, documents: Vec<IndexDocument>) -> trc::Result<()> {
let mut conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let trx = conn
.build_transaction()
.isolation_level(IsolationLevel::ReadCommitted)
@@ -85,8 +87,8 @@ impl PostgresStore {
if let Some(value) = fields.get(field) {
let value_ref = format!("${}", values.len() + 1);
let (text_len, language) = if let SearchValue::Text { value, language } = value
{
let (text_len, language) =
if let SearchValue::Text { value, language } = value {
(value.len(), self.ts_config(language))
} else {
(0, PG_UNSTEMMED_LANG)
@@ -155,6 +157,9 @@ impl PostgresStore {
}
trx.commit().await.map_err(into_error)
})
.await;
bounded(conn, result, limit)
}
pub async fn query<R: SearchDocumentId>(
@@ -170,6 +175,8 @@ impl PostgresStore {
build_sort(&mut query, sort);
}
let conn = self.conn_pool.get().await.map_err(into_pool_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())
@@ -180,6 +187,9 @@ impl PostgresStore {
.collect::<Result<Vec<R>, _>>()
})
.map_err(into_error)
})
.await;
bounded(conn, result, limit)
}
pub async fn unindex(&self, filter: SearchQuery) -> trc::Result<u64> {
@@ -189,6 +199,8 @@ impl PostgresStore {
let params = self.build_filter(&mut where_clause, &filter.filters);
let params = params.iter().map(SqlParam::as_sql).collect::<Vec<_>>();
let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
let limit = self.timeouts.maintenance;
let result = tokio::time::timeout(limit, async {
let s = conn
.prepare_cached(&format!("DELETE FROM {table}{where_clause}"))
.await
@@ -223,6 +235,9 @@ impl PostgresStore {
}
}
}
})
.await;
bounded(conn, result, limit)
}
fn build_filter<'x>(
crates/store/src/backend/postgres/write.rs
+18 -2
View File
@@ -2,9 +2,11 @@
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
*
* 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::{
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
SUBSPACE_REGISTRY_IDX,
@@ -30,6 +32,8 @@ enum CommitError {
impl PostgresStore {
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 limit = self.timeouts.query;
let result = tokio::time::timeout(limit, async {
let start = Instant::now();
let mut retry_count = 0;
@@ -72,6 +76,9 @@ impl PostgresStore {
}
}
}
})
.await;
bounded(conn, result, limit)
}
async fn write_trx(
@@ -393,16 +400,22 @@ impl PostgresStore {
pub(crate) async fn purge_store(&self) -> trc::Result<()> {
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?;
}
Ok(())
})
.await;
bounded(conn, result, limit)
}
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 limit = self.timeouts.maintenance;
let result = tokio::time::timeout(limit, async {
let table = char::from(from.subspace());
let mut from = from.serialize(0);
let to = to.serialize(0);
@@ -459,6 +472,9 @@ impl PostgresStore {
}
}
}
})
.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
//! good. No database is needed: a local listener that never answers plays
//! 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 std::time::{Duration, Instant};
use store::{Store, ValueKey, write::ValueClass};
use tokio::net::TcpListener;
use std::{
sync::{
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.
async fn silent_server() -> u16 {
@@ -94,3 +114,263 @@ pub async fn mysql_pool_timeout() {
)
.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();
}
}