diff --git a/crates/etl-postgres/src/application_name.rs b/crates/etl-postgres/src/application_name.rs new file mode 100644 index 000000000..ba72848c9 --- /dev/null +++ b/crates/etl-postgres/src/application_name.rs @@ -0,0 +1,139 @@ +//! Naming helpers for per-worker Postgres `application_name` values. +//! +//! Worker connections are tagged as `{base}:apply:{pipeline_id}` and +//! `{base}:table_sync:{pipeline_id}:{table_oid}` so they can be identified in +//! `pg_stat_activity` and targeted individually (e.g. by +//! `pg_terminate_backend` in tests). Connections not owned by a worker keep +//! the plain base name. + +use tracing::debug; + +use crate::schema::TableId; + +/// Maximum length of a Postgres `application_name` in bytes. +/// +/// Postgres silently truncates longer values, so names are clamped locally to +/// keep the worker suffix intact. +const MAX_APPLICATION_NAME_LENGTH: usize = 63; + +/// Separator between the base application name and worker suffix parts. +/// +/// `:` keeps suffixes unambiguous to parse, since base names contain `_`. +const SEPARATOR: char = ':'; + +/// Tag identifying apply worker connections. +pub const APPLY_WORKER_TAG: &str = "apply"; +/// Tag identifying table sync worker connections. +pub const TABLE_SYNC_WORKER_TAG: &str = "table_sync"; + +/// Returns the `application_name` for an apply worker connection. +pub fn apply_worker_application_name(base: &str, pipeline_id: u64) -> String { + with_worker_suffix(base, &format!("{SEPARATOR}{APPLY_WORKER_TAG}{SEPARATOR}{pipeline_id}")) +} + +/// Returns the `application_name` for a table sync worker connection. +pub fn table_sync_worker_application_name( + base: &str, + pipeline_id: u64, + table_id: TableId, +) -> String { + with_worker_suffix( + base, + &format!( + "{SEPARATOR}{TABLE_SYNC_WORKER_TAG}{SEPARATOR}{pipeline_id}{SEPARATOR}{}", + table_id.into_inner() + ), + ) +} + +/// Appends a worker suffix to `base`, clamping `base` so the result fits the +/// Postgres limit with the suffix intact. +/// +/// The suffix carries the worker identity that tests and operators match on, +/// so on overflow the base is truncated instead of letting Postgres cut the +/// tail. In-repo bases fit without clamping while pipeline id and table oid +/// stay within 15 digits combined; beyond that the suffix still survives. +fn with_worker_suffix(base: &str, suffix: &str) -> String { + let max_allowed_len = MAX_APPLICATION_NAME_LENGTH.saturating_sub(suffix.len()); + if base.len() > max_allowed_len { + debug!(base, suffix, "application_name base clamped to fit worker suffix"); + } + + let end = base.floor_char_boundary(max_allowed_len); + + format!("{}{suffix}", &base[..end]) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn apply_worker_name_appends_tag_and_pipeline_id() { + // --- GIVEN: a base name and pipeline id --- + let name = apply_worker_application_name("supabase_etl_replicator_replication", 42); + + // --- THEN: the name carries the apply tag and pipeline id --- + assert_eq!(name, "supabase_etl_replicator_replication:apply:42"); + } + + #[test] + fn table_sync_worker_name_appends_tag_pipeline_id_and_table_oid() { + // --- GIVEN: a base name, pipeline id, and table id --- + let name = table_sync_worker_application_name( + "supabase_etl_replicator_replication", + 42, + TableId::new(16384), + ); + + // --- THEN: the name carries the table_sync tag, pipeline id, and table oid --- + assert_eq!(name, "supabase_etl_replicator_replication:table_sync:42:16384"); + } + + #[test] + fn name_fits_without_clamping_up_to_15_combined_id_digits() { + // --- GIVEN: the longest in-repo base with pipeline id and table oid summing to + // 15 digits, the documented no-clamp bound --- + let base = "supabase_etl_replicator_replication"; + let name = table_sync_worker_application_name(base, 99_999, TableId::new(u32::MAX)); + + // --- THEN: the name fits the Postgres limit with the base intact --- + assert!(name.len() <= MAX_APPLICATION_NAME_LENGTH); + assert!(name.starts_with(base)); + } + + #[test] + fn name_beyond_no_clamp_bound_keeps_suffix_and_prefix_filterable_base() { + // --- GIVEN: the longest in-repo base with ids one digit past the no-clamp + // bound --- + let base = "supabase_etl_replicator_replication"; + let name = table_sync_worker_application_name(base, 999_999, TableId::new(u32::MAX)); + + // --- THEN: the base clamps but the suffix and the etl prefix survive --- + assert_eq!(name.len(), MAX_APPLICATION_NAME_LENGTH); + assert!(name.ends_with(&format!(":table_sync:999999:{}", u32::MAX))); + assert!(name.starts_with("supabase_etl_")); + } + + #[test] + fn overlong_base_is_clamped_with_suffix_preserved() { + // --- GIVEN: a base longer than the Postgres limit allows --- + let base = "x".repeat(80); + let name = table_sync_worker_application_name(&base, u64::MAX, TableId::new(u32::MAX)); + + // --- THEN: the name fits and the full worker suffix survives --- + assert_eq!(name.len(), MAX_APPLICATION_NAME_LENGTH); + assert!(name.ends_with(&format!(":table_sync:{}:{}", u64::MAX, u32::MAX))); + } + + #[test] + fn clamping_respects_multibyte_char_boundaries() { + // --- GIVEN: an overlong base of multibyte characters --- + let base = "é".repeat(60); + let name = apply_worker_application_name(&base, u64::MAX); + + // --- THEN: clamping does not split a character and the suffix survives --- + assert!(name.len() <= MAX_APPLICATION_NAME_LENGTH); + assert!(name.ends_with(&format!(":apply:{}", u64::MAX))); + } +} diff --git a/crates/etl-postgres/src/lib.rs b/crates/etl-postgres/src/lib.rs index f6254061e..b618a5879 100644 --- a/crates/etl-postgres/src/lib.rs +++ b/crates/etl-postgres/src/lib.rs @@ -4,6 +4,7 @@ //! crates: source database access, replication slot naming, ETL metadata-store //! queries, schema primitives, value wrappers, and type conversion helpers. +pub mod application_name; pub mod default_expression; pub mod lag; pub mod numeric; diff --git a/crates/etl/src/postgres/client/raw.rs b/crates/etl/src/postgres/client/raw.rs index 995edb18f..c735b1488 100644 --- a/crates/etl/src/postgres/client/raw.rs +++ b/crates/etl/src/postgres/client/raw.rs @@ -1,11 +1,10 @@ -use std::{ - fmt, - num::NonZeroI32, - sync::{Arc, LazyLock}, - time::Duration, -}; +use std::{fmt, num::NonZeroI32, sync::Arc, time::Duration}; -use etl_postgres::{source::extract_server_version, tokio::tls::MakeRustlsConnect}; +use etl_postgres::{ + application_name::{apply_worker_application_name, table_sync_worker_application_name}, + source::extract_server_version, + tokio::tls::MakeRustlsConnect, +}; use pg_escape::{quote_identifier, quote_literal}; use postgres_replication::LogicalReplicationStream; use rustls::{ @@ -34,6 +33,7 @@ use crate::{ config::{IntoConnectOptions, PgConnectionConfig, PgConnectionOptions}, error::{ErrorKind, EtlResult}, etl_error, + pipeline::PipelineId, schema::{TableId, TableName}, }; @@ -44,21 +44,24 @@ use crate::{ const DELETE_SLOT_TIMEOUT: Duration = Duration::from_secs(30); /// Default duration unit used when `pg_settings.unit` is empty. const PG_SETTINGS_DEFAULT_DURATION_UNIT: &str = "ms"; -/// Application name for ETL logical replication connections. -const APP_NAME_REPLICATOR_REPLICATION: &str = "supabase_etl_replicator_replication"; +/// Base application name for ETL logical replication connections. +/// +/// Worker connections append a suffix (see [`etl_postgres::application_name`]) +/// so they can be told apart in `pg_stat_activity`. +pub const APP_NAME_REPLICATOR_REPLICATION: &str = "supabase_etl_replicator_replication"; -/// Connection options for logical replication. +/// Builds connection options for logical replication connections. /// /// Disables statement, lock, and idle-in-transaction timeouts because /// replication streams, slot creation, and initial table synchronization can /// legitimately run for a long time. -static REPLICATION_OPTIONS: LazyLock = LazyLock::new(|| { - PgConnectionOptions::builder(APP_NAME_REPLICATOR_REPLICATION) +fn replication_options(application_name: String) -> PgConnectionOptions { + PgConnectionOptions::builder(application_name) .statement_timeout(0) .lock_timeout(0) .idle_in_transaction_session_timeout(0) .build() -}); +} /// The kind of PostgreSQL connection to create. #[derive(Debug, Clone, Copy)] @@ -99,13 +102,21 @@ fn tls_client_config_from_root_certs(trusted_root_certs: &str) -> EtlResult, + /// Session options, including the per-worker `application_name`. + /// + /// Child connections reuse these options, so they inherit the parent + /// worker's `application_name`. + options: Arc, /// Parsed rustls client config, present when TLS is enabled. tls_client_config: Option>, } impl PgReplicationConnectionConfig { /// Creates shared connection settings from an owned PostgreSQL config. - fn new(pg_connection_config: PgConnectionConfig) -> EtlResult { + fn new( + pg_connection_config: PgConnectionConfig, + options: PgConnectionOptions, + ) -> EtlResult { let tls_client_config = if pg_connection_config.tls.enabled { Some(Arc::new(tls_client_config_from_root_certs( &pg_connection_config.tls.trusted_root_certs, @@ -114,7 +125,11 @@ impl PgReplicationConnectionConfig { None }; - Ok(Self { pg_connection_config: Arc::new(pg_connection_config), tls_client_config }) + Ok(Self { + pg_connection_config: Arc::new(pg_connection_config), + options: Arc::new(options), + tls_client_config, + }) } /// Returns the raw PostgreSQL connection settings. @@ -122,6 +137,11 @@ impl PgReplicationConnectionConfig { &self.pg_connection_config } + /// Returns the session options for connections created from this config. + fn options(&self) -> &PgConnectionOptions { + &self.options + } + /// Returns the shared rustls client config. fn tls_client_config(&self) -> Option> { self.tls_client_config.as_ref().map(Arc::clone) @@ -213,9 +233,55 @@ impl PgReplicationClient { /// Establishes a connection to Postgres. The connection uses TLS if /// configured in the supplied [`PgConnectionConfig`]. /// - /// The connection is configured for logical replication mode. + /// The connection is configured for logical replication mode and uses the + /// base replication `application_name`. Worker connections should use + /// [`PgReplicationClient::connect_for_apply_worker`] or + /// [`PgReplicationClient::connect_for_table_sync_worker`] instead. pub async fn connect(pg_connection_config: PgConnectionConfig) -> EtlResult { - let connection_config = PgReplicationConnectionConfig::new(pg_connection_config)?; + Self::connect_with_application_name( + pg_connection_config, + APP_NAME_REPLICATOR_REPLICATION.to_owned(), + ) + .await + } + + /// Establishes a replication connection tagged for an apply worker. + pub async fn connect_for_apply_worker( + pg_connection_config: PgConnectionConfig, + pipeline_id: PipelineId, + ) -> EtlResult { + Self::connect_with_application_name( + pg_connection_config, + apply_worker_application_name(APP_NAME_REPLICATOR_REPLICATION, pipeline_id), + ) + .await + } + + /// Establishes a replication connection tagged for a table sync worker. + pub async fn connect_for_table_sync_worker( + pg_connection_config: PgConnectionConfig, + pipeline_id: PipelineId, + table_id: TableId, + ) -> EtlResult { + Self::connect_with_application_name( + pg_connection_config, + table_sync_worker_application_name( + APP_NAME_REPLICATOR_REPLICATION, + pipeline_id, + table_id, + ), + ) + .await + } + + /// Establishes a replication connection with the supplied + /// `application_name`. + async fn connect_with_application_name( + pg_connection_config: PgConnectionConfig, + application_name: String, + ) -> EtlResult { + let options = replication_options(application_name); + let connection_config = PgReplicationConnectionConfig::new(pg_connection_config, options)?; PgReplicationClient::connect_with_kind(connection_config, ConnectionKind::Replication).await } @@ -237,7 +303,7 @@ impl PgReplicationClient { kind: ConnectionKind, ) -> EtlResult { let mut config: Config = - connection_config.pg_connection_config().with_db(Some(&REPLICATION_OPTIONS)); + connection_config.pg_connection_config().with_db(Some(connection_config.options())); kind.configure(&mut config); let (client, connection) = config.connect(NoTls).await?; @@ -258,7 +324,7 @@ impl PgReplicationClient { kind: ConnectionKind, ) -> EtlResult { let mut config: Config = - connection_config.pg_connection_config().with_db(Some(&REPLICATION_OPTIONS)); + connection_config.pg_connection_config().with_db(Some(connection_config.options())); kind.configure(&mut config); let tls_config = connection_config.tls_client_config().ok_or_else(|| { @@ -694,3 +760,23 @@ impl PgReplicationClient { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn replication_base_name_fits_worker_suffix_without_clamping() { + // --- GIVEN: a worker suffix with pipeline id and table oid summing to 15 + // digits, the documented no-clamp bound for the replication base name --- + let name = table_sync_worker_application_name( + APP_NAME_REPLICATOR_REPLICATION, + 99_999, + TableId::new(u32::MAX), + ); + + // --- THEN: the name fits Postgres's 63-byte limit with the base intact --- + assert!(name.len() <= 63, "worker application_name '{name}' exceeds Postgres limit"); + assert!(name.starts_with(APP_NAME_REPLICATOR_REPLICATION)); + } +} diff --git a/crates/etl/src/runtime/apply/worker.rs b/crates/etl/src/runtime/apply/worker.rs index 35142e571..50deaec2c 100644 --- a/crates/etl/src/runtime/apply/worker.rs +++ b/crates/etl/src/runtime/apply/worker.rs @@ -283,8 +283,11 @@ where /// Runs a single apply worker attempt. async fn run_apply_worker(self) -> EtlResult<()> { - let replication_client = - PgReplicationClient::connect(self.config.pg_connection.clone()).await?; + let replication_client = PgReplicationClient::connect_for_apply_worker( + self.config.pg_connection.clone(), + self.pipeline_id, + ) + .await?; let start_lsn = get_start_lsn( self.pipeline_id, diff --git a/crates/etl/src/runtime/table_sync/worker.rs b/crates/etl/src/runtime/table_sync/worker.rs index 23e05da76..da798599a 100644 --- a/crates/etl/src/runtime/table_sync/worker.rs +++ b/crates/etl/src/runtime/table_sync/worker.rs @@ -729,8 +729,12 @@ where // Note that this connection must be tied to the lifetime of this worker, // otherwise there will be problems when cleaning up the replication // slot. - let mut replication_client = - PgReplicationClient::connect(self.config.pg_connection.clone()).await?; + let mut replication_client = PgReplicationClient::connect_for_table_sync_worker( + self.config.pg_connection.clone(), + self.pipeline_id, + self.table_id, + ) + .await?; let table_sync_result = start_table_sync( self.pipeline_id, diff --git a/crates/etl/tests/pipeline_with_failpoints.rs b/crates/etl/tests/pipeline_with_failpoints.rs index 9c4038424..f5846c424 100644 --- a/crates/etl/tests/pipeline_with_failpoints.rs +++ b/crates/etl/tests/pipeline_with_failpoints.rs @@ -1,5 +1,7 @@ #![cfg(all(feature = "test-utils", feature = "failpoints"))] +use std::time::Duration; + use etl::{ event::{Event, EventType, InsertEvent}, failpoints::{ @@ -27,7 +29,12 @@ use etl::{ }, }, }; -use etl_postgres::{below_version, tokio::test_utils::TableModification, version::POSTGRES_15}; +use etl_postgres::{ + application_name::{apply_worker_application_name, table_sync_worker_application_name}, + below_version, + tokio::test_utils::TableModification, + version::POSTGRES_15, +}; use etl_telemetry::tracing::init_test_tracing; use fail::FailScenario; use rand::random; @@ -1545,3 +1552,82 @@ async fn table_schema_replication_masks_are_consistent_after_restart() { pipeline.shutdown_and_wait().await.unwrap(); } + +#[tokio::test(flavor = "multi_thread")] +async fn worker_connections_are_tagged_with_per_worker_application_names() { + let _scenario = FailScenario::setup(); + fail::cfg(START_TABLE_SYNC_AFTER_FINISHED_COPY_FP, "pause").unwrap(); + + init_test_tracing(); + + // --- GIVEN: a pipeline whose table sync worker is paused after copy, so both + // worker connections are alive --- + let mut database = spawn_source_database().await; + let database_schema = setup_test_database_schema(&database, TableSelection::UsersOnly).await; + let table_id = database_schema.users_schema().id; + + insert_users_data(&mut database, &database_schema.users_schema().name, 1..=1).await; + + let store = NotifyingStore::new(); + let destination = TestDestinationWrapper::wrap(MemoryDestination::new(store.clone())); + + let pipeline_id: PipelineId = random(); + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + database_schema.publication_name(), + store.clone(), + destination.clone(), + ); + + let finished_copy_notify = + store.notify_on_table_state_type(table_id, TableStateType::FinishedCopy).await; + + pipeline.start().await.unwrap(); + + finished_copy_notify.notified().await; + + // --- THEN: pg_stat_activity shows both connections under their per-worker + // application names --- + // + // Names are computed with the same helpers used by the workers because the + // random test pipeline ids can be long enough to clamp the base name. + let apply_name = + apply_worker_application_name("supabase_etl_replicator_replication", pipeline_id); + let table_sync_name = table_sync_worker_application_name( + "supabase_etl_replicator_replication", + pipeline_id, + table_id, + ); + + let client = database.client.as_ref().unwrap(); + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let rows = client + .query( + "select application_name from pg_stat_activity where datname = \ + current_database() and application_name in ($1, $2)", + &[&apply_name, &table_sync_name], + ) + .await + .unwrap(); + let names: Vec = rows.iter().map(|row| row.get(0)).collect(); + + if names.contains(&apply_name) && names.contains(&table_sync_name) { + break; + } + + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await + .expect("timed out waiting for per-worker application names in pg_stat_activity"); + + let ready_notify = store.notify_on_table_state_type(table_id, TableStateType::Ready).await; + + fail::remove(START_TABLE_SYNC_AFTER_FINISHED_COPY_FP); + + ready_notify.notified().await; + + pipeline.shutdown_and_wait().await.unwrap(); +}