diff --git a/Justfile b/Justfile index 1bd021a0a6..4fe7d1abbf 100644 --- a/Justfile +++ b/Justfile @@ -445,6 +445,10 @@ test-unit: # #[ignore]d, so --lib runs only the infra-free set. Without this gate a # stray file in migrations/ or a broken lint ships green. cargo nextest run -p buzz-db --lib + # buzz-search unit tests: pure query-builder and cancel-outcome checks + # (no infra). Its Postgres suite runs in the PostgreSQL lane. Keep + # scripts/run-tests.sh in step. + cargo nextest run -p buzz-search --lib # buzz-db source-policy tests: infra-free scans of the crate sources # that enforce the event-write admission chokepoint, metric/provenance # contracts, and their fixtures. They live in an integration-test diff --git a/crates/buzz-db/src/runtime/mod.rs b/crates/buzz-db/src/runtime/mod.rs index 061076cae1..8bfe71786e 100644 --- a/crates/buzz-db/src/runtime/mod.rs +++ b/crates/buzz-db/src/runtime/mod.rs @@ -1099,14 +1099,7 @@ impl Db { .test_before_acquire(false) .acquire_timeout(Self::READER_ACQUIRE_TIMEOUT) .after_connect(|conn, _meta| { - Box::pin(async move { - sqlx::query(crate::event::TIGHTEN_STATEMENT_TIMEOUT_SQL) - .bind(crate::event::READ_STATEMENT_TIMEOUT) - .bind(false) - .execute(&mut *conn) - .await - .map(|_| ()) - }) + Box::pin(crate::event::tighten_session_read_statement_timeout(conn)) }) .after_release(|conn, _meta| Box::pin(Self::reader_release_check(conn))) } diff --git a/crates/buzz-db/src/store/event.rs b/crates/buzz-db/src/store/event.rs index bf80e302d4..55e21b5ec8 100644 --- a/crates/buzz-db/src/store/event.rs +++ b/crates/buzz-db/src/store/event.rs @@ -893,6 +893,21 @@ pub(crate) const TIGHTEN_STATEMENT_TIMEOUT_SQL: &str = "SELECT set_config('state OR current_setting('statement_timeout')::interval > $1::text::interval \ THEN $1::text ELSE current_setting('statement_timeout') END, $2)"; +/// Lower `conn`'s session `statement_timeout` to [`READ_STATEMENT_TIMEOUT`] +/// unless it is already tighter. For pools whose every statement is a read: +/// run it once per new connection (`after_connect`), so no statement on that +/// pool is unbounded. +pub async fn tighten_session_read_statement_timeout( + conn: &mut sqlx::PgConnection, +) -> sqlx::Result<()> { + sqlx::query(TIGHTEN_STATEMENT_TIMEOUT_SQL) + .bind(READ_STATEMENT_TIMEOUT) + .bind(false) + .execute(conn) + .await + .map(|_| ()) +} + /// Run an event read under [`READ_STATEMENT_TIMEOUT`] without changing the /// deadline for anything else on `conn`. `begin()` opens a transaction on an /// autocommit writer connection or a savepoint inside the replica's read diff --git a/crates/buzz-relay/src/api/bridge.rs b/crates/buzz-relay/src/api/bridge.rs index 10c88e3697..2f9544515a 100644 --- a/crates/buzz-relay/src/api/bridge.rs +++ b/crates/buzz-relay/src/api/bridge.rs @@ -23,7 +23,9 @@ use crate::nip_fi_http::{admit_nip_fi_http_on_state, Nip98Proof}; use crate::state::AppState; use super::relay_members::MembershipRefusal; -use super::{api_error, db_read_error, internal_error, not_found, parse_query_or_400}; +use super::{ + api_error, db_read_error, internal_error, not_found, parse_query_or_400, search_read_error, +}; mod thread_roots; mod thread_window; @@ -2582,7 +2584,7 @@ async fn handle_bridge_search( .search .search(&search_query) .await - .map_err(|e| internal_error(&format!("search error: {e}")))?; + .map_err(|e| search_read_error("search error", &e))?; // Fetch full events from DB by ID. Hit ids are already raw 32-byte // arrays from the FTS layer — no hex decode. @@ -9008,6 +9010,9 @@ pub(crate) mod postgres_tests { admin: sqlx::PgPool, writer: sqlx::PgPool, replica: sqlx::PgPool, + /// The production search pool options on the writer, with a 300ms + /// session `statement_timeout` they must keep (it is tighter). + search: sqlx::PgPool, names: [String; 2], state: Arc, host: String, @@ -9082,6 +9087,16 @@ pub(crate) mod postgres_tests { } let replica = pools.pop().expect("replica"); let writer = pools.pop().expect("writer"); + let search_url = format!("{base}/{}", names[0]); + let search = buzz_search::pool_options() + .connect_with( + search_url + .parse::() + .expect("search url") + .options([("statement_timeout", "300ms")]), + ) + .await + .expect("search pool"); let mut config = crate::config::Config::from_env().expect("config"); config.redis_url = @@ -9107,7 +9122,7 @@ pub(crate) mod postgres_tests { buzz_audit::AuditService::new(writer.clone()), pubsub, buzz_auth::AuthService::new(config.auth.clone()), - buzz_search::SearchService::new(writer.clone()), + buzz_search::SearchService::new(search.clone()), Arc::new(buzz_workflow::WorkflowEngine::new( db, buzz_workflow::WorkflowConfig::default(), @@ -9120,6 +9135,7 @@ pub(crate) mod postgres_tests { admin, writer, replica, + search, names: [names[0].clone(), names[1].clone()], state: Arc::new(state), host, @@ -9158,20 +9174,43 @@ pub(crate) mod postgres_tests { } async fn http(&self, uri: &str, filter: &serde_json::Value) -> (StatusCode, Value) { + let body = serde_json::to_vec(&[filter]).expect("json"); + self.post(uri, body, axum::http::HeaderMap::new()).await + } + + /// `POST /profiles/search` for `needle`, NIP-98 signed by the reader. + async fn profile_search(&self) -> (StatusCode, Value) { + let path = super::super::profiles::SEARCH_PATH; + let body = serde_json::to_vec(&serde_json::json!({"query": "needle"})).expect("json"); + let url = super::nip98_expected_url( + &self.state.config.relay_url, + &TenantContext::resolved(self.community, self.host.clone()), + path, + ); + let auth = make_nip98_headers(&self.reader, &url, "POST", &body); + self.post(path, body, auth).await + } + + async fn post( + &self, + uri: &str, + body: Vec, + headers: axum::http::HeaderMap, + ) -> (StatusCode, Value) { use axum::body::Body; use axum::http::{header, Request}; use tower::ServiceExt; + let mut request = Request::builder() + .method("POST") + .uri(uri) + .header(header::HOST, &self.host) + .header("x-pubkey", self.reader.public_key().to_hex()) + .header("content-type", "application/json"); + if let Some(map) = request.headers_mut() { + map.extend(headers); + } let resp = crate::router::build_router(self.state.clone()) - .oneshot( - Request::builder() - .method("POST") - .uri(uri) - .header(header::HOST, &self.host) - .header("x-pubkey", self.reader.public_key().to_hex()) - .header("content-type", "application/json") - .body(Body::from(serde_json::to_vec(&[filter]).expect("json"))) - .expect("request"), - ) + .oneshot(request.body(Body::from(body)).expect("request")) .await .expect("oneshot"); let status = resp.status(); @@ -9257,6 +9296,7 @@ pub(crate) mod postgres_tests { async fn drop(self) { drop(self.state); + self.search.close().await; for (pool, name) in [ (self.writer, &self.names[0]), (self.replica, &self.names[1]), @@ -9358,6 +9398,90 @@ pub(crate) mod postgres_tests { fx.drop().await; } + /// The search itself (not its hydrate) cancelled by the search pool's + /// statement timeout answers with the timeout contract on every search + /// caller: HTTP `/query` and `/profiles/search` 503 `query timed out`, WS + /// CLOSED `error: query timed out`. The ordinary-error control (42P01) + /// keeps the generic 500 and the retryable `error: search failed`. + /// Mutations: drop the 57014 check in any caller, or map every search + /// error to the timeout → RED. + #[tokio::test] + #[ignore = "requires Postgres"] + async fn search_statement_timeout_maps_to_timeout_contract() { + use crate::handlers::req::{QUERY_TIMED_OUT_CLOSED, SEARCH_FAILED_CLOSED}; + let fx = cancel_fixture().await; + let closed = |reason: &str| format!(r#"["CLOSED","s","{reason}"]"#); + + // Healthy controls. + let (status, body) = fx.http("/query", &fx.search_filter()).await; + assert_eq!(status, StatusCode::OK, "healthy HTTP search: {body}"); + assert!( + body.to_string().contains("root needle"), + "search hit: {body}" + ); + let (status, body) = fx.profile_search().await; + assert_eq!(status, StatusCode::OK, "healthy profile search: {body}"); + let frame = fx.ws_search().await; + assert!(frame.contains("root needle"), "healthy WS search: {frame}"); + + // `events` locked on the search pool's database: the search waits past + // its 300ms `statement_timeout` and is cancelled with 57014. + let mut locker = fx.writer.begin().await.expect("begin locker"); + sqlx::query("LOCK TABLE events IN ACCESS EXCLUSIVE MODE") + .execute(&mut *locker) + .await + .expect("lock search events"); + let timed_out = ( + StatusCode::SERVICE_UNAVAILABLE, + Some(super::super::QUERY_TIMED_OUT), + ); + let (status, body) = fx.http("/query", &fx.search_filter()).await; + assert_eq!((status, body["error"].as_str()), timed_out, "HTTP search"); + let (status, body) = fx.profile_search().await; + assert_eq!( + (status, body["error"].as_str()), + timed_out, + "HTTP profile search" + ); + assert_eq!( + fx.ws_search().await, + closed(QUERY_TIMED_OUT_CLOSED), + "WS search" + ); + locker.rollback().await.expect("unlock"); + + // Ordinary-error controls: `events` gone → 42P01. Not a timeout. + for pool in [&fx.writer, &fx.replica] { + sqlx::query("ALTER TABLE events RENAME TO events_gone") + .execute(pool) + .await + .expect("rename events"); + } + let failed = ( + StatusCode::INTERNAL_SERVER_ERROR, + Some("internal server error"), + ); + let (status, body) = fx.http("/query", &fx.search_filter()).await; + assert_eq!( + (status, body["error"].as_str()), + failed, + "HTTP search ordinary error" + ); + let (status, body) = fx.profile_search().await; + assert_eq!( + (status, body["error"].as_str()), + failed, + "HTTP profile search ordinary error" + ); + assert_eq!( + fx.ws_search().await, + closed(SEARCH_FAILED_CLOSED), + "WS search ordinary error" + ); + + fx.drop().await; + } + /// A search REQ reusing a live subscription's ID retires it (NIP-01 /// replacement) before its hydrate is cancelled, so the timeout CLOSED /// leaves nothing registered under the ID. diff --git a/crates/buzz-relay/src/api/mod.rs b/crates/buzz-relay/src/api/mod.rs index 129d7dd732..d9059d03ae 100644 --- a/crates/buzz-relay/src/api/mod.rs +++ b/crates/buzz-relay/src/api/mod.rs @@ -41,7 +41,24 @@ pub(crate) fn db_read_error( context: &str, e: &buzz_db::DbError, ) -> (StatusCode, Json) { - if e.is_statement_cancelled() { + read_error(context, e.is_statement_cancelled(), e) +} + +/// [`db_read_error`] for a search: a search cancelled by the search pool's +/// statement timeout gets the same 503 `query timed out`. +pub(crate) fn search_read_error( + context: &str, + e: &buzz_search::SearchError, +) -> (StatusCode, Json) { + read_error(context, e.is_statement_cancelled(), e) +} + +fn read_error( + context: &str, + statement_cancelled: bool, + e: &dyn std::fmt::Display, +) -> (StatusCode, Json) { + if statement_cancelled { tracing::warn!("{context}: {e}"); return api_error(StatusCode::SERVICE_UNAVAILABLE, QUERY_TIMED_OUT); } @@ -769,6 +786,16 @@ mod db_read_error_tests { assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR); assert_eq!(body.0["error"], "internal server error"); } + + #[test] + fn non_cancel_search_error_stays_generic_500() { + let (status, body) = search_read_error( + "ctx", + &buzz_search::SearchError::Db(sqlx::Error::PoolTimedOut), + ); + assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(body.0["error"], "internal server error"); + } } #[cfg(test)] diff --git a/crates/buzz-relay/src/api/profiles.rs b/crates/buzz-relay/src/api/profiles.rs index 66885cf283..e099d896b3 100644 --- a/crates/buzz-relay/src/api/profiles.rs +++ b/crates/buzz-relay/src/api/profiles.rs @@ -19,7 +19,7 @@ use serde_json::Value; use crate::state::AppState; -use super::{api_error, bridge, db_read_error, internal_error}; +use super::{api_error, bridge, db_read_error, search_read_error}; /// Route path for profile name search. pub(crate) const SEARCH_PATH: &str = "/profiles/search"; @@ -98,7 +98,7 @@ async fn search_inner( .search .search_profile_names(tenant.community(), &query, page, limit) .await - .map_err(|e| internal_error(&format!("profile search error: {e}")).into_response())? + .map_err(|e| search_read_error("profile search error", &e).into_response())? .hits; if hits.is_empty() { return Ok(Json(Value::Array(Vec::new()))); diff --git a/crates/buzz-relay/src/handlers/req.rs b/crates/buzz-relay/src/handlers/req.rs index bddaa55e15..c0ff3b5356 100644 --- a/crates/buzz-relay/src/handlers/req.rs +++ b/crates/buzz-relay/src/handlers/req.rs @@ -894,7 +894,7 @@ async fn handle_search_req( super::close::close_if_owner( sub_id, owner, - Some(SEARCH_FAILED_CLOSED), + Some(search_closed_reason(&e)), conn, state, ) @@ -1665,6 +1665,17 @@ fn history_read_closed_reason(e: &buzz_db::DbError) -> &'static str { /// (e.g. the read replica it queries is down). pub(crate) const SEARCH_FAILED_CLOSED: &str = "error: search failed"; +/// CLOSED reason for a failed NIP-50 search: a search cancelled by the search +/// pool's statement timeout gets the stable timeout reason, like any other +/// cancelled read; anything else is the retryable [`SEARCH_FAILED_CLOSED`]. +fn search_closed_reason(e: &buzz_search::SearchError) -> &'static str { + if e.is_statement_cancelled() { + QUERY_TIMED_OUT_CLOSED + } else { + SEARCH_FAILED_CLOSED + } +} + /// CLOSED reason for a failed one-shot DB read (COUNT): a statement cancel gets /// the stable timeout reason; anything else keeps the raw error. pub(crate) fn db_read_closed_reason(e: &buzz_db::DbError) -> String { diff --git a/crates/buzz-relay/src/main.rs b/crates/buzz-relay/src/main.rs index f48926f72d..59803a0aa8 100644 --- a/crates/buzz-relay/src/main.rs +++ b/crates/buzz-relay/src/main.rs @@ -61,7 +61,7 @@ fn connect_search_pool(config: &DbConfig) -> anyhow::Result { .read_database_url .as_deref() .unwrap_or(&config.database_url); - sqlx::postgres::PgPoolOptions::new() + buzz_search::pool_options() .acquire_timeout(std::time::Duration::from_secs(config.acquire_timeout_secs)) .connect_lazy(url) .map_err(|e| anyhow::anyhow!("Search DB URL is invalid: {e}")) diff --git a/crates/buzz-search/Cargo.toml b/crates/buzz-search/Cargo.toml index 2d491e7056..73a324791a 100644 --- a/crates/buzz-search/Cargo.toml +++ b/crates/buzz-search/Cargo.toml @@ -9,13 +9,11 @@ description = "Postgres full-text search for Buzz, scoped by community" [dependencies] buzz-core = { workspace = true } +buzz-db = { workspace = true } buzz-datastore-tracing = { workspace = true } sqlx = { workspace = true } uuid = { workspace = true } thiserror = { workspace = true } tracing = { workspace = true } metrics = { workspace = true } - -[dev-dependencies] -buzz-db = { workspace = true } tokio = { workspace = true } diff --git a/crates/buzz-search/src/cancel.rs b/crates/buzz-search/src/cancel.rs new file mode 100644 index 0000000000..1ca4bf2090 --- /dev/null +++ b/crates/buzz-search/src/cancel.rs @@ -0,0 +1,163 @@ +//! Cancel a search statement server-side when its caller goes away. +//! +//! Dropping a sqlx query future does not stop its statement: Postgres runs it +//! until it next writes to the socket, and a ranked `LIMIT` query writes +//! nothing until it has matched and ranked every candidate row. Typeahead +//! abandons a search on almost every keystroke (the client aborts its +//! request and the relay drops the handler), so without this each abandoned +//! search would still run to the end. +//! +//! sqlx does not expose a connection's cancel key, so [`fetch_all`] reads the +//! backend's `(pid, backend_start)` before the statement. If the caller drops +//! the future first, a spawned task signals that backend with +//! `pg_cancel_backend` from another pool connection. The pair names one +//! backend for its lifetime: a reader endpoint can route the cancelling +//! connection to a different replica, where the pair matches no backend, so a +//! cancel never reaches an unrelated statement; the task tries another +//! connection instead. The abandoned connection stays checked out until the +//! task is done, so its backend cannot start another caller's statement +//! before the signal lands. A cancel that cannot reach the backend leaves the +//! statement to the pool's statement timeout ([`crate::pool_options`]). + +use std::time::Duration; + +use sqlx::pool::PoolConnection; +use sqlx::postgres::{PgArguments, PgRow}; +use sqlx::query::Query; +use sqlx::{PgPool, Postgres}; + +/// Pool connections a cancel tries before it gives up on reaching the +/// backend. Each miss is a connection to another replica. +const CANCEL_ATTEMPTS: usize = 3; + +/// Bound on each cancel attempt's pool checkout; a saturated pool gives up +/// rather than queueing behind the searches it is trying to stop. +const CANCEL_ACQUIRE_TIMEOUT: Duration = Duration::from_millis(500); + +/// `backend_start` as integer microseconds, so it round-trips exactly. +macro_rules! backend_start_us { + () => { + "(EXTRACT(EPOCH FROM backend_start) * 1000000)::bigint" + }; +} + +/// This connection's backend: `(pid, backend_start in microseconds)`. +const BACKEND_SQL: &str = concat!( + "SELECT pid, ", + backend_start_us!(), + " FROM pg_stat_activity WHERE pid = pg_backend_pid()" +); + +/// Signal the backend `($1 pid, $2 backend_start in microseconds)`: one row +/// when it is on this server, none otherwise. Public only for its tests. +#[doc(hidden)] +pub const CANCEL_SQL: &str = concat!( + "SELECT pg_cancel_backend(pid) FROM pg_stat_activity WHERE pid = $1 AND ", + backend_start_us!(), + " = $2" +); + +/// Run `query` on a connection from `pool`, cancelling its statement on the +/// server if the returned future is dropped before it completes. +pub(crate) async fn fetch_all( + pool: &PgPool, + query: Query<'_, Postgres, PgArguments>, +) -> Result, sqlx::Error> { + let mut conn = pool.acquire().await?; + let backend: (i32, i64) = sqlx::query_as(BACKEND_SQL).fetch_one(&mut *conn).await?; + let mut armed = CancelOnDrop { + pool: pool.clone(), + backend, + conn: Some(conn), + }; + // `conn` is `Some` until this function or `Drop` takes it. + let Some(conn) = armed.conn.as_deref_mut() else { + return Err(sqlx::Error::PoolClosed); + }; + let rows = query.fetch_all(conn).await; + // Finished (rows or error): nothing left to cancel, and the connection + // goes back to the pool. + armed.conn = None; + rows +} + +/// The connection of a statement still in flight, and how to signal its +/// backend. Dropping it with the connection still held cancels the statement. +struct CancelOnDrop { + pool: PgPool, + backend: (i32, i64), + conn: Option>, +} + +impl Drop for CancelOnDrop { + fn drop(&mut self) { + let Some(conn) = self.conn.take() else { + return; + }; + // No runtime (process shutdown): the session ends with the process. + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + let pool = self.pool.clone(); + let backend = self.backend; + runtime.spawn(async move { + let outcome = cancel(&pool, backend).await; + metrics::counter!("buzz_search_cancel_total", "outcome" => outcome).increment(1); + // Only now may the backend take another statement. + drop(conn); + }); + } +} + +/// Signal `backend` from another connection of `pool`; returns the outcome +/// label for `buzz_search_cancel_total`. +async fn cancel(pool: &PgPool, (pid, started_us): (i32, i64)) -> &'static str { + // Connections that reached another replica stay checked out, so each + // attempt gets a different connection. + let mut elsewhere = Vec::with_capacity(CANCEL_ATTEMPTS); + for _ in 0..CANCEL_ATTEMPTS { + let mut conn = match tokio::time::timeout(CANCEL_ACQUIRE_TIMEOUT, pool.acquire()).await { + Ok(Ok(conn)) => conn, + Ok(Err(e)) => { + tracing::warn!(error = %e, pid, "search cancel could not get a connection"); + return "error"; + } + Err(_) => return "pool_busy", + }; + match sqlx::query_scalar::<_, bool>(CANCEL_SQL) + .bind(pid) + .bind(started_us) + .fetch_optional(&mut *conn) + .await + { + Ok(None) => elsewhere.push(conn), + Ok(Some(signalled)) => return signal_outcome(signalled), + Err(e) => { + tracing::warn!(error = %e, pid, "search cancel failed"); + return "error"; + } + } + } + "missed" +} + +/// Outcome label for a [`CANCEL_SQL`] row: `pg_cancel_backend` returns false +/// when the signal could not be sent (the backend exited in between). +fn signal_outcome(signalled: bool) -> &'static str { + if signalled { + "cancelled" + } else { + "signal_failed" + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn only_a_sent_signal_counts_as_cancelled() { + assert_eq!(signal_outcome(true), "cancelled"); + assert_eq!(signal_outcome(false), "signal_failed"); + } +} diff --git a/crates/buzz-search/src/error.rs b/crates/buzz-search/src/error.rs index 4a71a01daa..f608b03b9b 100644 --- a/crates/buzz-search/src/error.rs +++ b/crates/buzz-search/src/error.rs @@ -7,3 +7,18 @@ pub enum SearchError { #[error("database error: {0}")] Db(#[from] sqlx::Error), } + +impl SearchError { + /// Whether Postgres cancelled the search statement (SQLSTATE 57014): the + /// search pool's statement timeout ([`crate::pool_options`]) ended it. + /// Callers answer it with the same terminal deadline response as any + /// other cancelled read, so clients do not re-run the same expensive + /// search. A cancel the search sent itself never reaches a caller: it is + /// sent only after the caller dropped the search. + pub fn is_statement_cancelled(&self) -> bool { + matches!( + self, + SearchError::Db(sqlx::Error::Database(db)) if db.code().as_deref() == Some("57014") + ) + } +} diff --git a/crates/buzz-search/src/lib.rs b/crates/buzz-search/src/lib.rs index 4af282dacc..a7264b6ed7 100644 --- a/crates/buzz-search/src/lib.rs +++ b/crates/buzz-search/src/lib.rs @@ -21,6 +21,8 @@ //! `community_id = $ctx` as the first predicate. A query bound to community //! A cannot return events stored under community B, by construction. +#[doc(hidden)] +pub mod cancel; /// Search error types. pub mod error; /// Search query execution. @@ -33,8 +35,20 @@ pub use query::{ SearchSort, PROFILE_NAMES_PER_PAGE_MAX, }; +use sqlx::postgres::PgPoolOptions; use sqlx::PgPool; +/// Options for the pool behind [`SearchService`]: every new connection's +/// session `statement_timeout` is lowered to buzz-db's event-read deadline +/// ([`buzz_db::event::READ_STATEMENT_TIMEOUT`]) unless it is already +/// tighter. Searches cancel their statement when abandoned (see +/// [`cancel`]); this bounds the ones a cancel cannot reach. +pub fn pool_options() -> PgPoolOptions { + PgPoolOptions::new().after_connect(|conn, _meta| { + Box::pin(buzz_db::event::tighten_session_read_statement_timeout(conn)) + }) +} + /// Thin handle around a `PgPool` for community-scoped FTS. /// /// Holds nothing the pool itself doesn't already own. The whole purpose of diff --git a/crates/buzz-search/src/query.rs b/crates/buzz-search/src/query.rs index ec2ce027ea..6f55d2300d 100644 --- a/crates/buzz-search/src/query.rs +++ b/crates/buzz-search/src/query.rs @@ -331,7 +331,7 @@ pub async fn search(pool: &PgPool, query: &SearchQuery) -> Result = row.try_get("id")?; diff --git a/crates/buzz-search/tests/postgres_fts_integration.rs b/crates/buzz-search/tests/postgres_fts_integration.rs index 4fe4b69b1b..b4c92c6e6f 100644 --- a/crates/buzz-search/tests/postgres_fts_integration.rs +++ b/crates/buzz-search/tests/postgres_fts_integration.rs @@ -1775,3 +1775,203 @@ async fn profile_names_survive_unparsable_content() { teardown(pool, &schema).await; } + +/// Connect the way the relay does: [`buzz_search::pool_options`]. +async fn search_pool() -> PgPool { + let url = std::env::var("BUZZ_TEST_DATABASE_URL").unwrap_or_else(|_| TEST_DB_URL.to_string()); + buzz_search::pool_options() + .max_connections(4) + .connect(&url) + .await + .expect("connect search pool") +} + +/// Backends in this test's database, other than the caller's, that wait on +/// a lock inside a search. Scoped to the database so concurrent tests in +/// sibling databases cannot change the count. +async fn waiting_searches(pool: &PgPool) -> i64 { + sqlx::query_scalar( + "SELECT count(*) FROM pg_stat_activity WHERE pid <> pg_backend_pid() \ + AND datname = current_database() \ + AND wait_event_type = 'Lock' AND query LIKE '%search_query.query%'", + ) + .fetch_one(pool) + .await + .expect("count waiting searches") +} + +/// Poll until [`waiting_searches`] is `want`, failing after five seconds. +async fn await_waiting_searches(pool: &PgPool, want: i64, why: &str) { + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + let got = waiting_searches(pool).await; + if got == want { + return; + } + assert!( + tokio::time::Instant::now() < deadline, + "{why}: {got} waiting searches" + ); + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } +} + +/// Hold `start()`'s search on a table lock, drop it while it waits, and +/// require that its statement stops; then run it again unlocked. +async fn abandon_under_lock(pool: &PgPool, name: &str, start: F) +where + F: Fn() -> Fut, + Fut: std::future::Future>, +{ + let mut lock = pool.begin().await.expect("begin lock"); + sqlx::query("LOCK TABLE events IN ACCESS EXCLUSIVE MODE") + .execute(&mut *lock) + .await + .expect("lock events"); + // Boxed so `drop` drops the search itself, not a pinned reference to it. + let mut running = Box::pin(start()); + tokio::select! { + result = &mut running => panic!("{name} finished under the lock: {result:?}"), + () = await_waiting_searches(pool, 1, "search never reached the lock") => {} + } + drop(running); + await_waiting_searches(pool, 0, &format!("abandoned {name} still running")).await; + lock.rollback().await.expect("release lock"); + // The cancelled connection is usable again. + start() + .await + .unwrap_or_else(|e| panic!("{name} after cancel: {e}")); +} + +#[tokio::test] +#[ignore = "requires Postgres"] +async fn search_pool_bounds_every_statement() { + let (pool, schema) = setup().await; + let search = search_pool().await; + let mut conns = Vec::new(); + for _ in 0..2 { + let mut conn = search.acquire().await.expect("search connection"); + let timeout: String = sqlx::query_scalar("SHOW statement_timeout") + .fetch_one(&mut *conn) + .await + .expect("read statement_timeout"); + assert_eq!(timeout, "20s", "every search connection is bounded"); + conns.push(conn); + } + drop(conns); + search.close().await; + teardown(pool, &schema).await; +} + +#[tokio::test] +#[ignore = "requires Postgres"] +async fn abandoned_search_cancels_its_statement() { + let (pool, schema) = setup().await; + let c = mk_community(&pool, "abandoned-search.example").await; + insert_event( + &pool, + c, + rand_bytes32(), + rand_bytes32(), + 9, + "project plan", + None, + 1_700_000_000, + ) + .await; + let svc = SearchService::new(search_pool().await); + let message_query = SearchQuery { + community: c, + q: "project".into(), + channel_scope: ChannelScope::Any, + kinds: Some(vec![9]), + authors: None, + since: None, + until: None, + page: 1, + per_page: 10, + mode: buzz_search::SearchMode::Prefix, + sort: buzz_search::SearchSort::Relevance, + }; + abandon_under_lock(&pool, "search", || svc.search(&message_query)).await; + abandon_under_lock(&pool, "search_profile_names", || { + svc.search_profile_names(c, "proj", 1, 10) + }) + .await; + teardown(pool, &schema).await; +} + +#[tokio::test] +#[ignore = "requires Postgres"] +async fn cancel_signals_only_the_named_backend() { + let (pool, schema) = setup().await; + let mut victim = pool.acquire().await.expect("victim connection"); + let (pid, started_us): (i32, i64) = sqlx::query_as( + "SELECT pid, (EXTRACT(EPOCH FROM backend_start) * 1000000)::bigint \ + FROM pg_stat_activity WHERE pid = pg_backend_pid()", + ) + .fetch_one(&mut *victim) + .await + .expect("victim backend"); + let sleeping = tokio::spawn(async move { + sqlx::query("SELECT pg_sleep(30)") + .execute(&mut *victim) + .await + .map(|_| ()) + }); + // Signal only once the victim's statement is running: a cancel that + // lands first is a no-op, and the named cancel below would then pass + // without stopping anything. + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + let running: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity \ + WHERE pid = $1 AND state = 'active' AND query LIKE '%pg_sleep(30)%')", + ) + .bind(pid) + .fetch_one(&pool) + .await + .expect("victim activity"); + if running { + break; + } + assert!( + tokio::time::Instant::now() < deadline, + "victim statement never started" + ); + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + + // A same-pid backend that started at another time (another replica, or a + // reused pid) is not signalled. + let other: Option = sqlx::query_scalar(buzz_search::cancel::CANCEL_SQL) + .bind(pid) + .bind(started_us + 1) + .fetch_optional(&pool) + .await + .expect("cancel other backend"); + assert_eq!(other, None); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + assert!( + !sleeping.is_finished(), + "an unmatched cancel must not stop the statement" + ); + + let named: Option = sqlx::query_scalar(buzz_search::cancel::CANCEL_SQL) + .bind(pid) + .bind(started_us) + .fetch_optional(&pool) + .await + .expect("cancel named backend"); + assert_eq!(named, Some(true)); + let err = tokio::time::timeout(std::time::Duration::from_secs(5), sleeping) + .await + .expect("cancelled statement ends") + .expect("join") + .expect_err("statement was cancelled"); + assert_eq!( + err.as_database_error().and_then(|e| e.code()).as_deref(), + Some("57014") + ); + teardown(pool, &schema).await; +} diff --git a/scripts/run-tests.sh b/scripts/run-tests.sh index ec10d1d1a5..83c2591782 100755 --- a/scripts/run-tests.sh +++ b/scripts/run-tests.sh @@ -140,6 +140,11 @@ run_unit_tests() { run_test_step "buzz-db source-policy tests" \ cargo test -p buzz-db --test observability_source -- --nocapture + # buzz-search unit tests (no infra); its Postgres suite runs in the + # PostgreSQL lane. + run_test_step "buzz-search unit tests" \ + cargo test -p buzz-search --lib -- --nocapture + run_test_step "buzz-media storage snapshot serialization test" \ cargo test -p buzz-media --lib bucket_index::tests::bucket_snapshot_json_round_trip_preserves_community_keys -- --exact --nocapture