diff --git a/crates/adapters/src/controller.rs b/crates/adapters/src/controller.rs index ea8cea43b19..04f81a6f0b0 100644 --- a/crates/adapters/src/controller.rs +++ b/crates/adapters/src/controller.rs @@ -179,6 +179,10 @@ use feldera_types::constants::{STATE_FILE, STEPS_FILE}; use feldera_types::format::json::{JsonFlavor, JsonParserConfig, JsonUpdateFormat}; pub use feldera_types::pipeline_diff::compute_pipeline_diff; use feldera_types::program_schema::{SqlIdentifier, canonical_identifier}; +// Connector tests assert that a connector's own error messages stay under the +// bound the endpoint status enforces. +#[cfg(test)] +pub(crate) use stats::MAX_CONNECTOR_ERROR_LEN; pub use stats::{CompletionToken, ControllerStatus, ControllerStatusContext, InputEndpointStatus}; /// Maximal number of concurrent API connections per circuit diff --git a/crates/adapters/src/controller/stats.rs b/crates/adapters/src/controller/stats.rs index 8f28f958e29..81398e6674d 100644 --- a/crates/adapters/src/controller/stats.rs +++ b/crates/adapters/src/controller/stats.rs @@ -38,6 +38,7 @@ use crate::{ checkpoint::{CheckpointInputEndpointMetrics, CheckpointOutputEndpointMetrics}, journal::{InputChecksums, InputLog}, }, + util::truncate_ellipse_middle, }; use anyhow::Error as AnyError; use anyhow::anyhow; @@ -91,6 +92,19 @@ use utoipa::ToSchema; /// There are separate counters for parse and transport errors and for every tag. pub(crate) const MAX_CONNECTOR_ERRORS: usize = 100; +/// Maximum length of an error message stored as part of the endpoint status. +/// +/// An endpoint reports up to `MAX_CONNECTOR_ERRORS` messages per tag, so without +/// a per-message bound a connector that quotes its payload in an error message +/// can grow `/stats` and `/output_endpoints/{name}/stats` past the size limit +/// the pipeline manager applies while proxying the response, which turns the +/// whole request into an error. A connector is expected to bound the data it +/// quotes; this cap is the backstop for when it does not. +/// +/// The bound is generous enough that a message describing a single record, its +/// error chain and a backtrace survives intact. +pub(crate) const MAX_CONNECTOR_ERROR_LEN: usize = 8 * 1024; + /// Kind of input buffered by an endpoint. /// /// Used by [ControllerStatus::input_batch_global] to decide whether input @@ -2330,7 +2344,7 @@ impl InputEndpointStatus { if fatal { let mut fatal_error = self.fatal_error.lock().unwrap(); if fatal_error.is_none() { - *fatal_error = Some(error.to_string()); + *fatal_error = Some(bound_error_message(&error.to_string())); } } } @@ -2634,6 +2648,16 @@ pub struct ProcessedRecords { pub total_processed_steps: Step, } +/// Caps an error message at [`MAX_CONNECTOR_ERROR_LEN`] before it is stored in +/// the endpoint status. +/// +/// The middle is what gets dropped: a message built from an `anyhow` error chain +/// leads with the outermost context and ends with the root cause, so keeping +/// both ends preserves the diagnosis. +pub(crate) fn bound_error_message(message: &str) -> String { + truncate_ellipse_middle(message, MAX_CONNECTOR_ERROR_LEN).into_owned() +} + /// Recent connector errors. /// /// Stores up to MAX_CONNECTOR_ERRORS most recent errors for each tag. @@ -2664,7 +2688,7 @@ impl ConnectorErrorList { timestamp: Utc::now(), tag: tag.map(|tag| tag.to_string()), index, - message: error.to_string(), + message: bound_error_message(&error.to_string()), }); if entry.len() > MAX_CONNECTOR_ERRORS { entry.pop_front(); @@ -2692,16 +2716,19 @@ impl ConnectorErrorList { /// `OutputEndpointStatus::encode_error` / `transport_error`). Do not /// renumber on restore. /// - /// The per-tag [`MAX_CONNECTOR_ERRORS`] bound is re-enforced so that - /// a malformed checkpoint cannot make a list grow unbounded. If the - /// input exceeds the bound for any tag, the excess oldest entries - /// are dropped and a warning is logged. + /// The per-tag [`MAX_CONNECTOR_ERRORS`] bound and the + /// [`MAX_CONNECTOR_ERROR_LEN`] message bound are both re-enforced so + /// that a malformed checkpoint, or one written by a build that lacked + /// either bound, cannot make a list grow unbounded. If the input + /// exceeds the count bound for any tag, the excess oldest entries are + /// dropped and a warning is logged. pub fn from_api_type(errors: Vec) -> Self { let mut list = Self::new(); let mut dropped: usize = 0; - for error in errors { + for mut error in errors { let entry: &mut VecDeque = list.errors.entry(error.tag.clone()).or_default(); + error.message = bound_error_message(&error.message); entry.push_back(error); if entry.len() > MAX_CONNECTOR_ERRORS { entry.pop_front(); @@ -3020,7 +3047,7 @@ impl OutputEndpointStatus { if fatal { let mut fatal_error = self.fatal_error.lock().unwrap(); if fatal_error.is_none() { - *fatal_error = Some(error.to_string()); + *fatal_error = Some(bound_error_message(&error.to_string())); } } } @@ -3042,7 +3069,12 @@ impl OutputEndpointStatus { #[cfg(test)] mod test { - use super::InputEndpointMetrics; + use super::{ + ConnectorError, ConnectorErrorList, InputEndpointMetrics, MAX_CONNECTOR_ERROR_LEN, + bound_error_message, + }; + use anyhow::anyhow; + use chrono::Utc; #[test] fn latency_p99_absent_without_samples() { @@ -3120,4 +3152,67 @@ mod test { .record(750_000u64); assert_eq!(metrics.to_api_type().processing_latency_p99_micros, None); } + + /// An error message the size of a Postgres output batch, shaped like the + /// `anyhow` chain the connector produces: context first, root cause last. + fn oversized_error() -> String { + format!( + "while executing insert statement: [{}]\n\nCaused by:\n \ + null value in column \"freshness_timestamp\" violates not-null constraint", + "{\"a\":1},".repeat(150_000) + ) + } + + #[test] + fn add_error_bounds_the_message() { + let error = oversized_error(); + assert!(error.len() > 1024 * 1024, "test needs an oversized message"); + + let mut list = ConnectorErrorList::new(); + list.add_error(Some("pg_exec"), anyhow!("{error}"), 1); + + let stored = list.to_api_type(); + assert_eq!(stored.len(), 1); + assert!( + stored[0].message.len() <= MAX_CONNECTOR_ERROR_LEN + 64, + "stored message is {} bytes", + stored[0].message.len() + ); + } + + /// Truncation must keep the root cause, which sits at the end of the chain. + #[test] + fn bounding_keeps_both_ends_of_the_message() { + let bounded = bound_error_message(&oversized_error()); + + assert!(bounded.starts_with("while executing insert statement:")); + assert!(bounded.ends_with("violates not-null constraint")); + assert!(bounded.contains("bytes elided")); + } + + #[test] + fn bounding_leaves_a_normal_message_intact() { + let error = "postgres error: permanent: SqlState: Some(SqlState(E23502))"; + assert_eq!(bound_error_message(error), error); + } + + /// A checkpoint written before the bound existed must not reintroduce an + /// unbounded message on restore. + #[test] + fn from_api_type_bounds_restored_messages() { + let restored = ConnectorErrorList::from_api_type(vec![ConnectorError { + timestamp: Utc::now(), + index: 1, + tag: Some("pg_exec".to_string()), + message: oversized_error(), + }]); + + let stored = restored.to_api_type(); + assert_eq!(stored.len(), 1); + assert!( + stored[0].message.len() <= MAX_CONNECTOR_ERROR_LEN + 64, + "restored message is {} bytes", + stored[0].message.len() + ); + } } diff --git a/crates/adapters/src/integrated/postgres/error.rs b/crates/adapters/src/integrated/postgres/error.rs index 7b809a86600..ca4d4e00c5f 100644 --- a/crates/adapters/src/integrated/postgres/error.rs +++ b/crates/adapters/src/integrated/postgres/error.rs @@ -1,4 +1,5 @@ use anyhow::anyhow; +use postgres::error::SqlState; pub(super) enum BackoffError { Temporary(anyhow::Error), @@ -6,6 +7,38 @@ pub(super) enum BackoffError { } impl BackoffError { + /// Classifies a failure to open a connection. + /// + /// Almost every way that connecting can fail is worth waiting out: the + /// server may be starting up or recovering from a crash, refusing + /// connections during maintenance, out of connection slots, or unreachable. + /// So this denies a short list rather than allowing one, unlike + /// [`From`], which classifies statements failing on a + /// connection that is already established. + /// + /// Only a configuration the connector cannot outlast is permanent. Retrying + /// a wrong password or a database that does not exist would spin until + /// someone changes the connector's configuration, which restarts it anyway. + pub fn connecting(value: postgres::Error) -> Self { + let permanent = value.code().is_some_and(|code| { + [ + SqlState::INVALID_PASSWORD, + SqlState::INVALID_AUTHORIZATION_SPECIFICATION, + SqlState::INVALID_CATALOG_NAME, + SqlState::INSUFFICIENT_PRIVILEGE, + ] + .contains(code) + }); + + // Chain rather than interpolate, so that the server's message survives: + // see the note in `From`. + if permanent { + Self::Permanent(anyhow::Error::new(value).context("cannot connect to postgres")) + } else { + Self::Temporary(anyhow::Error::new(value).context("cannot connect to postgres yet")) + } + } + pub fn should_retry(&self) -> bool { match self { BackoffError::Temporary(_) => true, @@ -30,12 +63,15 @@ impl BackoffError { } } +/// Classifies a statement failing on an established connection. +/// +/// Use [`BackoffError::connecting`] for a failure to open the connection in the +/// first place, where far more of the failures are transient. impl From for BackoffError { fn from(value: postgres::Error) -> Self { - use postgres::error::SqlState; - - if value.is_closed() - || value.code().is_some_and(|c| { + let code = value.code().cloned(); + let temporary = value.is_closed() + || code.as_ref().is_some_and(|c| { [ SqlState::CONNECTION_FAILURE, SqlState::CONNECTION_DOES_NOT_EXIST, @@ -46,14 +82,22 @@ impl From for BackoffError { .contains(c) }) // value.code() is none when connection is refused by the OS - || value.code().is_none() - { - Self::Temporary(anyhow!("failed to connect to postgres: {value}")) + || code.is_none(); + + // Keep the postgres error as the source instead of interpolating it: + // `Display for postgres::Error` reports only the kind ("db error"), and + // the server's message and DETAIL reach the report solely through the + // error chain, which `BackoffError::inner` formats in full. + if temporary { + Self::Temporary( + anyhow::Error::new(value) + .context(format!("postgres error: transient: SqlState: {code:?}")), + ) } else { - Self::Permanent(anyhow!( - "postgres error: permanent: SqlState: {:?}: {value}", - value.code() - )) + Self::Permanent( + anyhow::Error::new(value) + .context(format!("postgres error: permanent: SqlState: {code:?}")), + ) } } } diff --git a/crates/adapters/src/integrated/postgres/output.rs b/crates/adapters/src/integrated/postgres/output.rs index e0d2b16b23a..d241fd120bc 100644 --- a/crates/adapters/src/integrated/postgres/output.rs +++ b/crates/adapters/src/integrated/postgres/output.rs @@ -6,7 +6,7 @@ use std::{io::Write, str::FromStr, sync::Weak, time::Duration}; use super::{ error::BackoffError, prepared_statements::PreparedStatements, tls::make_tls_connector, }; -use crate::{ControllerError, util::indexed_operation_type}; +use crate::{ControllerError, PipelineState, util::indexed_operation_type}; use crate::{ buffer_op, catalog::{RecordFormat, SerBatchReader, SerCursor}, @@ -14,10 +14,11 @@ use crate::{ flush_op, format::{Encoder, OutputConsumer}, transport::OutputEndpoint, - util::IndexedOperationType, + util::{IndexedOperationType, truncate_ellipse}, }; use anyhow::{Context, Result as AnyResult, anyhow, bail}; use feldera_adapterlib::catalog::SplitCursorBuilder; +use feldera_adapterlib::format::MAX_RECORD_LEN_IN_ERRMSG; use feldera_adapterlib::transport::{AsyncErrorCallback, CommandHandler, OutputBatchType, Step}; use feldera_types::{ format::json::JsonFlavor, @@ -46,6 +47,47 @@ enum WorkerResult { Err(anyhow::Error), } +/// Maximum number of times a worker writes one batch. +/// +/// This bounds repeated failures once a connection has come back; waiting for +/// postgres to return happens inside each attempt, in +/// [`PostgresWorker::retry_connecting_with_backoff`], and is not bounded. +const MAX_BATCH_WRITE_ATTEMPTS: usize = 10; + +/// The partition of a batch assigned to a worker, retained until the batch +/// commits so that the worker can write it again. +/// +/// [`SplitCursorBuilder`] holds the batch and hands out a fresh cursor per +/// call, and the controller keeps the batch alive across +/// `batch_start`/`encode`/`batch_end` regardless, so retaining this costs a +/// reference count rather than a copy of the data. +struct PendingBatch { + cursor: SplitCursorBuilder, + /// The extra column values captured when the batch was first encoded. + /// Writing it again reuses them, so a rewrite produces the same rows even if + /// the configured values have changed since. + extra_columns: BTreeMap>, +} + +/// What has become of the transaction a worker is writing its batch into. +/// +/// Postgres aborts a transaction on any statement error and answers a subsequent +/// `COMMIT` with `ROLLBACK` rather than an error, so a commit reporting success +/// proves nothing on its own: only this tells the worker whether the rows it +/// counted reached the table. +#[derive(Clone, Copy, PartialEq, Eq)] +enum TransactionState { + /// Every statement has succeeded, so the transaction can commit. + Open, + /// The connection dropped, taking the transaction and every row it held. + /// + /// Any rows written as part of this transaction are lost and must be written again. + Lost, + /// Postgres rejected a statement and aborted the transaction. Writing the + /// batch again would hit the same rejection. + Rejected, +} + /// A single postgres worker that owns a connection and runs on a dedicated thread. struct PostgresWorker { worker_idx: usize, @@ -66,8 +108,26 @@ struct PostgresWorker { key_schema: Relation, value_schema: Relation, controller: Weak, + /// Bytes written by statements in the current transaction. + /// + /// Counted only after `execute` succeeds, and reported to the endpoint only + /// if the transaction then commits cleanly, so that it measures what reached + /// the table rather than what the worker attempted. num_bytes: usize, + /// Rows written by statements in the current transaction, counted and + /// reported on the same terms as [`Self::num_bytes`]. num_rows: usize, + /// Rows encoded for the current batch, counted as they are buffered rather + /// than as statements succeed, so that a failed batch can report how many + /// rows never reached the table. Unlike [`Self::num_rows`], this count + /// covers the rows whose statements never ran because an earlier failure had + /// already aborted the transaction. + batch_rows: usize, + /// What has become of [`Self::transaction`], which decides whether the `pending_batch` + /// commits, is written again, or is dropped. + transaction_state: TransactionState, + /// The batch being written, kept for as long as it may need writing again. + pending_batch: Option, /// Shared counter of records sent to postgres in the current batch. /// Updated atomically after each successful `execute()` call across all /// workers; reset to 0 by the endpoint at batch boundaries. @@ -97,7 +157,8 @@ fn connect(config: &PostgresWriterConfig, endpoint_name: &str) -> Result return Err(BackoffError::Permanent(e)), - }?; + } + .map_err(BackoffError::connecting)?; Ok(client) } @@ -135,6 +196,9 @@ impl PostgresWorker { key_schema: key_schema.clone(), num_rows: 0, num_bytes: 0, + batch_rows: 0, + transaction_state: TransactionState::Open, + pending_batch: None, inserts: 0, upserts: 0, deletes: 0, @@ -171,29 +235,32 @@ impl PostgresWorker { name: &str, num_records: usize, ) { - loop { - match self.exec_statement_inner(stmt.clone(), &mut value, name, num_records) { - Ok(_) => return, - Err(e) => { - let retry = e.should_retry(); - let Some(controller) = self.controller.upgrade() else { - tracing::warn!("controller is shutting down: aborting"); - return; - }; - controller.output_transport_error( - self.endpoint_id, - &self.endpoint_name, - true, - e.inner(), - Some("pg_exec"), - ); - if !retry { - return; - } - self.retry_connecting_with_backoff(); - } - } + // Nothing more can reach the table through this transaction, so pushing + // the rest of the batch at it only produces more failures to report. + // The whole transaction will eventually be retried. + if self.transaction_state != TransactionState::Open { + return; } + + let Err(e) = self.exec_statement_inner(stmt, &mut value, name, num_records) else { + return; + }; + + // Retrying the statement here cannot work: recovering takes a new + // transaction and a fresh encode, which `write_batch` drives. + let recoverable = e.should_retry(); + + let Some(controller) = self.controller.upgrade() else { + tracing::warn!("controller is shutting down: aborting"); + return; + }; + controller.output_transport_error( + self.endpoint_id, + &self.endpoint_name, + !recoverable, + e.inner(), + Some("pg_exec"), + ); } fn exec_statement_inner( @@ -211,18 +278,39 @@ impl PostgresWorker { return Ok(()); } - self.num_bytes += value.len(); - let v: &str = std::str::from_utf8(value.as_slice()).map_err(|e| { BackoffError::Permanent(anyhow!("record contains non utf-8 characters: {e}")) })?; - self.transaction() + let result = self + .transaction() .map_err(BackoffError::Permanent)? - .execute(&stmt, &[&v]) - .map_err(|e| { - BackoffError::from(e).context(format!("while executing {name} statement: {v}")) - })?; + .execute(&stmt, &[&v]); + + if let Err(e) = result { + // Report only a prefix of the payload: it holds a whole batch of + // records, up to `max_buffer_size_bytes`, which is far too large + // to keep in the error list served by `/stats`. + let error = BackoffError::from(e).context(format!( + "while executing {name} statement for {num_records} record(s) ({} bytes), which start with: {}", + v.len(), + truncate_ellipse(v, MAX_RECORD_LEN_IN_ERRMSG, "...") + )); + + // Transaction state is decided by whether the error is retry-able. + self.transaction_state = if error.should_retry() { + TransactionState::Lost + } else { + TransactionState::Rejected + }; + + return Err(error); + } + + // Count what this statement wrote. The transaction still has to commit + // for any of it to survive, which `batch_end_inner` decides. + self.num_bytes += v.len(); + self.num_rows += num_records; // Report progress: these records have now been sent to postgres within // the open transaction. The endpoint resets the counter to 0 at batch @@ -317,7 +405,12 @@ impl PostgresWorker { return Ok(()); }; + // Losing the connection rolled back the open transaction, so drop what + // it had written from the count. `batch_end_inner` fails for the rest of + // this batch, since there is no transaction left to commit. self.transaction = None; + self.num_bytes = 0; + self.num_rows = 0; self.client = connect(&self.config, &self.endpoint_name)?; self.prepared_statements = PreparedStatements::new( @@ -346,12 +439,12 @@ These statements were successfully prepared before reconnecting. Does the table let backoff = 1000; let mut n_retries = 1; - let Some(controller) = self.controller.upgrade() else { - tracing::warn!("controller is shutting down: aborting"); - return; - }; - loop { + if self.shutting_down() { + tracing::warn!("controller is shutting down: aborting"); + return; + } + tracing::info!( "worker-thread-{} retrying to connect to postgres", self.worker_idx @@ -359,27 +452,59 @@ These statements were successfully prepared before reconnecting. Does the table match self.retry_connecting() { Ok(_) => return, Err(e) => { + // Only a failure that ends the loop leaves the endpoint + // unable to reach postgres. let retry = e.should_retry(); - controller.output_transport_error( - self.endpoint_id, - &self.endpoint_name, - true, - e.inner(), - Some("pg_conn_retry"), - ); + if let Some(controller) = self.controller.upgrade() { + controller.output_transport_error( + self.endpoint_id, + &self.endpoint_name, + !retry, + e.inner(), + Some("pg_conn_retry"), + ); + } if !retry { return; } } } - std::thread::sleep( + if !self.backoff_unless_shutdown( Duration::from_millis(backoff * n_retries).min(Duration::from_secs(60)), - ); + ) { + tracing::warn!("controller is shutting down: aborting"); + return; + } n_retries += 1; } } + /// Whether the pipeline has been asked to stop. + fn shutting_down(&self) -> bool { + match self.controller.upgrade() { + None => true, + Some(controller) => controller.status.state() == PipelineState::Terminated, + } + } + + /// Waits `duration` before the next connection attempt, giving up as soon as + /// the pipeline is asked to stop. Returns whether the wait ran to completion. + fn backoff_unless_shutdown(&self, duration: Duration) -> bool { + const SLICE: Duration = Duration::from_millis(100); + + let mut remaining = duration; + while !remaining.is_zero() { + if self.shutting_down() { + return false; + } + let slice = remaining.min(SLICE); + std::thread::sleep(slice); + remaining -= slice; + } + true + } + fn batch_start_inner(&mut self) -> Result<(), BackoffError> { // Skip the controller check in test/bench mode. #[cfg(not(any(test, feature = "bench-mode")))] @@ -388,14 +513,30 @@ These statements were successfully prepared before reconnecting. Does the table return Ok(()); } + // Roll back and release anything a previous batch left open before + // borrowing the client again. The transmute below hides that borrow from + // the compiler, so it is on this function to keep two of them from + // overlapping. + self.transaction = None; + let txn = self.client.transaction()?; // SAFETY: The transaction borrows `self.client`. Both live on this - // worker's dedicated thread and never move. The transaction is committed - // or rolled back in `batch_end_inner` before the next batch. + // worker's dedicated thread and never move. The borrow ends before the + // next one begins: `batch_end_inner` commits or rolls the transaction + // back, and the line above releases it on the paths that do not reach + // `batch_end_inner`. let transaction: postgres::Transaction<'static> = unsafe { std::mem::transmute(txn) }; self.transaction = Some(transaction); + // Start counting afresh: whatever the previous transaction wrote was + // either reported at its commit or lost with it. Writing a batch again + // encodes it from the start, so its rows are counted from the start too. + self.num_bytes = 0; + self.num_rows = 0; + self.batch_rows = 0; + self.transaction_state = TransactionState::Open; + Ok(()) } @@ -403,6 +544,19 @@ These statements were successfully prepared before reconnecting. Does the table fn batch_end_inner(&mut self) -> Result<(usize, usize), BackoffError> { self.flush(); + // Do not try to commit a transaction the server has already discarded: + // report a retryable failure so that the batch is written again. + if self.transaction_state == TransactionState::Lost { + self.transaction = None; + self.num_bytes = 0; + self.num_rows = 0; + return Err(BackoffError::Temporary(anyhow!( + "the connection to PostgreSQL dropped before the rows for table {:?} \ + could be committed", + self.table + ))); + } + let transaction = self .transaction .take() @@ -410,6 +564,24 @@ These statements were successfully prepared before reconnecting. Does the table "postgres: attempted to commit a transaction that hasn't been started" )))?; + // Roll back explicitly rather than committing: Postgres would answer the + // commit with `ROLLBACK` and no error, which would report the rows this + // transaction wrote as written when they were all discarded. + if self.transaction_state == TransactionState::Rejected { + let rollback = match transaction.rollback() { + Ok(()) => String::new(), + Err(e) => format!(" Rolling the transaction back also failed: {e}."), + }; + return Err(BackoffError::Permanent(anyhow!( + "PostgreSQL rejected one of the statements writing to table {:?} and \ + aborted the transaction, so {} row(s) from the current batch were \ + dropped. The rejected statement is reported as a separate \ + error.{rollback}", + self.table, + self.batch_rows, + ))); + } + transaction.commit()?; let num_bytes = std::mem::take(&mut self.num_bytes); @@ -418,6 +590,105 @@ These statements were successfully prepared before reconnecting. Does the table Ok((num_bytes, num_rows)) } + /// Write the current batch, and write it again from the start if the + /// connection drops partway. + /// + /// Losing the connection rolls the transaction back, so the rows it held are + /// gone and the payloads that carried them have been drained from the + /// buffers. Writing the batch again from [`Self::pending_batch`] is what + /// keeps it from being lost, and is safe precisely because the rollback + /// returned the table to its pre-batch state. + /// + /// A failure that is not retryable ends the batch: rewriting it would hit + /// the same error, so the caller reports it instead. + /// + /// One case stays ambiguous. When the commit itself fails on a lost + /// connection, whether the server applied it is unknowable, and this writes + /// the batch again. In `materialized` mode that is harmless, since an insert + /// resolves conflicts and updates and deletes are idempotent. In `cdc` mode + /// the target is append-only, so a batch that did commit gains a duplicate + /// set of rows. Duplicates are the lesser evil for an event log, and this + /// connector does not claim to be fault tolerant either way. + fn write_batch(&mut self) -> AnyResult<(usize, usize)> { + let mut last_error = None; + + for attempt in 0..MAX_BATCH_WRITE_ATTEMPTS { + if attempt > 0 { + self.retry_connecting_with_backoff(); + + // The connection can drop again while restarting, which is just + // another attempt rather than the end of the batch. + match self.restart_batch() { + Ok(()) => (), + Err(e) if e.should_retry() => { + last_error = Some(e.inner()); + continue; + } + Err(e) => return Err(e.inner()), + } + } + + match self.batch_end_inner() { + Ok(totals) => return Ok(totals), + Err(e) if e.should_retry() => { + let error = e.inner(); + tracing::error!( + "postgres: worker-thread-{} lost its connection while writing a batch \ + (attempt {} of {MAX_BATCH_WRITE_ATTEMPTS}), writing it again: {error}", + self.worker_idx, + attempt + 1, + ); + last_error = Some(error); + } + Err(e) => return Err(e.inner()), + } + } + + Err(anyhow!( + "gave up writing to table {:?} after {MAX_BATCH_WRITE_ATTEMPTS} attempts, so \ + these rows were dropped; the last failure was: {}", + self.table, + last_error.expect("the loop records a failure before it gives up") + )) + } + + /// Open a fresh transaction and encode the retained batch into it again. + fn restart_batch(&mut self) -> Result<(), BackoffError> { + // The buffers may hold a partly built payload from the failed attempt. + // `batch_start_inner` resets the counts and the transaction state. + self.insert_buf.clear(); + self.upsert_buf.clear(); + self.delete_buf.clear(); + self.inserts = 0; + self.upserts = 0; + self.deletes = 0; + + self.batch_start_inner()?; + + // A worker that was handed no partition of this batch has nothing to + // write again; its transaction is simply empty. + let Some(pending) = self.pending_batch.take() else { + return Ok(()); + }; + let result = { + let mut cursor = pending.cursor.build(); + self.encode_cursor(&mut cursor, pending.extra_columns.clone()) + }; + self.pending_batch = Some(pending); + + // Encoding fails only on serialization, which will fail the same way + // however many times the batch is written again. + if let Err(e) = result { + // This ends the batch, and no one downstream commits or rolls back + // the transaction just opened, so release it here rather than leave + // the session idle in a transaction until the next batch starts. + self.transaction = None; + return Err(BackoffError::Permanent(e)); + } + + Ok(()) + } + /// Encode records from the cursor into postgres within the current transaction. fn encode_cursor( &mut self, @@ -481,7 +752,8 @@ These statements were successfully prepared before reconnecting. Does the table } }; - self.num_rows += 1; + // One row buffered, counted whether or not its statement runs. + self.batch_rows += 1; } cursor.step_key(); @@ -578,43 +850,46 @@ impl PostgresWorker { } }, WorkerCommand::Encode(cursor_builder) => { - let mut cursor = cursor_builder.build(); let extra_columns = self.extra_columns.read().clone(); - match self.encode_cursor(&mut cursor, extra_columns) { - Ok(()) => { - let _ = result_tx.send(WorkerResult::Ok { - num_bytes: 0, - num_rows: 0, - }); - } - Err(e) => { - let _ = result_tx.send(WorkerResult::Err(e)); - } - } + let result = { + let mut cursor = cursor_builder.build(); + self.encode_cursor(&mut cursor, extra_columns.clone()) + }; + // Hold on to the batch: if the connection drops before the + // transaction commits, `write_batch` encodes it again. + self.pending_batch = Some(PendingBatch { + cursor: cursor_builder, + extra_columns, + }); + let _ = match result { + Ok(()) => result_tx.send(WorkerResult::Ok { + num_bytes: 0, + num_rows: 0, + }), + Err(e) => result_tx.send(WorkerResult::Err(e)), + }; + } + WorkerCommand::Broadcast(BroadcastCommand::BatchEnd) => { + let _ = match self.write_batch() { + Ok((num_bytes, num_rows)) => result_tx.send(WorkerResult::Ok { + num_bytes, + num_rows, + }), + Err(e) => result_tx.send(WorkerResult::Err(e)), + }; + // Release the batch, and with it the worker's reference to + // the data the controller handed it. + self.pending_batch = None; + + // The next batch borrows the client again behind a + // transmute, so no transaction may outlive this one. Every + // test that writes a batch checks this. + debug_assert!( + self.transaction.is_none(), + "worker-thread-{} finished a batch with its transaction still open", + self.worker_idx + ); } - WorkerCommand::Broadcast(BroadcastCommand::BatchEnd) => loop { - match self.batch_end_inner() { - Ok((num_bytes, num_rows)) => { - let _ = result_tx.send(WorkerResult::Ok { - num_bytes, - num_rows, - }); - break; - } - Err(e) => { - if e.should_retry() { - tracing::error!( - "error when trying to commit transaction, retrying with backoff: {}", - e.inner() - ); - self.retry_connecting_with_backoff(); - continue; - } - let _ = result_tx.send(WorkerResult::Err(e.inner())); - break; - } - } - }, WorkerCommand::Broadcast(BroadcastCommand::Shutdown) => break, } } @@ -675,7 +950,15 @@ impl PostgresOutputEndpointCommandHandler { impl CommandHandler for PostgresOutputEndpointCommandHandler { fn command(&self, command: serde_json::Value) -> AnyResult { let command = serde_json::from_value::(command.clone()) - .map_err(|e| anyhow!("Postgres output connector failed to parse command '{command}' with the following error: {e}"))?; + .map_err(|e| { + let command = truncate_ellipse( + &command.to_string(), + MAX_RECORD_LEN_IN_ERRMSG, + "...", + ) + .into_owned(); + anyhow!("Postgres output connector failed to parse command '{command}' with the following error: {e}") + })?; match command { PostgresOutputEndpointCommand::SetExtraColumns(extra_columns) => { @@ -876,7 +1159,15 @@ impl PostgresOutputEndpoint { .map(|e| format!("{e:#}")) .collect::>() .join("; "); - bail!("{} worker(s) failed: {msg}", errors.len()); + // With one writer thread, which is the default, naming the count + // adds nothing to what the failure itself already says. + if errors.len() == 1 { + bail!("{msg}"); + } + bail!( + "{} of the connector's writer threads failed: {msg}", + errors.len() + ); } Ok(()) @@ -930,26 +1221,30 @@ impl OutputConsumer for PostgresOutputEndpoint { // the records have been committed; on failure the counter would // otherwise leak across batches. self.records_written.store(0, Ordering::Relaxed); + + // Report what committed even when a worker failed. Each worker commits + // its own transaction, so a batch that failed in one worker still wrote + // the rows the others committed, and `broadcast_and_collect` accumulated + // only those. Taking the totals here, rather than on the success path + // alone, also keeps a failed batch from carrying them into the next one. + let elapsed = self.txn_start.elapsed(); + let num_bytes = std::mem::take(&mut self.num_bytes); + let num_rows = std::mem::take(&mut self.num_rows); + let Some(controller) = self.controller.upgrade() else { + tracing::warn!("controller is shutting down: aborting"); + return; + }; + controller + .status + .output_buffer(self.endpoint_id, num_bytes, num_rows); + match result { Ok(()) => { - let elapsed = self.txn_start.elapsed(); - let num_bytes = std::mem::take(&mut self.num_bytes); - let num_rows = std::mem::take(&mut self.num_rows); tracing::debug!( "postgres: flushed {num_rows} rows and {num_bytes} bytes in {elapsed:?}", ); - - if let Some(controller) = self.controller.upgrade() { - controller - .status - .output_buffer(self.endpoint_id, num_bytes, num_rows); - }; } Err(err) => { - let Some(controller) = self.controller.upgrade() else { - tracing::warn!("controller is shutting down: aborting"); - return; - }; controller.output_transport_error( self.endpoint_id, &self.endpoint_name, @@ -1014,7 +1309,15 @@ impl Encoder for PostgresOutputEndpoint { .map(|e| format!("{e:#}")) .collect::>() .join("; "); - bail!("{} worker(s) failed: {msg}", errors.len()); + // With one writer thread, which is the default, naming the count + // adds nothing to what the failure itself already says. + if errors.len() == 1 { + bail!("{msg}"); + } + bail!( + "{} of the connector's writer threads failed: {msg}", + errors.len() + ); } Ok(()) diff --git a/crates/adapters/src/integrated/postgres/prepared_statements.rs b/crates/adapters/src/integrated/postgres/prepared_statements.rs index 6dc15ce5feb..99ea07c0600 100644 --- a/crates/adapters/src/integrated/postgres/prepared_statements.rs +++ b/crates/adapters/src/integrated/postgres/prepared_statements.rs @@ -1,4 +1,5 @@ use super::error::BackoffError; +use crate::util::truncate_ellipse_middle; use feldera_types::{ program_schema::Relation, transport::postgres::{PostgresWriteMode, PostgresWriterConfig}, @@ -6,6 +7,12 @@ use feldera_types::{ use itertools::Itertools; use postgres::Statement; +/// Maximum length of a generated query echoed back in an error message. +/// +/// A query names every column of the target table, so a wide table yields a +/// query too long to keep in the error list served by `/stats`. +const MAX_QUERY_LEN_IN_ERRMSG: usize = 2048; + #[derive(Debug, Default)] struct RawQueries { insert: String, @@ -158,7 +165,7 @@ impl PreparedStatements { .map_err(|e| { BackoffError::from(e).context(format!( "failed to prepare insert statement: `{}`: {err_msg}", - &raw_queries.insert + truncate_ellipse_middle(&raw_queries.insert, MAX_QUERY_LEN_IN_ERRMSG) )) })?; let upsert = client @@ -166,7 +173,7 @@ impl PreparedStatements { .map_err(|e| { BackoffError::from(e).context(format!( "failed to prepare update statement: `{}`: {err_msg}", - &raw_queries.upsert + truncate_ellipse_middle(&raw_queries.upsert, MAX_QUERY_LEN_IN_ERRMSG) )) })?; let delete = client @@ -174,7 +181,7 @@ impl PreparedStatements { .map_err(|e| { BackoffError::from(e).context(format!( "failed to prepare delete statement: `{}`: {err_msg}", - raw_queries.delete + truncate_ellipse_middle(&raw_queries.delete, MAX_QUERY_LEN_IN_ERRMSG) )) })?; diff --git a/crates/adapters/src/integrated/postgres/test.rs b/crates/adapters/src/integrated/postgres/test.rs index fbd1eabab5d..8bc5daf9129 100644 --- a/crates/adapters/src/integrated/postgres/test.rs +++ b/crates/adapters/src/integrated/postgres/test.rs @@ -21,6 +21,7 @@ use tempfile::NamedTempFile; use crate::{ Catalog, CircuitCatalog, Controller, + controller::MAX_CONNECTOR_ERROR_LEN, integrated::postgres::test::pg::PostgresTestStructCdc, test::{TestStruct, wait}, }; @@ -1069,6 +1070,892 @@ fn test_pg_insert_omitted_default_column() { ); } +/// Bytes reported by the `while executing ... statement for N record(s) (M +/// bytes)` context that the connector attaches to a failed statement. +fn payload_bytes_in_error(message: &str) -> Option { + message + .split(" record(s) (") + .nth(1)? + .split(" bytes)") + .next()? + .parse() + .ok() +} + +/// A statement Postgres rejects must not carry the whole batch into the error +/// message. +/// +/// The payload of one statement is as large as `max_buffer_size_bytes` (1 MiB by +/// default) and an endpoint retains up to `MAX_CONNECTOR_ERRORS` messages per +/// tag, so quoting the payload in full grew `/stats` and +/// `/output_endpoints/{name}/stats` past the size limit the pipeline manager +/// applies while proxying them, failing the whole request. +#[test] +#[serial] +fn test_pg_error_message_is_bounded() { + let table_name = unique_pg_name("test_pg_errmsg"); + let url = postgres_url(); + + // One record is about 590 bytes, so 10000 records fill the connector's + // default 1 MiB buffer several times over. + let data: Vec = (0..10000).map(|_| rand::random()).collect(); + + let mut temp_file = NamedTempFile::new().unwrap(); + for datum in data.iter() { + let mut serializer = serde_json::Serializer::new(Vec::new()); + datum + .serialize_with_context(&mut serializer, &SqlSerdeConfig::default()) + .unwrap(); + temp_file + .as_file_mut() + .write_all(&serializer.into_inner()) + .unwrap(); + temp_file.write_all(b"\n").unwrap(); + } + + let config = serde_json::from_value(json!({ + "name": "test", + "workers": 4, + // Step on large batches so that the connector fills its 1 MiB buffer, + // instead of flushing a handful of records at a time. + "min_batch_size_records": 5000, + "max_buffering_delay_usecs": 5000000, + "inputs": { + "ins": { + "stream": "test_input1", + "transport": { "name": "file_input", "config": { "path": temp_file.path(), "byte_size_buffer": 8388608 } }, + "format": { "name": "json", "config": { "update_format": "raw", "array": false } } + } + }, + "outputs": { + "test_output1": { + "stream": "test_output1", + "transport": { "name": "postgres_output", "config": { "uri": url, "table": &table_name } }, + "index": "idx" + } + } + })) + .unwrap(); + + // `freshness_timestamp` is NOT NULL, has no DEFAULT, and is absent from the + // Feldera view, so every INSERT the connector issues violates the + // constraint. + let _table = PostgresTestStruct::create_table_with_extra_columns( + &table_name, + url, + true, + &None, + "freshness_timestamp TIMESTAMP NOT NULL", + ); + + let (controller, _err_receiver) = PostgresTestStruct::test_circuit(config); + controller.start(); + + wait( + || { + controller + .status() + .output_status() + .values() + .any(|endpoint| { + endpoint + .metrics + .num_transport_errors + .load(std::sync::atomic::Ordering::Relaxed) + > 0 + }) + }, + 60_000, + ) + .expect("timeout: connector reported no transport error"); + + // Snapshot the errors and shut the pipeline down before asserting, so that a + // failing assertion does not race the table teardown. + let (errors, fatal_error) = { + let status = controller.status(); + let outputs = status.output_status(); + let endpoint = outputs.values().next().expect("output endpoint"); + ( + endpoint.transport_errors.lock().unwrap().to_api_type(), + endpoint.fatal_error(), + ) + }; + let _ = controller.stop(); + + assert!(!errors.is_empty(), "no transport error was recorded"); + + for error in &errors { + assert!( + !error.message.contains("bytes elided"), + "connector produced a {}-byte message that the status layer had to \ + truncate: {}", + error.message.len(), + error.message + ); + assert!( + error.message.len() <= MAX_CONNECTOR_ERROR_LEN, + "error message is {} bytes", + error.message.len() + ); + } + + // Truncating must keep the diagnosis, which Postgres reports at the end of + // the error chain: the SQLSTATE, the server's message, and the DETAIL + // naming the offending row. + let diagnosed = errors.iter().find(|e| e.message.contains("E23502")); + let diagnosed = diagnosed + .unwrap_or_else(|| panic!("no error reports the SQLSTATE: {errors:#?}")) + .message + .as_str(); + assert!( + diagnosed.contains("violates not-null constraint"), + "error drops the server's message: {diagnosed}" + ); + assert!( + diagnosed.contains("DETAIL: Failing row contains"), + "error drops the server's DETAIL: {diagnosed}" + ); + + // At least one failed statement has + // to have carried a payload far larger than the message describing it. + let largest_payload = errors + .iter() + .filter_map(|e| payload_bytes_in_error(&e.message)) + .max() + .expect("no error reports its payload size"); + assert!( + largest_payload > 64 * 1024, + "the connector never buffered a large batch (largest was {largest_payload} bytes), \ + so this test would pass without truncation" + ); + + let fatal_error = fatal_error.expect("endpoint reported no fatal error"); + assert!( + fatal_error.len() <= MAX_CONNECTOR_ERROR_LEN, + "fatal_error is {} bytes", + fatal_error.len() + ); +} + +/// Serializes `data` as the newline-delimited JSON into a temp file the file input reads. +fn write_records_to_temp_file(data: &[PostgresTestStruct]) -> NamedTempFile { + let mut temp_file = NamedTempFile::new().unwrap(); + for datum in data { + let mut serializer = serde_json::Serializer::new(Vec::new()); + datum + .serialize_with_context(&mut serializer, &SqlSerdeConfig::default()) + .unwrap(); + temp_file + .as_file_mut() + .write_all(&serializer.into_inner()) + .unwrap(); + temp_file.write_all(b"\n").unwrap(); + } + temp_file +} + +/// Pipeline that feeds the records in `path` to a postgres output connector +/// writing to `table_name` through `threads` connections. +/// +/// `min_batch_size_records`, when set, makes the pipeline ingest that +/// many records in one large batch. +fn pg_output_test_config( + path: &Path, + url: &str, + table_name: &str, + threads: usize, + min_batch_size_records: Option, +) -> PipelineConfig { + let mut config = json!({ + "name": "test", + "workers": 4, + "inputs": { + "ins": { + "stream": "test_input1", + "transport": { "name": "file_input", "config": { "path": path } }, + "format": { "name": "json", "config": { "update_format": "raw", "array": false } } + } + }, + "outputs": { + "test_output1": { + "stream": "test_output1", + "transport": { "name": "postgres_output", "config": { "uri": url, "table": table_name, "threads": threads } }, + "index": "idx" + } + } + }); + + if let Some(records) = min_batch_size_records { + config["min_batch_size_records"] = json!(records); + config["max_buffering_delay_usecs"] = json!(5_000_000); + config["inputs"]["ins"]["transport"]["config"]["byte_size_buffer"] = json!(8_388_608); + } + + serde_json::from_value(config).unwrap() +} + +fn pg_transmitted_totals(controller: &Controller) -> (u64, u64) { + let status = controller.status(); + let outputs = status.output_status(); + let endpoint = outputs.values().next().expect("output endpoint"); + ( + endpoint + .metrics + .transmitted_records + .load(std::sync::atomic::Ordering::Relaxed), + endpoint + .metrics + .transmitted_bytes + .load(std::sync::atomic::Ordering::Relaxed), + ) +} + +fn pg_total_processed_input_records(controller: &Controller) -> u64 { + let status = controller.status(); + let outputs = status.output_status(); + let endpoint = outputs.values().next().expect("output endpoint"); + endpoint + .metrics + .total_processed_input_records + .load(std::sync::atomic::Ordering::Relaxed) +} + +fn pg_row_count(url: &str, table_name: &str) -> u64 { + let mut client = pg::pg_connect(url, &None); + let count: i64 = client + .query_one(&format!("SELECT count(*) FROM {table_name}"), &[]) + .expect("failed to count rows") + .get(0); + count as u64 +} + +/// `transmitted_records` counts what Postgres committed, not what the connector +/// buffered. +/// +/// A failing statement aborts the whole transaction, and Postgres answers the +/// following COMMIT with ROLLBACK instead of an error, so a commit that reports +/// success proves nothing on its own. +#[test] +#[serial] +fn test_pg_transmitted_records_excludes_rolled_back_rows() { + let table_name = unique_pg_name("test_pg_rollback_count"); + let url = postgres_url(); + let verify_url = url.clone(); + + let data: Vec = (0..1000).map(|_| rand::random()).collect(); + let temp_file = write_records_to_temp_file(&data); + let config = pg_output_test_config(temp_file.path(), &url, &table_name, 1, None); + + // `freshness_timestamp` is NOT NULL, has no DEFAULT, and is absent from the + // Feldera view, so every INSERT the connector issues aborts its transaction + // and the batch writes nothing. + let _table = PostgresTestStruct::create_table_with_extra_columns( + &table_name, + url, + true, + &None, + "freshness_timestamp TIMESTAMP NOT NULL", + ); + + let (controller, _err_receiver) = PostgresTestStruct::test_circuit(config); + controller.start(); + + // Wait until the connector has finished handling every record. `total_processed_input_records` + // reports what the connector handled whether or not the write succeeded, so + // it advances even though nothing reaches the table, which keeps the + // assertions below from passing against counters that have yet to move. + wait( + || pg_total_processed_input_records(&controller) >= data.len() as u64, + 60_000, + ) + .expect("timeout: connector did not finish handling the input"); + + let (records, bytes) = pg_transmitted_totals(&controller); + let _ = controller.stop(); + + assert_eq!( + records, 0, + "reported {records} records as transmitted, but every transaction was rolled back" + ); + assert_eq!( + bytes, 0, + "reported {bytes} bytes as transmitted, but every transaction was rolled back" + ); + assert_eq!( + pg_row_count(&verify_url, &table_name), + 0, + "the table should have rejected every row" + ); +} + +/// Records reported by the `... statement for N record(s)` context that the +/// connector attaches to a failed statement. +fn records_in_error(message: &str) -> Option { + message + .split(" statement for ") + .nth(1)? + .split(" record(s)") + .next()? + .parse() + .ok() +} + +/// A transaction that fails partway discards the statements that already +/// succeeded, so none of the batch counts as transmitted. +#[test] +#[serial] +fn test_pg_transmitted_records_excludes_rows_lost_to_a_later_failure() { + const RECORDS: usize = 3000; + const REJECTED: &str = "__reject__"; + + let table_name = unique_pg_name("test_pg_partial_rollback"); + let url = postgres_url(); + let verify_url = url.clone(); + + // Enough records to fill the connector's 1 MiB buffer more than once, so the + // batch issues several statements to postgres. Keys ascend and the cursor walks them in + // order, which puts the rejected record in the last statement and leaves + // every statement before it valid. + let mut data: Vec = (0..RECORDS).map(|_| rand::random()).collect(); + for (i, datum) in data.iter_mut().enumerate() { + datum.bigint_ = i as i64; + } + data.last_mut().unwrap().varchar_ = REJECTED.to_string().into(); + + let temp_file = write_records_to_temp_file(&data); + + let config = serde_json::from_value(json!({ + "name": "test", + "workers": 4, + // Step once over the whole input, so it becomes a single transaction. + "min_batch_size_records": RECORDS, + "max_buffering_delay_usecs": 5000000, + "inputs": { + "ins": { + "stream": "test_input1", + "transport": { "name": "file_input", "config": { "path": temp_file.path(), "byte_size_buffer": 8388608 } }, + "format": { "name": "json", "config": { "update_format": "raw", "array": false } } + } + }, + "outputs": { + "test_output1": { + "stream": "test_output1", + "transport": { "name": "postgres_output", "config": { "uri": url, "table": &table_name } }, + "index": "idx" + } + } + })) + .unwrap(); + + let _table = PostgresTestStruct::create_table_with_extra_columns( + &table_name, + url, + true, + &None, + &format!("CHECK (varchar_ <> '{REJECTED}')"), + ); + + let (controller, _err_receiver) = PostgresTestStruct::test_circuit(config); + controller.start(); + + wait( + || pg_total_processed_input_records(&controller) >= RECORDS as u64, + 60_000, + ) + .expect("timeout: connector did not finish handling the input"); + + let (records, bytes) = pg_transmitted_totals(&controller); + let errors = { + let status = controller.status(); + let outputs = status.output_status(); + outputs + .values() + .next() + .expect("output endpoint") + .transport_errors + .lock() + .unwrap() + .to_api_type() + }; + let _ = controller.stop(); + + // Guard against the test going vacuous: the rejected record has to arrive in + // a statement that carries only part of the batch, so that an earlier + // statement succeeded and had rows to lose. + let failed_records = errors + .iter() + .filter_map(|e| records_in_error(&e.message)) + .min() + .expect("no statement failure was reported"); + assert!( + failed_records < RECORDS, + "the whole batch went out as one statement of {failed_records} records, so no earlier \ + statement succeeded and this test would pass without the commit guard" + ); + + assert_eq!( + records, 0, + "reported {records} records as transmitted, but the transaction was rolled back" + ); + assert_eq!( + bytes, 0, + "reported {bytes} bytes as transmitted, but the transaction was rolled back" + ); + assert_eq!( + pg_row_count(&verify_url, &table_name), + 0, + "the aborted transaction should have left the table empty" + ); +} + +/// Selects the connector's backends that are inside a batch transaction, which +/// is where the tests below inject their faults. +/// +/// Postgres serves each connection from a process of its own, a backend, lists +/// one row per backend in `pg_stat_activity`, and terminates one on +/// `pg_terminate_backend(pid)`. That is the whole mechanism: a thread of the +/// test connects as an ordinary client and terminates the connector's backends, +/// which the connector sees as its connection dropping, taking with it the +/// transaction it was writing the batch into. The backends to terminate are the +/// ones connected to this database (`datname`), inside a transaction +/// (`xact_start`), whose latest statement is one of the connector's writes, +/// which all call `jsonb_populate_recordset`; `pg_backend_pid` excludes the +/// injecting thread's own backend. The connector prepares those same statements +/// when it connects, so `query` alone would also match a backend that has yet to +/// write anything; `xact_start` is what excludes it, preparation running outside +/// a transaction. +/// +/// The window these tests wait for is the span in which terminating a backend +/// lands mid-batch: from the connector's first write statement until its +/// transaction ends. It covers the whole batch rather than the instant a +/// statement is running, because Postgres keeps reporting the last statement in +/// `query` while the backend sits idle inside its transaction. Terminating a +/// backend before the window opens costs the connector no more than an idle +/// connection; terminating one after the window closes finds the batch already +/// committed. Neither exercises the reconnect path under test. +const BATCH_BACKENDS: &str = "SELECT pid FROM pg_stat_activity \ + WHERE pid <> pg_backend_pid() AND datname = current_database() \ + AND xact_start IS NOT NULL \ + AND query ILIKE '%jsonb_populate_recordset%'"; + +/// Polls for up to eight seconds until the window is open, that is, until a +/// connector backend is inside a batch transaction. +fn wait_for_batch_transaction(client: &mut postgres::Client) -> bool { + for _ in 0..40_000 { + let found: i64 = client + .query_one( + &format!("SELECT count(*) FROM ({BATCH_BACKENDS}) backends"), + &[], + ) + .map(|row| row.get(0)) + .unwrap_or(0); + if found > 0 { + return true; + } + std::thread::sleep(std::time::Duration::from_micros(200)); + } + false +} + +/// Terminates every connector backend that is inside a batch transaction, and +/// reports whether it found one to terminate. +fn terminate_batch_backends(client: &mut postgres::Client) -> bool { + let terminated: i64 = client + .query_one( + &format!( + "SELECT count(*) FROM \ + (SELECT pg_terminate_backend(pid) FROM ({BATCH_BACKENDS}) backends) terminated" + ), + &[], + ) + .map(|row| row.get(0)) + .unwrap_or(0); + terminated > 0 +} + +/// Drops the connector's connections mid-batch, `rounds` times over, and reports +/// how many rounds landed. +/// +/// A round lands when it catches an open window and terminates a backend in it. +/// A round that misses ends the loop; the caller asserts on the count, so a test +/// whose fault never landed fails instead of passing. +fn terminate_connector_backends(url: &str, rounds: usize) -> usize { + let mut client = pg::pg_connect(url, &None); + let mut landed = 0; + + for _ in 0..rounds { + if !wait_for_batch_transaction(&mut client) || !terminate_batch_backends(&mut client) { + break; + } + landed += 1; + } + + landed +} + +/// Losing the connection in the middle of a batch must not lose the batch. +/// +/// Reconnecting is not enough on its own: the dropped connection rolled its +/// transaction back, and the payloads that carried those rows were drained from +/// the buffers as they flushed, so the batch has to be encoded again from the +/// cursor the worker retained. +/// +/// `threads` is the number of connections the connector writes through, each +/// with a transaction and a partition of the batch of its own. `kills` is how +/// many connection drops to inject, one round after another: any drop past the +/// first lands while the connector is writing the batch again. +fn pg_reconnect_test(name: &str, threads: usize, kills: usize) { + const RECORDS: usize = 3000; + + let table_name = unique_pg_name(name); + let url = postgres_url(); + let verify_url = url.clone(); + + // Enough records to fill the connector's 1 MiB buffer more than once, so + // that statements have already written rows into the transaction by the time + // the connection drops. + let mut data: Vec = (0..RECORDS).map(|_| rand::random()).collect(); + for (i, datum) in data.iter_mut().enumerate() { + datum.bigint_ = i as i64; + } + let temp_file = write_records_to_temp_file(&data); + + let config = pg_output_test_config(temp_file.path(), &url, &table_name, threads, Some(RECORDS)); + + let _table = PostgresTestStruct::create_table(&table_name, url, true, &None); + + // Build the pipeline before arming the killer: the connector connects and + // prepares its statements while the controller is being constructed. + let (controller, _err_receiver) = PostgresTestStruct::test_circuit(config); + + // The killer is the fault injector: a thread that terminates the connector's + // backends over SQL from a connection of its own, as `BATCH_BACKENDS` describes. + let killer_url = verify_url.clone(); + let killer = std::thread::spawn(move || terminate_connector_backends(&killer_url, kills)); + + controller.start(); + let landed = killer.join().expect("killer thread panicked"); + + // The endpoint reports the batch as processed whether or not it wrote it, so + // this settles the test in both outcomes. + wait( + || pg_total_processed_input_records(&controller) >= RECORDS as u64, + 120_000, + ) + .expect("timeout: connector did not finish handling the input"); + + let (records, _bytes) = pg_transmitted_totals(&controller); + let fatal_error = { + let status = controller.status(); + let outputs = status.output_status(); + outputs + .values() + .next() + .expect("output endpoint") + .fatal_error() + }; + let _ = controller.stop(); + + assert_eq!( + landed, kills, + "landed {landed} of {kills} connection drops, so this test proves less than it should" + ); + assert_eq!( + pg_row_count(&verify_url, &table_name), + RECORDS as u64, + "rows were lost when the connection dropped" + ); + assert_eq!( + records, RECORDS as u64, + "reported {records} of {RECORDS} rows as transmitted" + ); + + assert_eq!( + fatal_error, None, + "the connector recovered, yet the endpoint is marked as fatally failed" + ); +} + +#[test] +#[serial] +fn test_pg_reconnect_does_not_lose_the_batch() { + pg_reconnect_test("test_pg_reconnect", 1, 1); +} + +#[test] +#[serial] +fn test_pg_reconnect_does_not_lose_the_batch_multi_thread() { + pg_reconnect_test("test_pg_reconnect_mt", 3, 1); +} + +#[test] +#[serial] +fn test_pg_reconnect_survives_repeated_drops() { + pg_reconnect_test("test_pg_reconnect_repeat", 1, 3); +} + +/// Makes the connector's database refuse new connections, and lets them through +/// again when dropped. +/// +/// `ALTER DATABASE ... ALLOW_CONNECTIONS false` turns away even a superuser, so +/// it stands in for a server that is restarting or under maintenance. +struct RefusedConnections { + admin: postgres::Client, + database: String, + refusing: bool, +} + +impl RefusedConnections { + /// Opens the connection this needs up front, so that starting to refuse is a + /// single statement once the moment arrives. + fn prepare(url: &str) -> Self { + let database: String = pg::pg_connect(url, &None) + .query_one("SELECT current_database()", &[]) + .expect("failed to read the database name") + .get(0); + + // ALTER DATABASE cannot run from a session connected to that database. + let admin = pg::pg_connect(&Self::other_database_url(url), &None); + + Self { + admin, + database, + refusing: false, + } + } + + fn refuse(&mut self) { + let refuse = Self::allow_connections(&self.database, false); + self.admin + .execute(&refuse, &[]) + .expect("failed to make the database refuse connections"); + self.refusing = true; + } + + fn other_database_url(url: &str) -> String { + match url::Url::parse(url) { + Ok(mut parsed) => { + parsed.set_path("/template1"); + parsed.to_string() + } + Err(_) => format!("{url}/template1"), + } + } + + fn allow_connections(database: &str, allow: bool) -> String { + format!(r#"ALTER DATABASE "{database}" WITH ALLOW_CONNECTIONS {allow}"#) + } +} + +/// Restoring on drop matters: leaving the database closed would fail every later test. +impl Drop for RefusedConnections { + fn drop(&mut self) { + if !self.refusing { + return; + } + let restore = Self::allow_connections(&self.database, true); + self.admin + .execute(&restore, &[]) + .expect("failed to let connections through again"); + } +} + +/// A database that refuses the reconnect must be waited out, not given up on. +/// +/// This is the one path in the connector that dropping a backend does not reach: +/// there the reconnect succeeds on the first try, so `retry_connecting_with_backoff` never +/// enters its error branch. Refusing connections for a few seconds forces that +/// branch, and the batch still has to arrive once they are let through. +#[test] +#[serial] +fn test_pg_reconnect_waits_out_a_refusing_database() { + const RECORDS: usize = 3000; + + let table_name = unique_pg_name("test_pg_refused"); + let url = postgres_url(); + let verify_url = url.clone(); + + let mut data: Vec = (0..RECORDS).map(|_| rand::random()).collect(); + for (i, datum) in data.iter_mut().enumerate() { + datum.bigint_ = i as i64; + } + let temp_file = write_records_to_temp_file(&data); + let config = pg_output_test_config(temp_file.path(), &url, &table_name, 1, Some(RECORDS)); + + let _table = PostgresTestStruct::create_table(&table_name, url, true, &None); + let (controller, _err_receiver) = PostgresTestStruct::test_circuit(config); + + let injector_url = verify_url.clone(); + let injector = std::thread::spawn(move || { + // Every connection this thread needs has to exist before the database + // starts turning connections away. + let mut watcher = pg::pg_connect(&injector_url, &None); + let mut refused = RefusedConnections::prepare(&injector_url); + + if !wait_for_batch_transaction(&mut watcher) { + return false; + } + refused.refuse(); + let terminated = terminate_batch_backends(&mut watcher); + + // Outlast the first backoff steps, one and two seconds, so that more + // than one reconnect fails. + std::thread::sleep(std::time::Duration::from_secs(4)); + drop(refused); + + terminated + }); + + controller.start(); + let injected = injector.join().expect("injector thread panicked"); + + wait( + || pg_total_processed_input_records(&controller) >= RECORDS as u64, + 180_000, + ) + .expect("timeout: connector did not finish handling the input"); + + let (records, _bytes) = pg_transmitted_totals(&controller); + let (errors, fatal_error) = { + let status = controller.status(); + let outputs = status.output_status(); + let endpoint = outputs.values().next().expect("output endpoint"); + ( + endpoint.transport_errors.lock().unwrap().to_api_type(), + endpoint.fatal_error(), + ) + }; + let _ = controller.stop(); + + assert!( + injected, + "never dropped the connector's connection while the database was refusing them, \ + so this test proves nothing" + ); + // Only `retry_connecting_with_backoff` reports under this tag, so its error + // branch ran. + assert!( + errors + .iter() + .any(|e| e.tag.as_deref() == Some("pg_conn_retry")), + "no reconnect attempt failed, so the retry loop was never exercised: {errors:#?}" + ); + assert_eq!( + pg_row_count(&verify_url, &table_name), + RECORDS as u64, + "rows were lost while the database refused connections" + ); + assert_eq!( + records, RECORDS as u64, + "reported {records} of {RECORDS} rows as transmitted" + ); + assert_eq!( + fatal_error, None, + "the connector recovered, yet the endpoint is marked as fatally failed" + ); +} + +/// The pipeline must be able to stop while postgres is unreachable. +/// +/// A worker waiting out a database it cannot reach sits in a backoff loop, and +/// the endpoint's drop joins that thread, so anything that keeps the loop from +/// noticing the shutdown hangs the pipeline rather than merely stalling its +/// output. +#[test] +#[serial] +fn test_pg_stops_while_the_database_is_unreachable() { + const RECORDS: usize = 3000; + const STOP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + + let table_name = unique_pg_name("test_pg_stop_unreachable"); + let url = postgres_url(); + let verify_url = url.clone(); + + let mut data: Vec = (0..RECORDS).map(|_| rand::random()).collect(); + for (i, datum) in data.iter_mut().enumerate() { + datum.bigint_ = i as i64; + } + let temp_file = write_records_to_temp_file(&data); + let config = pg_output_test_config(temp_file.path(), &url, &table_name, 1, Some(RECORDS)); + + let _table = PostgresTestStruct::create_table(&table_name, url, true, &None); + let (controller, _err_receiver) = PostgresTestStruct::test_circuit(config); + + // Both connections have to exist before the database starts turning them + // away. `refused` lets connections through again when it drops, including + // while this test unwinds, which releases the worker either way. + let mut watcher = pg::pg_connect(&verify_url, &None); + let mut refused = RefusedConnections::prepare(&verify_url); + + controller.start(); + assert!( + wait_for_batch_transaction(&mut watcher), + "never caught the connector inside its batch transaction" + ); + refused.refuse(); + assert!( + terminate_batch_backends(&mut watcher), + "never dropped the connector's connection" + ); + + // Let the worker settle into the backoff loop before asking it to stop. + std::thread::sleep(std::time::Duration::from_secs(3)); + + let (stopped_sender, stopped) = crossbeam::channel::bounded(1); + std::thread::spawn(move || { + let _ = controller.stop(); + let _ = stopped_sender.send(()); + }); + + // Wait with a deadline rather than blocking: a regression here would + // otherwise hang the whole test run instead of failing this one test. + assert!( + stopped.recv_timeout(STOP_TIMEOUT).is_ok(), + "the pipeline did not stop within {STOP_TIMEOUT:?} while postgres was unreachable" + ); +} + +/// The `transmitted_records` counter still reports every row of a batch that commits. +#[test] +#[serial] +fn test_pg_transmitted_records_counts_committed_rows() { + let table_name = unique_pg_name("test_pg_commit_count"); + let url = postgres_url(); + let verify_url = url.clone(); + + let data: Vec = (0..1000).map(|_| rand::random()).collect(); + let temp_file = write_records_to_temp_file(&data); + let config = pg_output_test_config(temp_file.path(), &url, &table_name, 1, None); + + let _table = PostgresTestStruct::create_table(&table_name, url, true, &None); + + let (controller, err_receiver) = PostgresTestStruct::test_circuit(config); + controller.start(); + + wait( + || pg_transmitted_totals(&controller).0 >= data.len() as u64 || !err_receiver.is_empty(), + 60_000, + ) + .expect("timeout: connector did not report the committed rows"); + + let (records, bytes) = pg_transmitted_totals(&controller); + let _ = controller.stop(); + + assert!(err_receiver.is_empty(), "connector reported an error"); + assert_eq!( + records, + data.len() as u64, + "reported {records} of {} rows as transmitted", + data.len() + ); + assert!(bytes > 0, "reported 0 bytes for {records} rows"); + assert_eq!( + pg_row_count(&verify_url, &table_name), + records, + "the counter and the table disagree" + ); +} + #[test] #[serial] fn test_pg_insert() { diff --git a/crates/adapters/src/util.rs b/crates/adapters/src/util.rs index 9c94c1cead7..a58f3fd94bd 100644 --- a/crates/adapters/src/util.rs +++ b/crates/adapters/src/util.rs @@ -158,6 +158,28 @@ pub fn indexed_operation_type( }) } +/// Largest offset `<= index` that falls on a char boundary of `s`. +fn floor_char_boundary(s: &str, index: usize) -> usize { + let mut index = index.min(s.len()); + while index > 0 && !s.is_char_boundary(index) { + index -= 1; + } + index +} + +/// Smallest offset `>= index` that falls on a char boundary of `s`. +fn ceil_char_boundary(s: &str, index: usize) -> usize { + let mut index = index.min(s.len()); + while index < s.len() && !s.is_char_boundary(index) { + index += 1; + } + index +} + +/// Truncates `s` to `len` bytes, appending `ellipse` if anything was dropped. +/// +/// The cut lands on a char boundary, so the retained prefix never exceeds `len` +/// bytes. pub(crate) fn truncate_ellipse<'a>(s: &'a str, len: usize, ellipse: &str) -> Cow<'a, str> { if s.len() <= len { return Cow::Borrowed(s); @@ -165,10 +187,37 @@ pub(crate) fn truncate_ellipse<'a>(s: &'a str, len: usize, ellipse: &str) -> Cow return Cow::Borrowed(""); } - let result = s.chars().take(len).chain(ellipse.chars()).collect(); + let mut result = String::with_capacity(len + ellipse.len()); + result.push_str(&s[..floor_char_boundary(s, len)]); + result.push_str(ellipse); Cow::Owned(result) } +/// Truncates `s` to roughly `len` bytes by cutting out the middle, reporting how +/// many bytes were dropped. +/// +/// Use this instead of [`truncate_ellipse`] when the tail of the string carries +/// as much information as the head. An `anyhow` error chain formatted with +/// `{:?}` is the motivating case: it leads with the outermost context and ends +/// with the root cause, so dropping the tail throws away the actual failure. +pub(crate) fn truncate_ellipse_middle(s: &str, len: usize) -> Cow<'_, str> { + if s.len() <= len { + return Cow::Borrowed(s); + } + + let head_len = len / 2; + let head_end = floor_char_boundary(s, head_len); + let tail_start = ceil_char_boundary(s, s.len() - (len - head_len)); + + format!( + "{}\n... [{} bytes elided] ...\n{}", + &s[..head_end], + tail_start - head_end, + &s[tail_start..] + ) + .into() +} + pub(crate) fn missing_pipeline_identity_message(operation: &str) -> String { format!( "{operation}: pipeline has no system-assigned name (config.name), which is necessary to verify ownership" @@ -587,7 +636,78 @@ mod test { time::Duration, }; - use crate::util::{RateLimitCheckResult, TokenBucketRateLimiter}; + use crate::util::{ + RateLimitCheckResult, TokenBucketRateLimiter, truncate_ellipse, truncate_ellipse_middle, + }; + + // -------- Truncation tests -------- // + + #[test] + fn truncate_ellipse_leaves_short_strings_alone() { + assert_eq!(truncate_ellipse("abc", 3, "..."), "abc"); + assert_eq!(truncate_ellipse("abc", 4, "..."), "abc"); + assert_eq!(truncate_ellipse("", 0, "..."), ""); + } + + #[test] + fn truncate_ellipse_bounds_the_prefix_in_bytes() { + assert_eq!(truncate_ellipse("abcdef", 3, "..."), "abc..."); + assert_eq!(truncate_ellipse("abcdef", 0, "..."), ""); + + // A multi-byte string must be measured in bytes, not chars: 10 chars of + // 3 bytes each may not be reported as fitting in a 10-byte budget. + let s = "日".repeat(10); + let truncated = truncate_ellipse(&s, 10, "..."); + assert_eq!(truncated, "日日日..."); + assert!(truncated.len() - "...".len() <= 10); + } + + #[test] + fn truncate_ellipse_middle_keeps_both_ends() { + let s = format!("head{}tail", "x".repeat(1000)); + let truncated = truncate_ellipse_middle(&s, 100); + + assert!(truncated.starts_with("head"), "{truncated}"); + assert!(truncated.ends_with("tail"), "{truncated}"); + assert!(truncated.contains("bytes elided"), "{truncated}"); + // The retained text is bounded even though the marker adds to it. + assert!(truncated.len() < 100 + 64, "{}", truncated.len()); + } + + #[test] + fn truncate_ellipse_middle_leaves_short_strings_alone() { + assert_eq!(truncate_ellipse_middle("abcdef", 6), "abcdef"); + assert_eq!(truncate_ellipse_middle("abcdef", 7), "abcdef"); + assert!(!truncate_ellipse_middle("abcdef", 6).contains("elided")); + } + + // Cutting in the middle of a multi-byte char would panic on a slice. + #[test] + fn truncate_ellipse_middle_respects_char_boundaries() { + let s = "日".repeat(100); + for len in 0..32 { + let truncated = truncate_ellipse_middle(&s, len); + assert!(truncated.contains("bytes elided"), "len={len}"); + } + } + + // The elided-byte count must add up: head + elided + tail == original. + #[test] + fn truncate_ellipse_middle_reports_the_elided_byte_count() { + let s = "abcdefghijklmnopqrstuvwxyz".repeat(10); + let truncated = truncate_ellipse_middle(&s, 40); + + let elided: usize = truncated + .split("... [") + .nth(1) + .and_then(|s| s.split(" bytes elided] ...").next()) + .expect("elision marker") + .parse() + .expect("elided byte count"); + let retained = + truncated.len() - "\n... [ bytes elided] ...\n".len() - elided.to_string().len(); + assert_eq!(retained + elided, s.len()); + } // -------- Basic single-thread tests -------- //