SQL queries time out; readiness follows the data store

Cluster rehearsal 3: with PostgreSQL paused (docker pause, so its
kernel still answered TCP keepalives), requests on connections already
checked out hung until it came back, and /healthz/ready stayed 200
through the outage. #41 bounded getting a connection, not using one.

Client-side query limits (store::backend::query_timeout). Every
operation on a PostgreSQL or MySQL connection now runs under a time
limit. A server-side statement_timeout (or MySQL's MAX_EXECUTION_TIME,
which covers SELECTs only) can't do this: the server that would enforce
it is the one not answering. When an operation runs out, its connection
is closed instead of pooled, since a query may still be in flight on it
or a transaction open: deadpool's Object::take on PostgreSQL;
Conn::disconnect on MySQL, which marks the connection closed before it
sends anything, so the pool discards it even when the server never
answers.

- query, 2 minutes: reads, writes (the whole transaction with its
  retries), blobs, SQL lookups, search queries and indexing. These take
  milliseconds; two minutes leaves room for a large blob over a slow
  link and still ends a hang.
- maintenance, 30 minutes: range deletes (account removal, purges),
  unindexing, purge_store, and creating tables and indexes at startup,
  which can legitimately run long in one statement. Their existing
  chunked fallback for server-side statement timeouts is unchanged.
- iterate (exports, reindexing, maintenance scans) can run for hours,
  so the query limit bounds each wait for the database (preparing, the
  query starting, the next row) rather than the whole scan.

The limits are fixed, like the pool timeouts; the DataStore schema has
no field for them. Tests set them with Store::with_query_timeouts
(test_mode only).

Readiness. /healthz/ready answered 200 whenever a data store was
configured. It now reads one key from the data store with a 2 s limit
and reuses the answer for 2 s, so probes can't load the database;
while one probe runs, others get the last answer. The first failed
probe of an outage is logged. /healthz/live stays 200: restarting a
node doesn't bring its database back, and an orchestrator restarting on
failed liveness would restart every node at once. The container
HEALTHCHECK already uses /healthz/live.

Tests, store::pool_timeout (a proxy that stops forwarding while
keeping connections open plays the paused database):
- postgres_query_timeout, mysql_query_timeout (new): with four pooled
  connections open, a read, a scan and a write each fail with "Query
  timed out" 2.0 s after the pause (2 s test limit); once the proxy
  forwards again the store answers. With the limits set to an hour
  (upstream's behavior), the read was still waiting at the test's 20 s
  limit.
- postgres_readiness (new, STORE=PostgreSql): a node's data store
  goes through the proxy; /healthz/ready is 200, 503 about 4 s after
  the pause while /healthz/live stays 200, and 200 again about 2 s
  after it ends.
- postgres_pool_timeout, mysql_pool_timeout: pass as before.
store::store_tests (PostgreSql, MySql, including the MariaDB statement
timeout step) and store::task_locks (PostgreSql) pass;
store::search_tests (PostgreSql) fails at the same ordering assertion
(query.rs:684) as on main.

tests/src/store/pool_timeout.rs

@@ -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();
+    }
+}