diff --git a/integration/two_pc_crash_safety/wal_helper/Cargo.toml b/integration/two_pc_crash_safety/wal_helper/Cargo.toml index 9ec8096b4..57c73732a 100644 --- a/integration/two_pc_crash_safety/wal_helper/Cargo.toml +++ b/integration/two_pc_crash_safety/wal_helper/Cargo.toml @@ -10,5 +10,5 @@ path = "src/main.rs" [dependencies] bytes = "1" -pgdog = { path = "../../../pgdog", default-features = false } +pgdog = { path = "../../../pgdog" } tokio = { version = "1", features = ["macros", "rt-multi-thread"] } diff --git a/integration/two_pc_crash_safety/wal_helper/src/main.rs b/integration/two_pc_crash_safety/wal_helper/src/main.rs index 2481fe219..1b3396bc3 100644 --- a/integration/two_pc_crash_safety/wal_helper/src/main.rs +++ b/integration/two_pc_crash_safety/wal_helper/src/main.rs @@ -33,10 +33,13 @@ async fn main() { let mut segment = Segment::create(&dir, 0).await.expect("create segment"); let mut buf = BytesMut::new(); + // The prefix is left empty so recovery takes the numeric-ID + // fallback path, keeping that path covered end to end. Record::Begin(BeginPayload { txn, user, database, + prefix: String::new(), }) .encode(&mut buf) .expect("encode begin"); diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index a762b248d..d8deb68f2 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -5,14 +5,16 @@ use crate::{ ClientRequest, client::query_engine::{ TwoPcPhase, - two_pc::{TwoPcTransaction, statement::phase_control}, + two_pc::{TwoPcTransaction, TwoPcTransactionOnShard, statement::phase_control}, }, }, - net::{FrontendPid, ProtocolMessage, Query, parameter::Parameters}, + net::{DataRow, FrontendPid, ProtocolMessage, Query, parameter::Parameters}, state::State, + util::is_safe_identifier, }; use futures::future::join_all; +use tracing::warn; use super::*; @@ -364,20 +366,25 @@ impl Binding { pub(crate) async fn two_pc_on_guards( servers: &mut [Guard], + state: &MultiShard, transaction: TwoPcTransaction, phase: TwoPcPhase, ) -> Result<(), Error> { let skip_missing = matches!(phase, TwoPcPhase::Phase2 | TwoPcPhase::Rollback); let mut futures = Vec::new(); - for (shard, server) in servers.iter_mut().enumerate() { + for (position, server) in servers.iter_mut().enumerate() { + // Map positional index to actual shard number. + // When only a subset of shards is connected (Shard::Multi binding), + // positional indices don't match actual shard numbers. + let shard = state.shard_index(position); let query = phase_control(transaction, shard, phase); futures.push(server.execute(query)); } let results = join_all(futures).await; - for (shard, result) in results.into_iter().enumerate() { + for (position, result) in results.into_iter().enumerate() { match result { Err(Error::ExecutionError(err)) => { if !(skip_missing && err.code == "42704") { @@ -387,7 +394,7 @@ impl Binding { Err(err) => return Err(err), Ok(_) => { if phase == TwoPcPhase::Phase2 { - servers[shard].stats_mut().transaction_2pc(); + servers[position].stats_mut().transaction_2pc(); } } } @@ -401,10 +408,116 @@ impl Binding { &mut self, transaction: TwoPcTransaction, phase: TwoPcPhase, + ) -> Result<(), Error> { + match self { + Binding::MultiShard(servers, state) => { + Self::two_pc_on_guards(servers, state, transaction, phase).await + } + + _ => Err(Error::TwoPcMultiShardOnly), + } + } + + /// Resolve a two-phase transaction on every shard during cleanup + /// and crash recovery. + /// + /// `prefix` is the coordinator GID prefix recorded when the + /// transaction was created; combined with the transaction's numeric + /// ID it names the exact prepared transaction on each shard, and the + /// phase statement runs against the GIDs found in + /// `pg_prepared_xacts`. An empty `prefix` means the transaction was + /// restored from a WAL record that did not store it; matching then + /// falls back to the durable numeric transaction ID and shard index. + /// A shard with no matching GID is already resolved and is skipped. + pub(crate) async fn two_pc_cleanup( + &mut self, + transaction: TwoPcTransaction, + prefix: &str, + phase: TwoPcPhase, ) -> Result<(), Error> { match self { Binding::MultiShard(servers, _) => { - Self::two_pc_on_guards(servers, transaction, phase).await + let futures = servers + .iter_mut() + .enumerate() + .map(|(shard, server)| async move { + let target = TwoPcTransactionOnShard::new(transaction, shard); + // A GID rendered from the recorded prefix is + // checked against the identifier alphabet before + // it is embedded in a quoted literal; a prefix + // outside the alphabet falls back to numeric-ID + // matching, which performs the same check. + let expected = if prefix.is_empty() { + None + } else { + let gid = target.gid(prefix); + if is_safe_identifier(&gid) { + Some(gid) + } else { + warn!( + "[2pc] recorded gid {:?} contains characters \ + PgDog never generates; matching by numeric ID", + gid + ); + None + } + }; + let rows: Vec = server + .fetch_all( + "SELECT gid FROM pg_prepared_xacts WHERE database = current_database()", + ) + .await?; + + for row in rows { + let Some(gid) = row.get_text(0) else { + continue; + }; + let matched = match &expected { + Some(expected) => &gid == expected, + None => target.matches_gid(&gid), + }; + if !matched { + continue; + } + // Both match paths guarantee the gid contains + // no characters that need quoting. + let statement = match phase { + TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{gid}'"), + TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{gid}'"), + TwoPcPhase::Phase1 => { + unreachable!("cleanup resolves transactions; it never prepares") + } + }; + match server.execute(&statement).await { + Ok(_) => { + if phase == TwoPcPhase::Phase2 { + server.stats_mut().transaction_2pc(); + } + } + // Insufficient privilege: the prepared transaction + // is owned by a role the configured user can't act + // for. Retrying can't succeed until an operator + // intervenes, so leave the transaction to them + // rather than retrying forever. + Err(Error::ExecutionError(err)) if err.code == "42501" => { + warn!( + "[2pc] insufficient privilege to run {}; \ + resolve the prepared transaction manually", + statement + ); + } + Err(err) => return Err(err), + } + } + + Ok::<(), Error>(()) + }); + + for result in join_all(futures).await { + result?; + } + + Ok(()) } _ => Err(Error::TwoPcMultiShardOnly), diff --git a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs index ff22ffb23..914d7b5f8 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs @@ -205,12 +205,14 @@ impl Manager { identifier: &Arc, phase: TwoPcPhase, ) -> Result { + let prefix = TwoPcTransaction::global_prefix(); let prior = { let mut guard = self.inner.lock(); let prior = guard.transactions.get(&transaction).cloned(); let entry = guard.transactions.entry(transaction).or_default(); entry.identifier = identifier.clone(); entry.phase = phase; + entry.prefix = prefix.clone(); prior }; @@ -221,6 +223,7 @@ impl Manager { transaction, identifier.user.clone(), identifier.database.clone(), + prefix, ) .await } @@ -262,14 +265,20 @@ impl Manager { transaction: TwoPcTransaction, user: String, database: String, + prefix: String, phase: TwoPcPhase, ) { let identifier = Arc::new(User { user, database }); { let mut guard = self.inner.lock(); - guard - .transactions - .insert(transaction, TransactionInfo { phase, identifier }); + guard.transactions.insert( + transaction, + TransactionInfo { + phase, + identifier, + prefix, + }, + ); guard.queue.push_back(transaction); } self.stats.incr_recovered(); @@ -345,6 +354,13 @@ impl Manager { } /// Reconnect to cluster if available and rollback the two-phase transaction. + /// + /// The transaction may have been created by an earlier PgDog + /// process; the GID prefix recorded when it was created names the + /// exact prepared transactions on the shards. Cleanup scans each + /// shard's prepared transactions and acts on the GIDs that match, + /// falling back to the durable numeric transaction ID for WAL + /// records that did not store the prefix. async fn cleanup_phase(&self, transaction: TwoPcTransaction) -> Result<(), Error> { let state = match self.inner.lock().transactions.get(&transaction).cloned() { Some(state) => state, @@ -379,7 +395,9 @@ impl Manager { &Route::write(ShardWithPriority::new_override_transaction(Shard::All)), ) .await?; - connection.two_pc(transaction, phase).await?; + connection + .two_pc_cleanup(transaction, &state.prefix, phase) + .await?; connection.disconnect(); Ok(()) @@ -412,6 +430,10 @@ impl Manager { pub struct TransactionInfo { pub phase: TwoPcPhase, pub identifier: Arc, + /// Coordinator GID prefix the transaction's prepared names were + /// rendered with. Empty when restored from a WAL record written by + /// a version that did not store it. + pub prefix: String, } #[derive(Default, Debug)] diff --git a/pgdog/src/frontend/client/query_engine/two_pc/statement.rs b/pgdog/src/frontend/client/query_engine/two_pc/statement.rs index 9c59824e4..6ca310561 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/statement.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/statement.rs @@ -2,7 +2,10 @@ use std::{fmt::Display, str::FromStr}; +use tracing::warn; + use crate::frontend::client::query_engine::two_pc::TwoPcTransaction; +use crate::util::is_safe_identifier; use super::TwoPcPhase; @@ -24,6 +27,47 @@ impl TwoPcTransactionOnShard { pub(crate) fn transaction(&self) -> TwoPcTransaction { self.transaction } + + /// The GID this transaction is prepared under on this shard when + /// rendered with `prefix`, the coordinator GID prefix recorded at + /// transaction creation. + pub(crate) fn gid(&self, prefix: &str) -> String { + format!("{}{}_{}", prefix, self.transaction.number(), self.shard) + } + + /// Whether `gid`, as listed in `pg_prepared_xacts`, refers to this + /// transaction on this shard. GID prefixes embed the identity of the + /// PgDog process that created the transaction, which can change + /// across restarts, so matching uses the durable numeric transaction + /// ID and the shard index. This is the fallback for transactions + /// restored from WAL records that did not store the coordinator GID + /// prefix, and for recorded prefixes that fail the identifier + /// alphabet check; with a recorded prefix, cleanup matches the exact + /// GID from [`Self::gid`] instead. + /// + /// A matched GID is guaranteed to contain only [`is_safe_identifier`] + /// characters, so it can be embedded verbatim in a quoted SQL literal. + /// PgDog only generates such GIDs (`NODE_ID` and `DEPLOYMENT_ID` are + /// validated at startup); a GID that matches the numeric ID but not + /// the alphabet was not created by PgDog and is refused with a + /// warning. + pub(crate) fn matches_gid(&self, gid: &str) -> bool { + let matches = match gid.parse::() { + Ok(parsed) => parsed.transaction == self.transaction && parsed.shard == self.shard, + Err(()) => false, + }; + + if matches && !is_safe_identifier(gid) { + warn!( + "[2pc] prepared transaction {:?} matches transaction {} on shard {} \ + but contains characters PgDog never generates; refusing to resolve it", + gid, self.transaction, self.shard + ); + return false; + } + + matches + } } impl Display for TwoPcTransactionOnShard { @@ -101,6 +145,47 @@ mod test { ); } + #[test] + fn matches_gid_ignores_prefix() { + let transaction: TwoPcTransaction = "__pgdog_2pc_123".parse().unwrap(); + let target = TwoPcTransactionOnShard::new(transaction, 1); + + // Rendered by this process. + assert!(target.matches_gid(&target.to_string())); + // Rendered by a process with a different instance ID. + assert!(target.matches_gid("__pgdog_2pc_deadbeef_123_1")); + // Rendered by a process with a deployment ID. + assert!(target.matches_gid("__pgdog_2pc_prod_deadbeef_123_1")); + + // Wrong shard. + assert!(!target.matches_gid("__pgdog_2pc_deadbeef_123_0")); + // Wrong transaction ID. + assert!(!target.matches_gid("__pgdog_2pc_deadbeef_9123_1")); + // Prepared transactions from other applications. + assert!(!target.matches_gid("app_txn_123_1")); + assert!(!target.matches_gid("123_1")); + assert!(!target.matches_gid("")); + // Matching ID but characters PgDog never generates: such GIDs + // are refused so they can never reach a quoted SQL literal. + assert!(!target.matches_gid("__pgdog_2pc_it's_123_1")); + assert!(!target.matches_gid("__pgdog_2pc_a\\'b_123_1")); + assert!(!target.matches_gid("__pgdog_2pc_a b_123_1")); + } + + #[test] + fn gid_renders_recorded_prefix() { + let transaction: TwoPcTransaction = "__pgdog_2pc_123".parse().unwrap(); + let target = TwoPcTransactionOnShard::new(transaction, 1); + + assert_eq!( + target.gid("__pgdog_2pc_prod_deadbeef_"), + "__pgdog_2pc_prod_deadbeef_123_1" + ); + // The recorded prefix names one exact GID; the same number + // under a different prefix is a different transaction. + assert_ne!(target.gid("__pgdog_2pc_a_"), target.gid("__pgdog_2pc_b_")); + } + #[test] fn phase_control_statements() { let transaction = TwoPcTransaction::new(); diff --git a/pgdog/src/frontend/client/query_engine/two_pc/test.rs b/pgdog/src/frontend/client/query_engine/two_pc/test.rs index 0ad10f27d..3d32b0664 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/test.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/test.rs @@ -83,6 +83,197 @@ async fn test_cleanup_transaction_phase_one() { ); } +#[tokio::test] +async fn test_cleanup_transaction_foreign_prefix() { + config::load_test(); + logger(); + let cluster = databases().all().iter().next().unwrap().1.clone(); + + // A transaction restored from the WAL after a restart. The GIDs on + // the shards carry the previous process's prefix, known here from + // the recorded WAL Begin record. The numeric ID is random so a + // failed run's leftover prepared transaction can't block reruns. + let number: u64 = rand::random(); + let transaction: TwoPcTransaction = format!("__pgdog_2pc_{number}").parse().unwrap(); + let foreign_gid = format!("__pgdog_2pc_previousinstance_{number}_0"); + let table = format!("test_cleanup_foreign_prefix_{number}"); + + let mut conn = Connection::new(cluster.user(), cluster.name(), false).unwrap(); + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); + conn.execute("BEGIN").await.unwrap(); + conn.execute(format!("CREATE TABLE {table}(id BIGINT)")) + .await + .unwrap(); + conn.execute(format!("PREPARE TRANSACTION '{foreign_gid}'")) + .await + .unwrap(); + conn.disconnect(); + + let manager = Manager::get(); + manager.restore_transaction( + transaction, + cluster.user().to_string(), + cluster.name().to_string(), + "__pgdog_2pc_previousinstance_".to_string(), + TwoPcPhase::Phase1, + ); + + manager.wait_until_cleaned_up(transaction).await; + assert!(manager.transaction(&transaction).is_none()); + + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); + + let leftover = conn + .execute(format!( + "SELECT gid FROM pg_prepared_xacts WHERE gid = '{foreign_gid}'" + )) + .await + .unwrap(); + assert!( + leftover.iter().find(|p| p.code() == 'D').is_none(), + "prepared transaction with the previous process's prefix was not cleaned up" + ); + + // Phase 1 transactions are rolled back: the table doesn't exist. + let missing = conn + .execute(format!("SELECT * FROM {table}")) + .await + .err() + .unwrap(); + assert!(missing.to_string().contains("does not exist")); +} + +#[tokio::test] +async fn test_cleanup_transaction_legacy_record() { + config::load_test(); + logger(); + let cluster = databases().all().iter().next().unwrap().1.clone(); + + // A transaction restored from a WAL record written by a version + // that did not store the coordinator GID prefix: cleanup falls + // back to matching by the durable numeric ID. + let number: u64 = rand::random(); + let transaction: TwoPcTransaction = format!("__pgdog_2pc_{number}").parse().unwrap(); + let foreign_gid = format!("__pgdog_2pc_legacyinstance_{number}_0"); + + let mut conn = Connection::new(cluster.user(), cluster.name(), false).unwrap(); + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); + conn.execute("BEGIN").await.unwrap(); + conn.execute("SELECT 1").await.unwrap(); + conn.execute(format!("PREPARE TRANSACTION '{foreign_gid}'")) + .await + .unwrap(); + conn.disconnect(); + + let manager = Manager::get(); + manager.restore_transaction( + transaction, + cluster.user().to_string(), + cluster.name().to_string(), + String::new(), + TwoPcPhase::Phase1, + ); + + manager.wait_until_cleaned_up(transaction).await; + assert!(manager.transaction(&transaction).is_none()); + + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); + + let leftover = conn + .execute(format!( + "SELECT gid FROM pg_prepared_xacts WHERE gid = '{foreign_gid}'" + )) + .await + .unwrap(); + assert!( + leftover.iter().find(|p| p.code() == 'D').is_none(), + "prepared transaction from a legacy WAL record was not cleaned up" + ); +} + +#[tokio::test] +async fn test_cleanup_leaves_other_prefixes_alone() { + config::load_test(); + logger(); + let cluster = databases().all().iter().next().unwrap().1.clone(); + + // A prepared transaction that shares the numeric ID but carries a + // different coordinator prefix belongs to someone else: cleanup + // with a recorded prefix must not touch it. + let number: u64 = rand::random(); + let transaction: TwoPcTransaction = format!("__pgdog_2pc_{number}").parse().unwrap(); + let other_gid = format!("__pgdog_2pc_otherinstance_{number}_0"); + + let mut conn = Connection::new(cluster.user(), cluster.name(), false).unwrap(); + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); + conn.execute("BEGIN").await.unwrap(); + conn.execute("SELECT 1").await.unwrap(); + conn.execute(format!("PREPARE TRANSACTION '{other_gid}'")) + .await + .unwrap(); + conn.disconnect(); + + let manager = Manager::get(); + manager.restore_transaction( + transaction, + cluster.user().to_string(), + cluster.name().to_string(), + "__pgdog_2pc_recordedinstance_".to_string(), + TwoPcPhase::Phase1, + ); + + manager.wait_until_cleaned_up(transaction).await; + // The recorded GID exists nowhere, so the transaction resolves. + assert!(manager.transaction(&transaction).is_none()); + + conn.connect( + &Request::default(), + &Route::write(ShardWithPriority::new_default_unset(Shard::All)), + ) + .await + .unwrap(); + + let survivor = conn + .execute(format!( + "SELECT gid FROM pg_prepared_xacts WHERE gid = '{other_gid}'" + )) + .await + .unwrap(); + assert!( + survivor.iter().find(|p| p.code() == 'D').is_some(), + "prepared transaction under a different prefix was resolved by cleanup" + ); + + conn.execute(format!("ROLLBACK PREPARED '{other_gid}'")) + .await + .unwrap(); +} + #[tokio::test] async fn test_cleanup_transaction_phase_two() { config::load_test(); diff --git a/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs b/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs index 4466fa650..10bf62a43 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs @@ -16,9 +16,14 @@ impl TwoPcTransaction { Self(rng().random_range(0..usize::MAX)) } + /// The durable numeric ID, unique per transaction. + pub(super) fn number(&self) -> usize { + self.0 + } + /// A prefix to identify two-phase commit transactions generated /// by this PgDog process. - fn global_prefix() -> String { + pub(super) fn global_prefix() -> String { format!( "{PREFIX}{}{}_", if let Some(cluster_id) = deployment_id() { @@ -40,7 +45,16 @@ impl Display for TwoPcTransaction { impl FromStr for TwoPcTransaction { type Err = (); + /// Parse a transaction from a GID generated by any PgDog process. + /// The GID must carry the [`PREFIX`] marker; only the trailing + /// numeric ID is significant because the rest of the prefix embeds + /// process identity, which can differ between instances and + /// restarts. fn from_str(s: &str) -> Result { + if !s.starts_with(PREFIX) { + return Err(()); + } + let id = s.rsplit("_").next().map(|id| id.parse()); if let Some(Ok(id)) = id { @@ -65,6 +79,13 @@ mod test { assert_eq!(reverse.0, transaction.0); } + #[test] + fn test_reject_gid_without_prefix() { + assert!(TwoPcTransaction::from_str("123").is_err()); + assert!(TwoPcTransaction::from_str("app_txn_123").is_err()); + assert!(TwoPcTransaction::from_str("__pgdog_123").is_err()); + } + #[test] fn test_instance_id() { for id in [1024, 11111111, usize::MAX, usize::MIN] { diff --git a/pgdog/src/frontend/client/query_engine/two_pc/wal/record.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/record.rs index ead80e1c9..56ebea495 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/record.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/record.rs @@ -96,6 +96,12 @@ pub struct BeginPayload { pub txn: TwoPcTransaction, pub user: String, pub database: String, + /// Coordinator GID prefix the transaction's prepared names are + /// rendered with. Recovery combines it with `txn` to reconstruct + /// the exact GIDs on the shards. Empty when the record was written + /// by a version that did not store it. + #[serde(default)] + pub prefix: String, } /// Payload for records that carry only a transaction id @@ -118,6 +124,11 @@ pub struct CheckpointEntry { pub database: String, /// `true` iff a [`Record::Committing`] had been fsynced for this txn. pub decided: bool, + /// Coordinator GID prefix, carried over from the transaction's + /// [`BeginPayload`]. Empty when the originating record did not + /// store it. + #[serde(default)] + pub prefix: String, } /// A successfully decoded record and the number of bytes it consumed. @@ -227,9 +238,50 @@ mod tests { txn: TwoPcTransaction::new(), user: "alice".into(), database: "shop".into(), + prefix: "__pgdog_2pc_prod_deadbeef_".into(), })); } + #[test] + fn begin_without_prefix_decodes_with_empty_prefix() { + // The on-disk shape written by versions that predate the + // `prefix` field: same tag, payload lacking the key. + #[derive(Serialize)] + struct OldBeginPayload { + txn: TwoPcTransaction, + user: String, + database: String, + } + + let txn = TwoPcTransaction::new(); + let payload = rmp_serde::to_vec_named(&OldBeginPayload { + txn, + user: "alice".into(), + database: "shop".into(), + }) + .unwrap(); + + let mut buf = Vec::new(); + let tag = Tag::Begin as u8; + let mut crc = crc32c::crc32c(&[tag]); + crc = crc32c::crc32c_append(crc, &payload); + buf.put_u32_le(1 + payload.len() as u32); + buf.put_u32_le(crc); + buf.put_u8(tag); + buf.put_slice(&payload); + + let decoded = Record::decode(&buf).unwrap().unwrap(); + assert_eq!( + decoded.record, + Record::Begin(BeginPayload { + txn, + user: "alice".into(), + database: "shop".into(), + prefix: String::new(), + }) + ); + } + #[test] fn round_trip_committing() { round_trip(Record::Committing(TxnPayload { @@ -244,6 +296,57 @@ mod tests { })); } + #[test] + fn checkpoint_without_prefix_decodes_with_empty_prefix() { + // The on-disk shape written by versions that predate the + // `prefix` field on checkpoint entries. + #[derive(Serialize)] + struct OldCheckpointEntry { + txn: TwoPcTransaction, + user: String, + database: String, + decided: bool, + } + #[derive(Serialize)] + struct OldCheckpointPayload { + active: Vec, + } + + let txn = TwoPcTransaction::new(); + let payload = rmp_serde::to_vec_named(&OldCheckpointPayload { + active: vec![OldCheckpointEntry { + txn, + user: "alice".into(), + database: "shop".into(), + decided: true, + }], + }) + .unwrap(); + + let mut buf = Vec::new(); + let tag = Tag::Checkpoint as u8; + let mut crc = crc32c::crc32c(&[tag]); + crc = crc32c::crc32c_append(crc, &payload); + buf.put_u32_le(1 + payload.len() as u32); + buf.put_u32_le(crc); + buf.put_u8(tag); + buf.put_slice(&payload); + + let decoded = Record::decode(&buf).unwrap().unwrap(); + assert_eq!( + decoded.record, + Record::Checkpoint(CheckpointPayload { + active: vec![CheckpointEntry { + txn, + user: "alice".into(), + database: "shop".into(), + decided: true, + prefix: String::new(), + }], + }) + ); + } + #[test] fn round_trip_checkpoint() { round_trip(Record::Checkpoint(CheckpointPayload { @@ -253,12 +356,14 @@ mod tests { user: "u1".into(), database: "d1".into(), decided: false, + prefix: String::new(), }, CheckpointEntry { txn: TwoPcTransaction::new(), user: "u2".into(), database: "d2".into(), decided: true, + prefix: "__pgdog_2pc_prod_a_".into(), }, ], })); @@ -313,6 +418,7 @@ mod tests { txn: TwoPcTransaction::new(), user: "u".into(), database: "d".into(), + prefix: "__pgdog_2pc_".into(), }); let b = Record::Committing(TxnPayload { txn: TwoPcTransaction::new(), diff --git a/pgdog/src/frontend/client/query_engine/two_pc/wal/recovery.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/recovery.rs index 5610fca02..84a52e105 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/recovery.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/recovery.rs @@ -26,6 +26,9 @@ struct Entry { user: String, database: String, decided: bool, + /// Coordinator GID prefix from the Begin record; empty for records + /// written by versions that did not store it. + prefix: String, } /// What [`recover_transactions`] hands back to the WAL setup path. @@ -129,9 +132,10 @@ pub(super) async fn recover_transactions( user: entry.user.clone(), database: entry.database.clone(), decided: entry.decided, + prefix: entry.prefix.clone(), }, ); - manager.restore_transaction(txn, entry.user, entry.database, phase); + manager.restore_transaction(txn, entry.user, entry.database, entry.prefix, phase); } if corruption { tracing::warn!( @@ -184,6 +188,7 @@ fn apply(working: &mut HashMap, record: Record) { user: p.user, database: p.database, decided: false, + prefix: p.prefix, }, ); } @@ -202,7 +207,7 @@ fn apply(working: &mut HashMap, record: Record) { user, database, decided, - .. + prefix, } in p.active { working.insert( @@ -211,6 +216,7 @@ fn apply(working: &mut HashMap, record: Record) { user, database, decided, + prefix, }, ); } @@ -238,11 +244,13 @@ mod tests { txn: txn(1), user: "alice".into(), database: "shop".into(), + prefix: "__pgdog_2pc_a_".into(), }), ); let entry = w.get(&txn(1)).unwrap(); assert_eq!(entry.user, "alice"); assert_eq!(entry.database, "shop"); + assert_eq!(entry.prefix, "__pgdog_2pc_a_"); assert!(!entry.decided); } @@ -255,6 +263,7 @@ mod tests { txn: txn(1), user: "u".into(), database: "d".into(), + prefix: "__pgdog_2pc_a_".into(), }), ); apply(&mut w, Record::Committing(TxnPayload { txn: txn(1) })); @@ -277,6 +286,7 @@ mod tests { txn: txn(1), user: "u".into(), database: "d".into(), + prefix: "__pgdog_2pc_a_".into(), }), ); apply(&mut w, Record::End(TxnPayload { txn: txn(1) })); @@ -292,6 +302,7 @@ mod tests { txn: txn(1), user: "u1".into(), database: "d1".into(), + prefix: "__pgdog_2pc_a_".into(), }), ); apply( @@ -300,6 +311,7 @@ mod tests { txn: txn(2), user: "u2".into(), database: "d2".into(), + prefix: "__pgdog_2pc_a_".into(), }), ); apply( @@ -310,12 +322,14 @@ mod tests { user: "u99".into(), database: "d99".into(), decided: true, + prefix: "__pgdog_2pc_b_".into(), }], }), ); assert_eq!(w.len(), 1); let entry = w.get(&txn(99)).unwrap(); assert_eq!(entry.user, "u99"); + assert_eq!(entry.prefix, "__pgdog_2pc_b_"); assert!(entry.decided); } } diff --git a/pgdog/src/frontend/client/query_engine/two_pc/wal/segment.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/segment.rs index d799380b3..096202347 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/segment.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/segment.rs @@ -410,6 +410,7 @@ mod tests { txn: TwoPcTransaction::new(), user: "u".into(), database: "d".into(), + prefix: "__pgdog_2pc_".into(), }); let r2 = Record::End(TxnPayload { txn: TwoPcTransaction::new(), @@ -442,6 +443,7 @@ mod tests { txn: TwoPcTransaction::new(), user: "u".into(), database: "d".into(), + prefix: "__pgdog_2pc_".into(), }) .encode(&mut buf) .unwrap(); @@ -489,6 +491,7 @@ mod tests { txn: TwoPcTransaction::new(), user: "u".into(), database: "d".into(), + prefix: "__pgdog_2pc_".into(), }); let mut buf = BytesMut::new(); r1.encode(&mut buf).unwrap(); @@ -556,6 +559,7 @@ mod tests { txn: TwoPcTransaction::new(), user: payload.clone(), database: "d".into(), + prefix: "__pgdog_2pc_".into(), }); r.encode(&mut buf).unwrap(); expected.push(r); diff --git a/pgdog/src/frontend/client/query_engine/two_pc/wal/writer.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/writer.rs index ed80e6c0f..1b1e0c7e3 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/writer.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/writer.rs @@ -181,18 +181,20 @@ impl Wal { } /// Log that `txn` is about to issue PREPARE TRANSACTION on its - /// participants. Must complete before any PREPARE leaves the - /// coordinator. + /// participants under GIDs rendered with `prefix`. Must complete + /// before any PREPARE leaves the coordinator. pub async fn append_begin( &self, txn: TwoPcTransaction, user: String, database: String, + prefix: String, ) -> Result> { self.append(Record::Begin(BeginPayload { txn, user, database, + prefix, })) .await } @@ -455,6 +457,7 @@ fn apply_to_snapshot( user: p.user.clone(), database: p.database.clone(), decided: false, + prefix: p.prefix.clone(), }, ); undo.push(Undo { txn: p.txn, prior }); diff --git a/pgdog/src/main.rs b/pgdog/src/main.rs index 684c87e55..980c05a32 100644 --- a/pgdog/src/main.rs +++ b/pgdog/src/main.rs @@ -105,6 +105,9 @@ fn main() -> Result<(), Box> { } async fn pgdog(command: Option) -> Result<(), Box> { + // Fail fast on identifiers that would produce malformed 2pc GIDs. + pgdog::util::validate_instance_identity()?; + // Run atexit handlers on SIGTERM (e.g. llvm-cov profile flushing). #[cfg(unix)] install_sigterm_handler(); diff --git a/pgdog/src/util.rs b/pgdog/src/util.rs index 8ea220448..535676b81 100644 --- a/pgdog/src/util.rs +++ b/pgdog/src/util.rs @@ -116,10 +116,9 @@ pub fn random_string(n: usize) -> String { } // Generate a unique 8-character hex instance ID on first access -static INSTANCE_ID: Lazy = Lazy::new(|| { - if let Ok(node_id) = env::var("NODE_ID") { - node_id - } else { +static INSTANCE_ID: Lazy = Lazy::new(|| match env::var("NODE_ID") { + Ok(node_id) if !node_id.is_empty() => node_id, + _ => { let mut rng = rand::rng(); (0..8) .map(|_| { @@ -132,10 +131,52 @@ static INSTANCE_ID: Lazy = Lazy::new(|| { /// Get the instance ID for this pgdog instance. /// This is generated once at startup and persists for the lifetime of the process. +/// An unset or empty `NODE_ID` auto-generates a random ID; otherwise the value +/// is restricted to [`is_safe_identifier`] characters, which +/// [`validate_instance_identity`] enforces at startup. pub fn instance_id() -> &'static str { &INSTANCE_ID } +/// True if `s` contains only ASCII alphanumerics, underscores and hyphens: +/// the alphabet PgDog allows in instance identifiers and, by extension, in +/// the two-phase commit GIDs they are embedded in. +pub fn is_safe_identifier(s: &str) -> bool { + s.bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-') +} + +/// Environment-provided instance identifier rejected by +/// [`validate_instance_identity`]. +#[derive(Debug, thiserror::Error)] +#[error( + "{var} may only contain ASCII letters, digits, '_' and '-' \ + (it is embedded in two-phase commit transaction identifiers), got: {value:?}" +)] +pub struct IdentityError { + pub var: &'static str, + pub value: String, +} + +/// Validate the environment-provided identifiers that end up embedded in +/// two-phase commit GIDs. Called at startup so a misconfigured `NODE_ID` +/// or `DEPLOYMENT_ID` fails fast; GIDs built from validated identifiers +/// contain no characters that need quoting, which cleanup relies on when +/// it embeds recovered GIDs into transaction control statements. +/// +/// Unset and empty values are valid: both mean the identifier is absent. +pub fn validate_instance_identity() -> Result<(), IdentityError> { + for var in ["NODE_ID", "DEPLOYMENT_ID"] { + if let Ok(value) = env::var(var) + && !is_safe_identifier(&value) + { + return Err(IdentityError { var, value }); + } + } + + Ok(()) +} + /// Get an externally assigned, unique, node identifier /// for this instance of PgDog. /// @@ -148,12 +189,16 @@ pub fn node_id() -> Result { instance_id().split("-").last().unwrap().parse() } -static DEPLOYMENT_ID: Lazy> = Lazy::new(|| env::var("DEPLOYMENT_ID").ok()); +static DEPLOYMENT_ID: Lazy> = + Lazy::new(|| env::var("DEPLOYMENT_ID").ok().filter(|id| !id.is_empty())); /// Get the ID of this PgDog deployment. /// /// This should be _globally_ unique /// and is used to differentiate 2pc transactions. +/// An unset or empty `DEPLOYMENT_ID` means no deployment ID; otherwise the +/// value is restricted to [`is_safe_identifier`] characters, which +/// [`validate_instance_identity`] enforces at startup. /// pub(crate) fn deployment_id() -> Option<&'static str> { DEPLOYMENT_ID.as_deref() @@ -341,6 +386,42 @@ mod test { use super::*; use crate::test_utils::*; + #[test] + fn test_is_safe_identifier() { + assert!(is_safe_identifier("pgdog-node-1")); + assert!(is_safe_identifier("prod_us_east_2")); + assert!(is_safe_identifier("deadbeef")); + assert!(!is_safe_identifier("it's")); + assert!(!is_safe_identifier("node\\1")); + assert!(!is_safe_identifier("node 1")); + assert!(!is_safe_identifier("nöde")); + } + + #[test] + fn test_validate_instance_identity() { + let _node = set_env_var("NODE_ID", "pgdog-node-1"); + let _deployment = set_env_var("DEPLOYMENT_ID", "prod"); + assert!(validate_instance_identity().is_ok()); + + { + let _bad = set_env_var("NODE_ID", "it's"); + assert!(matches!( + validate_instance_identity(), + Err(IdentityError { var: "NODE_ID", .. }) + )); + } + + // Empty means absent, like unset. + let _empty = set_env_var("NODE_ID", ""); + assert!(validate_instance_identity().is_ok()); + } + + #[test] + fn test_empty_deployment_id_is_absent() { + let _empty = set_env_var("DEPLOYMENT_ID", ""); + assert_eq!(deployment_id(), None); + } + #[test] fn test_human_duration() { assert_eq!(human_duration(Duration::from_millis(500)), "500ms");