From cf30457ba9751f0e00e9bef57ea640aee6284fb9 Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Fri, 14 Aug 2026 09:26:15 -0700 Subject: [PATCH 01/10] [adapters] bound error messages stored in endpoint status Cap each connector error message at 8 KiB before storing it in the endpoint status, and cap `fatal_error` the same way. An endpoint retains up to `MAX_CONNECTOR_ERRORS` (100) messages per tag, and `output_endpoint_status` always serializes both error lists, so a connector that quotes its payload in an error message grows `/output_endpoints/{name}/stats` without limit. The pipeline manager proxies that response under a 50 MB limit and fails the whole request when the body exceeds it, which leaves the user unable to read the errors at all. `fatal_error` reaches plain `/stats` unconditionally, so an oversized message also weighs down every poll and every checkpoint that captures the error lists. Truncation drops the middle of the message. A message built from an `anyhow` chain leads with the outermost context and ends with the root cause, so keeping only the head would discard the diagnosis. A connector is still expected to bound the data it quotes; this cap is the backstop for when it does not. Also fix `truncate_ellipse`, which compared byte lengths but then took that many chars, so its documented byte bound overshot by up to 4x on multi-byte text. Signed-off-by: Leonid Ryzhyk --- crates/adapters/src/controller/stats.rs | 113 +++++++++++++++++++-- crates/adapters/src/util.rs | 124 +++++++++++++++++++++++- 2 files changed, 226 insertions(+), 11 deletions(-) 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/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 -------- // From 62fae4b9f5c624e4ae704c60f0d54099f18b14f4 Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Fri, 14 Aug 2026 09:32:17 -0700 Subject: [PATCH 02/10] [adapters] pg-out: keep the postgres error detail in the error chain Chain the `postgres::Error` as the source of the `BackoffError` instead of interpolating it into a fresh `anyhow!` message. `Display for postgres::Error` reports only the kind, so interpolating it yielded "postgres error: permanent: SqlState: Some(SqlState(E23502)): db error" and nothing more: the server's message and DETAIL live in the error's source, which the interpolation discarded. Chaining keeps them, and `BackoffError::inner` already formats the whole chain. tokio-postgres appended the source to `Display` up to 0.7.13 and stopped in 0.7.14, so this silently emptied out every postgres connector error when we upgraded. Reported errors now carry the diagnosis: while executing insert statement for 1721 record(s) ... Caused by: 0: postgres error: permanent: SqlState: Some(SqlState(E23502)) 1: db error 2: ERROR: null value in column "freshness_timestamp" of relation "..." violates not-null constraint DETAIL: Failing row contains (...). Signed-off-by: Leonid Ryzhyk --- .../adapters/src/integrated/postgres/error.rs | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/crates/adapters/src/integrated/postgres/error.rs b/crates/adapters/src/integrated/postgres/error.rs index 7b809a86600..d1486660420 100644 --- a/crates/adapters/src/integrated/postgres/error.rs +++ b/crates/adapters/src/integrated/postgres/error.rs @@ -34,8 +34,9 @@ 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 +47,19 @@ 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("failed to connect to postgres")) } else { - Self::Permanent(anyhow!( - "postgres error: permanent: SqlState: {:?}: {value}", - value.code() - )) + Self::Permanent( + anyhow::Error::new(value) + .context(format!("postgres error: permanent: SqlState: {code:?}")), + ) } } } From 3dd80898313f764da1414269b097665c2a329d91 Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Fri, 14 Aug 2026 09:34:34 -0700 Subject: [PATCH 03/10] [adapters] pg-out: quote only a prefix of the batch in error messages Report the record count, the payload size and a 4 KiB prefix when a statement fails, instead of the entire payload. The payload holds one buffered batch, up to `max_buffer_size_bytes` (1 MiB by default), and the endpoint retains up to 100 messages per tag, so quoting it in full grew `/output_endpoints/{name}/stats` past the 50 MB limit the pipeline manager applies while proxying the response. A customer hit exactly that and could not read the 872 errors their connector reported; the offending `fatal_error` alone was 1,049,481 bytes, and it also reached every plain `/stats` poll. Bound the two other places that echo unbounded text into an error: the generated query, which names every column of the target table, and the JSON body of a `set_extra_columns` command. Signed-off-by: Leonid Ryzhyk --- crates/adapters/src/controller.rs | 4 + .../src/integrated/postgres/output.rs | 22 ++- .../postgres/prepared_statements.rs | 13 +- .../adapters/src/integrated/postgres/test.rs | 171 ++++++++++++++++++ 4 files changed, 204 insertions(+), 6 deletions(-) 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/integrated/postgres/output.rs b/crates/adapters/src/integrated/postgres/output.rs index e0d2b16b23a..fad08bbcb42 100644 --- a/crates/adapters/src/integrated/postgres/output.rs +++ b/crates/adapters/src/integrated/postgres/output.rs @@ -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, @@ -221,7 +222,14 @@ impl PostgresWorker { .map_err(BackoffError::Permanent)? .execute(&stmt, &[&v]) .map_err(|e| { - BackoffError::from(e).context(format!("while executing {name} statement: {v}")) + // 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`. + 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, "...") + )) })?; // Report progress: these records have now been sent to postgres within @@ -675,7 +683,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) => { 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..67243c3fc45 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,176 @@ 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. This is the customer-reported failure. + 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 { + // The bound enforced by the status layer is a backstop; the connector is + // expected to quote no more than a prefix of the payload on its own. + // "bytes elided" means the backstop had to fire. + 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}" + ); + + // Guard against the test going vacuous: 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() + ); +} + #[test] #[serial] fn test_pg_insert() { From 60543192b41f8ac4f890bd4dce4fc31d9afd4def Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Fri, 14 Aug 2026 13:33:52 -0700 Subject: [PATCH 04/10] [adapters] pg-out: count only the rows Postgres committed Count rows and bytes per statement that `execute` accepted, and report them to the endpoint only once the transaction commits cleanly. `transmitted_records` counted rows as `encode_cursor` buffered them, so it measured what the connector attempted rather than what reached the table. A failing statement aborts the whole transaction, and Postgres answers the following COMMIT with ROLLBACK instead of an error, so the connector took the commit as success and reported the batch as transmitted. Counting per successful statement is not enough on its own, because the rows of a statement that succeeded before a later one failed are rolled back too. Track whether the transaction has been poisoned and roll it back explicitly rather than issuing a commit whose success means nothing. Report the totals whether or not every worker succeeded. Each worker commits its own transaction, so a batch that failed in one worker still wrote the rows the others committed; reporting only on the success path also carried those rows into the next batch's total. The error the rollback reports names the target table and says that the rows PostgreSQL had already accepted went with it, since someone who reads only the rejected statement expects the rest of the rows to have landed. Reporting a single failed writer thread no longer prefixes the count, which is noise in the default single-thread configuration. Signed-off-by: Leonid Ryzhyk --- .../src/integrated/postgres/output.rs | 142 ++++++-- .../adapters/src/integrated/postgres/test.rs | 302 ++++++++++++++++++ 2 files changed, 413 insertions(+), 31 deletions(-) diff --git a/crates/adapters/src/integrated/postgres/output.rs b/crates/adapters/src/integrated/postgres/output.rs index fad08bbcb42..dc01385f745 100644 --- a/crates/adapters/src/integrated/postgres/output.rs +++ b/crates/adapters/src/integrated/postgres/output.rs @@ -67,8 +67,21 @@ 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, + /// Whether a statement in the current transaction has failed. + /// + /// 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. + txn_poisoned: bool, /// 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. @@ -136,6 +149,7 @@ impl PostgresWorker { key_schema: key_schema.clone(), num_rows: 0, num_bytes: 0, + txn_poisoned: false, inserts: 0, upserts: 0, deletes: 0, @@ -212,25 +226,34 @@ 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| { - // 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`. - 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, "...") - )) - })?; + .execute(&stmt, &[&v]); + + if let Err(e) = result { + // Postgres has aborted the transaction, so nothing written in it so + // far will reach the table either. + self.txn_poisoned = true; + + // 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`. + return Err(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, "...") + ))); + } + + // 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 @@ -325,7 +348,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( @@ -404,6 +432,12 @@ These statements were successfully prepared before reconnecting. Does the table 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. + self.num_bytes = 0; + self.num_rows = 0; + self.txn_poisoned = false; + Ok(()) } @@ -418,6 +452,34 @@ 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.txn_poisoned { + let rollback = match transaction.rollback() { + Ok(()) => String::new(), + Err(e) => format!(" Rolling the transaction back also failed: {e}."), + }; + // Spell out that the rows PostgreSQL accepted are gone too. A reader + // who saw only the rejected statement would expect the rest of the + // rows to have landed. + let accepted = if self.num_rows > 0 { + format!( + ": the {} row(s) PostgreSQL had already accepted were rolled back, \ + and the rest were dropped", + self.num_rows + ) + } else { + String::new() + }; + return Err(BackoffError::Permanent(anyhow!( + "PostgreSQL rejected one of the statements writing to table {:?} and \ + aborted the transaction, so no rows were written{accepted}. The \ + rejected statement is reported as a separate error.{rollback}", + self.table, + ))); + } + transaction.commit()?; let num_bytes = std::mem::take(&mut self.num_bytes); @@ -488,8 +550,6 @@ These statements were successfully prepared before reconnecting. Does the table self.upsert(buf); } }; - - self.num_rows += 1; } cursor.step_key(); @@ -892,7 +952,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(()) @@ -946,26 +1014,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, @@ -1030,7 +1102,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/test.rs b/crates/adapters/src/integrated/postgres/test.rs index 67243c3fc45..9b32d824629 100644 --- a/crates/adapters/src/integrated/postgres/test.rs +++ b/crates/adapters/src/integrated/postgres/test.rs @@ -1240,6 +1240,308 @@ fn test_pg_error_message_is_bounded() { ); } +/// Serializes `data` as the newline-delimited JSON 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`. +fn pg_output_test_config(path: &Path, url: &str, table_name: &str) -> PipelineConfig { + serde_json::from_value(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 } }, + "index": "idx" + } + } + })) + .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. Counting rows as they were encoded +/// therefore reported a batch as transmitted that the server had discarded. +#[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); + + // `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. This counter + // 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. +/// +/// This is the case the commit guard exists for. Counting rows per successful +/// statement is not enough on its own: the rows of the earlier statement would +/// still be reported, because the failure aborts the transaction and the +/// following COMMIT silently rolls them back. +#[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 becomes several statements. 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" + ); +} + +/// The 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); + + 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() { From 0b063a21ce58fdaca1b12a4428e49182bcf8979b Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Fri, 14 Aug 2026 13:33:59 -0700 Subject: [PATCH 05/10] [adapters] pg-out: write the batch again when the connection drops Retain each worker's partition of the batch until it commits, and encode it into a fresh transaction if the connection drops partway. A dropped connection lost the whole batch. `retry_connecting` clears the open transaction and nothing reopened one, so the statement being retried and every statement after it failed with "transaction that hasn't been created yet", and the commit failed too. Even with a transaction the retry could not have recovered the batch: the payloads that carried the earlier rows had been drained from the buffers as they flushed, so the connector no longer held them. `exec_statement` no longer retries a statement in place, which could never work, and stops issuing statements once the batch is doomed instead of reporting one failure per remaining flush. Two limits are worth knowing. A commit that fails on a lost connection leaves it unknowable whether the server applied it, and this writes the batch again: harmless in materialized mode, a duplicate set of rows in CDC mode, which is documented on `write_batch`. And the batch_records_written gauge counts the abandoned attempt's rows as well as the rewrite's until the next batch boundary resets it. Signed-off-by: Leonid Ryzhyk --- .../src/integrated/postgres/output.rs | 255 ++++++++++++++---- .../adapters/src/integrated/postgres/test.rs | 144 ++++++++++ 2 files changed, 342 insertions(+), 57 deletions(-) diff --git a/crates/adapters/src/integrated/postgres/output.rs b/crates/adapters/src/integrated/postgres/output.rs index dc01385f745..13f0cc79892 100644 --- a/crates/adapters/src/integrated/postgres/output.rs +++ b/crates/adapters/src/integrated/postgres/output.rs @@ -47,6 +47,28 @@ 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>, +} + /// A single postgres worker that owns a connection and runs on a dedicated thread. struct PostgresWorker { worker_idx: usize, @@ -82,6 +104,15 @@ struct PostgresWorker { /// subsequent `COMMIT` with `ROLLBACK` rather than an error, so a commit /// reporting success proves nothing on its own. txn_poisoned: bool, + /// Whether the connection dropped while writing the current batch, which + /// rolled the transaction back and lost every row it held. + /// + /// The batch is written again from [`Self::pending_batch`] rather than + /// resuming, since the payloads that carried those rows were drained from + /// the buffers as they flushed. + needs_replay: bool, + /// 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. @@ -150,6 +181,8 @@ impl PostgresWorker { num_rows: 0, num_bytes: 0, txn_poisoned: false, + needs_replay: false, + pending_batch: None, inserts: 0, upserts: 0, deletes: 0, @@ -186,29 +219,34 @@ 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. + if self.needs_replay || self.txn_poisoned { + return; + } + + let Err(e) = self.exec_statement_inner(stmt, &mut value, name, num_records) else { + return; + }; + + // Retrying the statement here cannot work: a retryable failure means the + // connection is gone, and recovering takes a new transaction and a fresh + // encode, which `write_batch` drives. + if e.should_retry() { + self.needs_replay = true; } + + 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"), + ); } fn exec_statement_inner( @@ -437,6 +475,7 @@ These statements were successfully prepared before reconnecting. Does the table self.num_bytes = 0; self.num_rows = 0; self.txn_poisoned = false; + self.needs_replay = false; Ok(()) } @@ -445,6 +484,23 @@ 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. + // + // A lost connection also poisons the transaction, so this is checked + // first: it is the case that can be recovered by writing the batch + // again, whereas a statement Postgres rejected cannot be. + if self.needs_replay { + 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() @@ -488,6 +544,97 @@ 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 failure flags. + 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. + result.map_err(BackoffError::Permanent) + } + /// Encode records from the cursor into postgres within the current transaction. fn encode_cursor( &mut self, @@ -646,43 +793,37 @@ 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; } - 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, } } diff --git a/crates/adapters/src/integrated/postgres/test.rs b/crates/adapters/src/integrated/postgres/test.rs index 9b32d824629..e165b3ad64b 100644 --- a/crates/adapters/src/integrated/postgres/test.rs +++ b/crates/adapters/src/integrated/postgres/test.rs @@ -1501,6 +1501,150 @@ fn test_pg_transmitted_records_excludes_rows_lost_to_a_later_failure() { ); } +/// Terminates connector backends that are inside a batch transaction, `rounds` +/// times, and reports how many rounds landed a kill. +/// +/// Requiring an open transaction whose latest query is a connector statement +/// keeps this off the connection while it prepares its statements, and widens +/// the window to the whole batch: Postgres keeps reporting the last query while +/// the backend sits idle in the transaction. +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 { + let mut caught = false; + for _ in 0..40_000 { + let terminated: i64 = client + .query_one( + "SELECT count(*) FROM (SELECT pg_terminate_backend(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%') terminated", + &[], + ) + .map(|row| row.get(0)) + .unwrap_or(0); + if terminated > 0 { + caught = true; + break; + } + std::thread::sleep(std::time::Duration::from_micros(200)); + } + if !caught { + 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 +/// batch the worker retained. +/// +/// `threads` covers the connector writing through one connection and through +/// several, which partition the batch. `kills` covers the connection dropping +/// again while the batch is being rewritten. +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: PipelineConfig = 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, "threads": threads } }, + "index": "idx" + } + } + })) + .unwrap(); + + 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); + + 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 handled 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 _ = 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" + ); +} + +#[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); +} + /// The counter still reports every row of a batch that commits. #[test] #[serial] From 706711308d46cc2fd367b6256a3572ae5506f5c4 Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Fri, 14 Aug 2026 14:28:14 -0700 Subject: [PATCH 06/10] [adapters] pg-out: stop reporting a retried failure as fatal Report a connection failure as fatal only when the connector gives up on it, not when it is about to retry. `fatal` does two things: it stamps `fatal_error` on the endpoint status, which is write-once and never cleared, and it prefixes the logged message with "FATAL". So a blip the connector recovered from within a second marked the endpoint as fatally failed for the life of the pipeline, and said so in the log. `retry_connecting_with_backoff` reported every failure as fatal, one per iteration, including the ones it went on to retry successfully. `exec_statement` did the same, which became wrong when a retryable statement failure started leading to the batch being written again rather than dropped. The reconnect tests now assert that an endpoint that recovered carries no fatal error. Signed-off-by: Leonid Ryzhyk --- crates/adapters/src/integrated/postgres/output.rs | 12 +++++++++--- crates/adapters/src/integrated/postgres/test.rs | 15 +++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/crates/adapters/src/integrated/postgres/output.rs b/crates/adapters/src/integrated/postgres/output.rs index 13f0cc79892..79e1ac46712 100644 --- a/crates/adapters/src/integrated/postgres/output.rs +++ b/crates/adapters/src/integrated/postgres/output.rs @@ -232,7 +232,8 @@ impl PostgresWorker { // Retrying the statement here cannot work: a retryable failure means the // connection is gone, and recovering takes a new transaction and a fresh // encode, which `write_batch` drives. - if e.should_retry() { + let recoverable = e.should_retry(); + if recoverable { self.needs_replay = true; } @@ -243,7 +244,7 @@ impl PostgresWorker { controller.output_transport_error( self.endpoint_id, &self.endpoint_name, - true, + !recoverable, e.inner(), Some("pg_exec"), ); @@ -433,11 +434,16 @@ 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. Reporting one it is about to + // retry as fatal stamps `fatal_error`, which never clears, + // so a blip the connector recovers from in a second would + // mark the endpoint failed for the life of the pipeline. let retry = e.should_retry(); controller.output_transport_error( self.endpoint_id, &self.endpoint_name, - true, + !retry, e.inner(), Some("pg_conn_retry"), ); diff --git a/crates/adapters/src/integrated/postgres/test.rs b/crates/adapters/src/integrated/postgres/test.rs index e165b3ad64b..46b3155e36c 100644 --- a/crates/adapters/src/integrated/postgres/test.rs +++ b/crates/adapters/src/integrated/postgres/test.rs @@ -1610,6 +1610,15 @@ fn pg_reconnect_test(name: &str, threads: usize, kills: usize) { .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!( @@ -1625,6 +1634,12 @@ fn pg_reconnect_test(name: &str, threads: usize, kills: usize) { records, RECORDS as u64, "reported {records} of {RECORDS} rows as transmitted" ); + // `fatal_error` never clears, so recording one for a drop the connector + // recovered from would leave the endpoint marked failed for good. + assert_eq!( + fatal_error, None, + "the connector recovered, yet the endpoint is marked as fatally failed" + ); } #[test] From bc207ce18bb8b85dcaaf89d5d7350ce9cdaa1039 Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Fri, 14 Aug 2026 15:43:55 -0700 Subject: [PATCH 07/10] [adapters] pg-out: wait out a database that refuses connections Classify a failure to open a connection with a new `BackoffError::connecting`, which denies a short list of permanent failures instead of allowing a short list of transient ones. The shared classification was written for statements failing on an established connection, where allowing a handful of connection-lost codes is right. Applied to connecting, it called everything else permanent, so the connector gave up on exactly the conditions the retry loop exists for. `57P03`, which is what postgres answers while it is starting up after a restart, was treated as permanent, as were `57P02` after a crash, `53300` when connection slots run out, and the refusal a database under maintenance returns. Measured: the connector abandoned the batch and wrote 0 of 3000 rows. Connecting keeps retrying unless the configuration is one the connector cannot outlast: a wrong password, a missing database, or insufficient privilege. Changing any of those restarts the connector anyway. Fold the two failure flags this leaves behind into one `TransactionState` of Open, Lost, or Rejected. They were never independent: any statement failure set `txn_poisoned`, and a retryable one then also set `needs_replay`, so a lost connection left both true and `batch_end_inner` had to test `needs_replay` first to reach the recoverable case. That precedence was load-bearing and only a comment said so. The states are now mutually exclusive by construction, one assignment replaces two, and it happens where the error is built and `should_retry` is already known. Signed-off-by: Leonid Ryzhyk --- .../adapters/src/integrated/postgres/error.rs | 44 ++- .../src/integrated/postgres/output.rs | 84 +++--- .../adapters/src/integrated/postgres/test.rs | 280 ++++++++++++++---- 3 files changed, 316 insertions(+), 92 deletions(-) diff --git a/crates/adapters/src/integrated/postgres/error.rs b/crates/adapters/src/integrated/postgres/error.rs index d1486660420..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,10 +63,12 @@ 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; - let code = value.code().cloned(); let temporary = value.is_closed() || code.as_ref().is_some_and(|c| { @@ -54,7 +89,10 @@ impl From for BackoffError { // 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("failed to connect to postgres")) + Self::Temporary( + anyhow::Error::new(value) + .context(format!("postgres error: transient: SqlState: {code:?}")), + ) } else { Self::Permanent( anyhow::Error::new(value) diff --git a/crates/adapters/src/integrated/postgres/output.rs b/crates/adapters/src/integrated/postgres/output.rs index 79e1ac46712..9255abd3d8c 100644 --- a/crates/adapters/src/integrated/postgres/output.rs +++ b/crates/adapters/src/integrated/postgres/output.rs @@ -69,6 +69,27 @@ struct PendingBatch { 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. + /// + /// The batch is written again from [`PostgresWorker::pending_batch`] rather + /// than resumed, since the payloads that carried those rows were drained + /// from the buffers as they flushed. + Lost, + /// Postgres rejected a statement and aborted the transaction. Writing the + /// batch again would hit the same rejection, so the rows are dropped. + Rejected, +} + /// A single postgres worker that owns a connection and runs on a dedicated thread. struct PostgresWorker { worker_idx: usize, @@ -98,19 +119,9 @@ struct PostgresWorker { /// Rows written by statements in the current transaction, counted and /// reported on the same terms as [`Self::num_bytes`]. num_rows: usize, - /// Whether a statement in the current transaction has failed. - /// - /// 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. - txn_poisoned: bool, - /// Whether the connection dropped while writing the current batch, which - /// rolled the transaction back and lost every row it held. - /// - /// The batch is written again from [`Self::pending_batch`] rather than - /// resuming, since the payloads that carried those rows were drained from - /// the buffers as they flushed. - needs_replay: bool, + /// What has become of [`Self::transaction`], which decides whether the 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. @@ -142,7 +153,8 @@ fn connect(config: &PostgresWriterConfig, endpoint_name: &str) -> Result return Err(BackoffError::Permanent(e)), - }?; + } + .map_err(BackoffError::connecting)?; Ok(client) } @@ -180,8 +192,7 @@ impl PostgresWorker { key_schema: key_schema.clone(), num_rows: 0, num_bytes: 0, - txn_poisoned: false, - needs_replay: false, + transaction_state: TransactionState::Open, pending_batch: None, inserts: 0, upserts: 0, @@ -221,7 +232,7 @@ impl PostgresWorker { ) { // Nothing more can reach the table through this transaction, so pushing // the rest of the batch at it only produces more failures to report. - if self.needs_replay || self.txn_poisoned { + if self.transaction_state != TransactionState::Open { return; } @@ -229,13 +240,9 @@ impl PostgresWorker { return; }; - // Retrying the statement here cannot work: a retryable failure means the - // connection is gone, and recovering takes a new transaction and a fresh - // encode, which `write_batch` drives. + // Retrying the statement here cannot work: recovering takes a new + // transaction and a fresh encode, which `write_batch` drives. let recoverable = e.should_retry(); - if recoverable { - self.needs_replay = true; - } let Some(controller) = self.controller.upgrade() else { tracing::warn!("controller is shutting down: aborting"); @@ -275,18 +282,24 @@ impl PostgresWorker { .execute(&stmt, &[&v]); if let Err(e) = result { - // Postgres has aborted the transaction, so nothing written in it so - // far will reach the table either. - self.txn_poisoned = true; - // 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`. - return Err(BackoffError::from(e).context(format!( + 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, "...") - ))); + )); + + // Either way the transaction is gone, along with everything it held. + // Whether it can be written again is what separates the two. + 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 @@ -480,8 +493,7 @@ These statements were successfully prepared before reconnecting. Does the table // either reported at its commit or lost with it. self.num_bytes = 0; self.num_rows = 0; - self.txn_poisoned = false; - self.needs_replay = false; + self.transaction_state = TransactionState::Open; Ok(()) } @@ -492,11 +504,7 @@ These statements were successfully prepared before reconnecting. Does the table // Do not try to commit a transaction the server has already discarded: // report a retryable failure so that the batch is written again. - // - // A lost connection also poisons the transaction, so this is checked - // first: it is the case that can be recovered by writing the batch - // again, whereas a statement Postgres rejected cannot be. - if self.needs_replay { + if self.transaction_state == TransactionState::Lost { self.transaction = None; self.num_bytes = 0; self.num_rows = 0; @@ -517,7 +525,7 @@ These statements were successfully prepared before reconnecting. Does the table // 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.txn_poisoned { + if self.transaction_state == TransactionState::Rejected { let rollback = match transaction.rollback() { Ok(()) => String::new(), Err(e) => format!(" Rolling the transaction back also failed: {e}."), @@ -615,7 +623,7 @@ These statements were successfully prepared before reconnecting. Does the table /// 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 failure flags. + // `batch_start_inner` resets the counts and the transaction state. self.insert_buf.clear(); self.upsert_buf.clear(); self.delete_buf.clear(); diff --git a/crates/adapters/src/integrated/postgres/test.rs b/crates/adapters/src/integrated/postgres/test.rs index 46b3155e36c..86d304bc22e 100644 --- a/crates/adapters/src/integrated/postgres/test.rs +++ b/crates/adapters/src/integrated/postgres/test.rs @@ -1258,9 +1258,19 @@ fn write_records_to_temp_file(data: &[PostgresTestStruct]) -> NamedTempFile { } /// Pipeline that feeds the records in `path` to a postgres output connector -/// writing to `table_name`. -fn pg_output_test_config(path: &Path, url: &str, table_name: &str) -> PipelineConfig { - serde_json::from_value(json!({ +/// writing to `table_name` through `threads` connections. +/// +/// `min_batch_size_records`, when set, makes the pipeline step once over that +/// many records so a test can act on one large batch, and has the file input +/// read the whole file at once to match. +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": { @@ -1273,12 +1283,19 @@ fn pg_output_test_config(path: &Path, url: &str, table_name: &str) -> PipelineCo "outputs": { "test_output1": { "stream": "test_output1", - "transport": { "name": "postgres_output", "config": { "uri": url, "table": table_name } }, + "transport": { "name": "postgres_output", "config": { "uri": url, "table": table_name, "threads": threads } }, "index": "idx" } } - })) - .unwrap() + }); + + 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) { @@ -1332,7 +1349,7 @@ fn test_pg_transmitted_records_excludes_rolled_back_rows() { 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); + 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 @@ -1501,37 +1518,58 @@ fn test_pg_transmitted_records_excludes_rows_lost_to_a_later_failure() { ); } -/// Terminates connector backends that are inside a batch transaction, `rounds` -/// times, and reports how many rounds landed a kill. +/// Selects the connector backends that are inside a batch transaction. /// /// Requiring an open transaction whose latest query is a connector statement /// keeps this off the connection while it prepares its statements, and widens /// the window to the whole batch: Postgres keeps reporting the last query while /// the backend sits idle in the transaction. +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 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. +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, and reports how +/// many rounds landed. 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 { - let mut caught = false; - for _ in 0..40_000 { - let terminated: i64 = client - .query_one( - "SELECT count(*) FROM (SELECT pg_terminate_backend(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%') terminated", - &[], - ) - .map(|row| row.get(0)) - .unwrap_or(0); - if terminated > 0 { - caught = true; - break; - } - std::thread::sleep(std::time::Duration::from_micros(200)); - } - if !caught { + if !wait_for_batch_transaction(&mut client) || !terminate_batch_backends(&mut client) { break; } landed += 1; @@ -1566,28 +1604,7 @@ fn pg_reconnect_test(name: &str, threads: usize, kills: usize) { } let temp_file = write_records_to_temp_file(&data); - let config: PipelineConfig = 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, "threads": threads } }, - "index": "idx" - } - } - })) - .unwrap(); + 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); @@ -1660,6 +1677,167 @@ 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. Restoring +/// on drop matters: leaving the database closed would fail every later test. +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}"#) + } +} + +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 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, which is the whole point of refusing the connections. + 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 counter still reports every row of a batch that commits. #[test] #[serial] @@ -1670,7 +1848,7 @@ fn test_pg_transmitted_records_counts_committed_rows() { 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); + let config = pg_output_test_config(temp_file.path(), &url, &table_name, 1, None); let _table = PostgresTestStruct::create_table(&table_name, url, true, &None); From c288817f76dc5a7b1e2b5bbbf1d03c05c1ce1c7e Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Fri, 14 Aug 2026 16:06:45 -0700 Subject: [PATCH 08/10] [adapters] pg-out: let the pipeline stop while postgres is unreachable Watch the pipeline state, not the controller's liveness, when deciding whether to keep waiting for postgres, and break the backoff into slices so the wait ends promptly. A worker that cannot reach postgres sits in `retry_connecting_with_backoff`. The loop meant to abandon that wait on shutdown by checking whether its `Weak` still upgraded, but that check could never fire. The endpoint's drop joins this thread, and the output thread is blocked on the endpoint, so the controller stays alive exactly because the worker has not finished: the loop waited for a drop that was waiting for the loop. `ControllerInner::stop` sets `Terminated` before it touches anything else, so that is the signal that actually arrives. The minute-long backoff was also one uninterruptible sleep, which would have delayed the stop even with a signal that worked. It now wakes every 100ms to check. Signed-off-by: Leonid Ryzhyk --- .../src/integrated/postgres/output.rs | 71 +++++++++++++++---- .../adapters/src/integrated/postgres/test.rs | 61 ++++++++++++++++ 2 files changed, 117 insertions(+), 15 deletions(-) diff --git a/crates/adapters/src/integrated/postgres/output.rs b/crates/adapters/src/integrated/postgres/output.rs index 9255abd3d8c..10c6c0fbb29 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}, @@ -434,12 +434,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 @@ -453,26 +453,67 @@ These statements were successfully prepared before reconnecting. Does the table // so a blip the connector recovers from in a second would // mark the endpoint failed for the life of the pipeline. let retry = e.should_retry(); - controller.output_transport_error( - self.endpoint_id, - &self.endpoint_name, - !retry, - 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. + /// + /// Waiting for the controller to be dropped would never work here: the + /// endpoint's drop joins this thread, and the endpoint itself is what the + /// output thread is blocked on, so the controller stays alive precisely + /// because this thread has not finished. `Terminated`, which + /// `ControllerInner::stop` sets before it touches anything else, is the + /// signal that does arrive. + 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. + /// + /// The wait reaches a minute, and the endpoint's drop waits on this thread, + /// so sleeping through it in one piece would delay every pipeline shutdown + /// that happens while postgres is unreachable. + 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")))] diff --git a/crates/adapters/src/integrated/postgres/test.rs b/crates/adapters/src/integrated/postgres/test.rs index 86d304bc22e..32f5ebfe0b9 100644 --- a/crates/adapters/src/integrated/postgres/test.rs +++ b/crates/adapters/src/integrated/postgres/test.rs @@ -1838,6 +1838,67 @@ fn test_pg_reconnect_waits_out_a_refusing_database() { ); } +/// 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. Holding a strong reference to the controller across the loop did +/// exactly that, since it kept alive the very thing the loop was watching for. +#[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 counter still reports every row of a batch that commits. #[test] #[serial] From 012e746ad588f8480ee5e90990d23019451132bb Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Fri, 14 Aug 2026 16:23:49 -0700 Subject: [PATCH 09/10] [adapters] pg-out: never carry a transaction into the next batch Release any transaction still held before opening the next one, and drop the one a failed rewrite leaves behind. `batch_start_inner` transmutes a borrow of `self.client` to 'static, so nothing stops two transactions from borrowing the client at once except the promise that each is finished before the next begins. Every path did honor that until `restart_batch` arrived: it opens a transaction and then encodes into it, and if the encode fails it returns an error that ends the batch, leaving the transaction on the worker. The next batch would then borrow the client while that one was still alive. Release it in two places. `restart_batch` drops it when the encode fails, which is the point where the batch actually ends, so the session does not sit idle in a transaction until the next batch starts. `batch_start_inner` drops whatever it finds before borrowing the client again, which makes overlapping borrows impossible however a future caller behaves, and the SAFETY comment now says so. Signed-off-by: Leonid Ryzhyk --- .../src/integrated/postgres/output.rs | 31 +++++++++++++++++-- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/crates/adapters/src/integrated/postgres/output.rs b/crates/adapters/src/integrated/postgres/output.rs index 10c6c0fbb29..d0fba42c1b0 100644 --- a/crates/adapters/src/integrated/postgres/output.rs +++ b/crates/adapters/src/integrated/postgres/output.rs @@ -522,11 +522,19 @@ 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); @@ -687,7 +695,15 @@ These statements were successfully prepared before reconnecting. Does the table // Encoding fails only on serialization, which will fail the same way // however many times the batch is written again. - result.map_err(BackoffError::Permanent) + 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. @@ -878,6 +894,15 @@ impl PostgresWorker { // 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::Shutdown) => break, } From 9243040fa115097adcee324508061d6b2843e556 Mon Sep 17 00:00:00 2001 From: Leonid Ryzhyk Date: Sat, 15 Aug 2026 18:39:17 -0700 Subject: [PATCH 10/10] [adapters] pg-out: address review feedback Signed-off-by: Leonid Ryzhyk --- .../src/integrated/postgres/output.rs | 60 ++++------- .../adapters/src/integrated/postgres/test.rs | 100 ++++++++++-------- 2 files changed, 81 insertions(+), 79 deletions(-) diff --git a/crates/adapters/src/integrated/postgres/output.rs b/crates/adapters/src/integrated/postgres/output.rs index d0fba42c1b0..d241fd120bc 100644 --- a/crates/adapters/src/integrated/postgres/output.rs +++ b/crates/adapters/src/integrated/postgres/output.rs @@ -81,12 +81,10 @@ enum TransactionState { Open, /// The connection dropped, taking the transaction and every row it held. /// - /// The batch is written again from [`PostgresWorker::pending_batch`] rather - /// than resumed, since the payloads that carried those rows were drained - /// from the buffers as they flushed. + /// 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, so the rows are dropped. + /// batch again would hit the same rejection. Rejected, } @@ -119,7 +117,13 @@ struct PostgresWorker { /// Rows written by statements in the current transaction, counted and /// reported on the same terms as [`Self::num_bytes`]. num_rows: usize, - /// What has become of [`Self::transaction`], which decides whether the batch + /// 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. @@ -192,6 +196,7 @@ 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, @@ -232,6 +237,7 @@ impl PostgresWorker { ) { // 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; } @@ -291,8 +297,7 @@ impl PostgresWorker { truncate_ellipse(v, MAX_RECORD_LEN_IN_ERRMSG, "...") )); - // Either way the transaction is gone, along with everything it held. - // Whether it can be written again is what separates the two. + // Transaction state is decided by whether the error is retry-able. self.transaction_state = if error.should_retry() { TransactionState::Lost } else { @@ -448,10 +453,7 @@ These statements were successfully prepared before reconnecting. Does the table Ok(_) => return, Err(e) => { // Only a failure that ends the loop leaves the endpoint - // unable to reach postgres. Reporting one it is about to - // retry as fatal stamps `fatal_error`, which never clears, - // so a blip the connector recovers from in a second would - // mark the endpoint failed for the life of the pipeline. + // unable to reach postgres. let retry = e.should_retry(); if let Some(controller) = self.controller.upgrade() { controller.output_transport_error( @@ -479,13 +481,6 @@ These statements were successfully prepared before reconnecting. Does the table } /// Whether the pipeline has been asked to stop. - /// - /// Waiting for the controller to be dropped would never work here: the - /// endpoint's drop joins this thread, and the endpoint itself is what the - /// output thread is blocked on, so the controller stays alive precisely - /// because this thread has not finished. `Terminated`, which - /// `ControllerInner::stop` sets before it touches anything else, is the - /// signal that does arrive. fn shutting_down(&self) -> bool { match self.controller.upgrade() { None => true, @@ -495,10 +490,6 @@ These statements were successfully prepared before reconnecting. Does the table /// 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. - /// - /// The wait reaches a minute, and the endpoint's drop waits on this thread, - /// so sleeping through it in one piece would delay every pipeline shutdown - /// that happens while postgres is unreachable. fn backoff_unless_shutdown(&self, duration: Duration) -> bool { const SLICE: Duration = Duration::from_millis(100); @@ -539,9 +530,11 @@ These statements were successfully prepared before reconnecting. Does the table self.transaction = Some(transaction); // Start counting afresh: whatever the previous transaction wrote was - // either reported at its commit or lost with it. + // 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(()) @@ -579,23 +572,13 @@ These statements were successfully prepared before reconnecting. Does the table Ok(()) => String::new(), Err(e) => format!(" Rolling the transaction back also failed: {e}."), }; - // Spell out that the rows PostgreSQL accepted are gone too. A reader - // who saw only the rejected statement would expect the rest of the - // rows to have landed. - let accepted = if self.num_rows > 0 { - format!( - ": the {} row(s) PostgreSQL had already accepted were rolled back, \ - and the rest were dropped", - self.num_rows - ) - } else { - String::new() - }; return Err(BackoffError::Permanent(anyhow!( "PostgreSQL rejected one of the statements writing to table {:?} and \ - aborted the transaction, so no rows were written{accepted}. The \ - rejected statement is reported as a separate error.{rollback}", + 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, ))); } @@ -768,6 +751,9 @@ These statements were successfully prepared before reconnecting. Does the table self.upsert(buf); } }; + + // One row buffered, counted whether or not its statement runs. + self.batch_rows += 1; } cursor.step_key(); diff --git a/crates/adapters/src/integrated/postgres/test.rs b/crates/adapters/src/integrated/postgres/test.rs index 32f5ebfe0b9..8bc5daf9129 100644 --- a/crates/adapters/src/integrated/postgres/test.rs +++ b/crates/adapters/src/integrated/postgres/test.rs @@ -1139,7 +1139,7 @@ fn test_pg_error_message_is_bounded() { // `freshness_timestamp` is NOT NULL, has no DEFAULT, and is absent from the // Feldera view, so every INSERT the connector issues violates the - // constraint. This is the customer-reported failure. + // constraint. let _table = PostgresTestStruct::create_table_with_extra_columns( &table_name, url, @@ -1185,9 +1185,6 @@ fn test_pg_error_message_is_bounded() { assert!(!errors.is_empty(), "no transport error was recorded"); for error in &errors { - // The bound enforced by the status layer is a backstop; the connector is - // expected to quote no more than a prefix of the payload on its own. - // "bytes elided" means the backstop had to fire. assert!( !error.message.contains("bytes elided"), "connector produced a {}-byte message that the status layer had to \ @@ -1219,7 +1216,7 @@ fn test_pg_error_message_is_bounded() { "error drops the server's DETAIL: {diagnosed}" ); - // Guard against the test going vacuous: at least one failed statement has + // At least one failed statement has // to have carried a payload far larger than the message describing it. let largest_payload = errors .iter() @@ -1240,7 +1237,7 @@ fn test_pg_error_message_is_bounded() { ); } -/// Serializes `data` as the newline-delimited JSON the file input reads. +/// 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 { @@ -1260,9 +1257,8 @@ fn write_records_to_temp_file(data: &[PostgresTestStruct]) -> NamedTempFile { /// 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 step once over that -/// many records so a test can act on one large batch, and has the file input -/// read the whole file at once to match. +/// `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, @@ -1338,8 +1334,7 @@ fn pg_row_count(url: &str, table_name: &str) -> u64 { /// /// 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. Counting rows as they were encoded -/// therefore reported a batch as transmitted that the server had discarded. +/// success proves nothing on its own. #[test] #[serial] fn test_pg_transmitted_records_excludes_rolled_back_rows() { @@ -1365,7 +1360,7 @@ fn test_pg_transmitted_records_excludes_rolled_back_rows() { let (controller, _err_receiver) = PostgresTestStruct::test_circuit(config); controller.start(); - // Wait until the connector has finished handling every record. This counter + // 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. @@ -1407,11 +1402,6 @@ fn records_in_error(message: &str) -> Option { /// A transaction that fails partway discards the statements that already /// succeeded, so none of the batch counts as transmitted. -/// -/// This is the case the commit guard exists for. Counting rows per successful -/// statement is not enough on its own: the rows of the earlier statement would -/// still be reported, because the failure aborts the transaction and the -/// following COMMIT silently rolls them back. #[test] #[serial] fn test_pg_transmitted_records_excludes_rows_lost_to_a_later_failure() { @@ -1423,7 +1413,7 @@ fn test_pg_transmitted_records_excludes_rows_lost_to_a_later_failure() { let verify_url = url.clone(); // Enough records to fill the connector's 1 MiB buffer more than once, so the - // batch becomes several statements. Keys ascend and the cursor walks them in + // 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(); @@ -1518,18 +1508,38 @@ fn test_pg_transmitted_records_excludes_rows_lost_to_a_later_failure() { ); } -/// Selects the connector backends that are inside a batch transaction. +/// 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. /// -/// Requiring an open transaction whose latest query is a connector statement -/// keeps this off the connection while it prepares its statements, and widens -/// the window to the whole batch: Postgres keeps reporting the last query while -/// the backend sits idle in the 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 until a connector backend is inside a batch transaction. +/// 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 @@ -1547,7 +1557,8 @@ fn wait_for_batch_transaction(client: &mut postgres::Client) -> bool { false } -/// Terminates every connector backend that is inside a batch transaction. +/// 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( @@ -1562,8 +1573,12 @@ fn terminate_batch_backends(client: &mut postgres::Client) -> bool { terminated > 0 } -/// Drops the connector's connections mid-batch, `rounds` times, and reports how -/// many rounds landed. +/// 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; @@ -1583,11 +1598,12 @@ fn terminate_connector_backends(url: &str, rounds: usize) -> usize { /// 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 -/// batch the worker retained. +/// cursor the worker retained. /// -/// `threads` covers the connector writing through one connection and through -/// several, which partition the batch. `kills` covers the connection dropping -/// again while the batch is being rewritten. +/// `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; @@ -1612,13 +1628,15 @@ fn pg_reconnect_test(name: &str, threads: usize, kills: usize) { // 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 handled whether or not it wrote it, so + // 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, @@ -1651,8 +1669,7 @@ fn pg_reconnect_test(name: &str, threads: usize, kills: usize) { records, RECORDS as u64, "reported {records} of {RECORDS} rows as transmitted" ); - // `fatal_error` never clears, so recording one for a drop the connector - // recovered from would leave the endpoint marked failed for good. + assert_eq!( fatal_error, None, "the connector recovered, yet the endpoint is marked as fatally failed" @@ -1681,8 +1698,7 @@ fn test_pg_reconnect_survives_repeated_drops() { /// 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. Restoring -/// on drop matters: leaving the database closed would fail every later test. +/// it stands in for a server that is restarting or under maintenance. struct RefusedConnections { admin: postgres::Client, database: String, @@ -1731,6 +1747,7 @@ impl RefusedConnections { } } +/// Restoring on drop matters: leaving the database closed would fail every later test. impl Drop for RefusedConnections { fn drop(&mut self) { if !self.refusing { @@ -1745,8 +1762,8 @@ impl Drop for RefusedConnections { /// A database that refuses the reconnect must be waited out, not given up on. /// -/// This is the one path that dropping a backend does not reach: there the -/// reconnect succeeds on the first try, so `retry_connecting_with_backoff` never +/// 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] @@ -1816,7 +1833,7 @@ fn test_pg_reconnect_waits_out_a_refusing_database() { so this test proves nothing" ); // Only `retry_connecting_with_backoff` reports under this tag, so its error - // branch ran, which is the whole point of refusing the connections. + // branch ran. assert!( errors .iter() @@ -1843,8 +1860,7 @@ fn test_pg_reconnect_waits_out_a_refusing_database() { /// 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. Holding a strong reference to the controller across the loop did -/// exactly that, since it kept alive the very thing the loop was watching for. +/// output. #[test] #[serial] fn test_pg_stops_while_the_database_is_unreachable() { @@ -1899,7 +1915,7 @@ fn test_pg_stops_while_the_database_is_unreachable() { ); } -/// The counter still reports every row of a batch that commits. +/// The `transmitted_records` counter still reports every row of a batch that commits. #[test] #[serial] fn test_pg_transmitted_records_counts_committed_rows() {