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..f62cdbf57 100644 --- a/integration/two_pc_crash_safety/wal_helper/src/main.rs +++ b/integration/two_pc_crash_safety/wal_helper/src/main.rs @@ -23,8 +23,8 @@ async fn main() { std::process::exit(2); } let dir = PathBuf::from(&args[1]); - let txn = TwoPcTransaction::from_str(&args[2]) - .unwrap_or_else(|_| panic!("invalid gid {:?}", args[2])); + let gid = args[2].clone(); + let txn = TwoPcTransaction::from_str(&gid).unwrap_or_else(|_| panic!("invalid gid {:?}", gid)); let user = args[3].clone(); let database = args[4].clone(); @@ -34,9 +34,10 @@ async fn main() { let mut buf = BytesMut::new(); Record::Begin(BeginPayload { - txn, + txn: txn.clone(), user, database, + gid, }) .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..7a77b7e4a 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -364,7 +364,7 @@ impl Binding { pub(crate) async fn two_pc_on_guards( servers: &mut [Guard], - transaction: TwoPcTransaction, + transaction: &TwoPcTransaction, phase: TwoPcPhase, ) -> Result<(), Error> { let skip_missing = matches!(phase, TwoPcPhase::Phase2 | TwoPcPhase::Rollback); @@ -397,9 +397,14 @@ impl Binding { } /// Execute two-phase commit transaction control statements. + /// + /// The per-shard participant names are derived from `transaction`. For a + /// transaction rebuilt by WAL recovery this reproduces the exact gid it + /// was prepared with, even after a restart with a different + /// `instance_id` (see [`TwoPcTransaction`]). pub(crate) async fn two_pc( &mut self, - transaction: TwoPcTransaction, + transaction: &TwoPcTransaction, phase: TwoPcPhase, ) -> Result<(), Error> { match self { diff --git a/pgdog/src/backend/pool/connection/binding_test.rs b/pgdog/src/backend/pool/connection/binding_test.rs index 0a6258b6e..6d59672a5 100644 --- a/pgdog/src/backend/pool/connection/binding_test.rs +++ b/pgdog/src/backend/pool/connection/binding_test.rs @@ -77,7 +77,7 @@ mod tests { let mut binding = Binding::Direct(guard, 0); let result = binding - .two_pc(TwoPcTransaction::new(), TwoPcPhase::Phase1) + .two_pc(&TwoPcTransaction::new(), TwoPcPhase::Phase1) .await; // Should fail with TwoPcMultiShardOnly error @@ -97,7 +97,7 @@ mod tests { let mut binding = Binding::Admin(admin_server); let result = binding - .two_pc(TwoPcTransaction::new(), TwoPcPhase::Phase1) + .two_pc(&TwoPcTransaction::new(), TwoPcPhase::Phase1) .await; // Should fail with TwoPcMultiShardOnly error @@ -113,7 +113,7 @@ mod tests { let transaction = TwoPcTransaction::new(); // Test Phase1 - PREPARE TRANSACTION - let result = binding.two_pc(transaction, TwoPcPhase::Phase1).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase1).await; // Should succeed if let Err(ref error) = result { @@ -122,7 +122,7 @@ mod tests { assert!(result.is_ok()); // Cleanup: Rollback the prepared transaction to avoid leaving dangling transactions - let _cleanup = binding.two_pc(transaction, TwoPcPhase::Rollback).await; + let _cleanup = binding.two_pc(&transaction, TwoPcPhase::Rollback).await; } #[tokio::test] @@ -132,12 +132,12 @@ mod tests { // First prepare the transaction binding - .two_pc(transaction, TwoPcPhase::Phase1) + .two_pc(&transaction, TwoPcPhase::Phase1) .await .expect("Phase1 should succeed"); // Then commit it - let result = binding.two_pc(transaction, TwoPcPhase::Phase2).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase2).await; assert!(result.is_ok()); } @@ -148,12 +148,12 @@ mod tests { // First prepare the transaction binding - .two_pc(transaction, TwoPcPhase::Phase1) + .two_pc(&transaction, TwoPcPhase::Phase1) .await .expect("Phase1 should succeed"); // Then rollback - let result = binding.two_pc(transaction, TwoPcPhase::Rollback).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Rollback).await; assert!(result.is_ok()); } @@ -164,18 +164,18 @@ mod tests { // First prepare the transaction binding - .two_pc(transaction, TwoPcPhase::Phase1) + .two_pc(&transaction, TwoPcPhase::Phase1) .await .expect("Phase1 should succeed"); // Then commit it binding - .two_pc(transaction, TwoPcPhase::Phase2) + .two_pc(&transaction, TwoPcPhase::Phase2) .await .expect("Phase2 should succeed"); // Try to commit again - should succeed because skip_missing is true for Phase2 - let result = binding.two_pc(transaction, TwoPcPhase::Phase2).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase2).await; assert!( result.is_ok(), "Committing non-existent prepared transaction should be skipped" @@ -189,18 +189,18 @@ mod tests { // First prepare the transaction binding - .two_pc(transaction, TwoPcPhase::Phase1) + .two_pc(&transaction, TwoPcPhase::Phase1) .await .expect("Phase1 should succeed"); // Then rollback it binding - .two_pc(transaction, TwoPcPhase::Rollback) + .two_pc(&transaction, TwoPcPhase::Rollback) .await .expect("Rollback should succeed"); // Try to rollback again - should succeed because skip_missing is true for Rollback - let result = binding.two_pc(transaction, TwoPcPhase::Rollback).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Rollback).await; assert!( result.is_ok(), "Rolling back non-existent prepared transaction should be skipped" @@ -216,19 +216,19 @@ mod tests { let transaction = TwoPcTransaction::new(); // 1. Prepare transaction - let result = binding.two_pc(transaction, TwoPcPhase::Phase1).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase1).await; assert!(result.is_ok(), "Phase1 preparation should succeed"); // 2. Try to prepare the same transaction again - PostgreSQL behavior may vary - let _result = binding.two_pc(transaction, TwoPcPhase::Phase1).await; + let _result = binding.two_pc(&transaction, TwoPcPhase::Phase1).await; // Note: PostgreSQL behavior for duplicate PREPARE TRANSACTION can vary depending on context // 3. Commit the prepared transaction - let result = binding.two_pc(transaction, TwoPcPhase::Phase2).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase2).await; assert!(result.is_ok(), "Phase2 commit should succeed"); // 4. Try to commit again - should succeed (skip_missing = true) - let result = binding.two_pc(transaction, TwoPcPhase::Phase2).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase2).await; assert!( result.is_ok(), "Committing non-existent transaction should be skipped" @@ -241,15 +241,15 @@ mod tests { let transaction = TwoPcTransaction::new(); // 1. Prepare transaction - let result = binding.two_pc(transaction, TwoPcPhase::Phase1).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase1).await; assert!(result.is_ok(), "Phase1 preparation should succeed"); // 2. Rollback the prepared transaction - let result = binding.two_pc(transaction, TwoPcPhase::Rollback).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Rollback).await; assert!(result.is_ok(), "Rollback should succeed"); // 3. Try to commit after rollback - should succeed (skip_missing = true) - let result = binding.two_pc(transaction, TwoPcPhase::Phase2).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase2).await; assert!( result.is_ok(), "Committing rolled back transaction should be skipped" diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index 49cbfb7b7..1f32c5065 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -255,7 +255,7 @@ impl Error { /// Two-phase commit transaction that still needs manager cleanup, if any. pub fn two_pc_cleanup_transaction(&self) -> Option { match self { - Self::TwoPcCleanupPending { transaction, .. } => Some(*transaction), + Self::TwoPcCleanupPending { transaction, .. } => Some(transaction.clone()), _ => None, } } diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index 6b7c5eff0..7ac2161b9 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -296,14 +296,14 @@ impl CopySubscriber { async { let _guard_phase_1 = manager - .transaction_state(txn, &identifier, TwoPcPhase::Phase1) + .transaction_state(txn.clone(), &identifier, TwoPcPhase::Phase1) .await?; - self.two_pc_on_shards(txn, TwoPcPhase::Phase1).await?; + self.two_pc_on_shards(&txn, TwoPcPhase::Phase1).await?; let _guard_phase_2 = manager - .transaction_state(txn, &identifier, TwoPcPhase::Phase2) + .transaction_state(txn.clone(), &identifier, TwoPcPhase::Phase2) .await?; - self.two_pc_on_shards(txn, TwoPcPhase::Phase2).await?; + self.two_pc_on_shards(&txn, TwoPcPhase::Phase2).await?; manager.done(&txn).await?; Ok(()) @@ -317,7 +317,7 @@ impl CopySubscriber { async fn two_pc_on_shards( &mut self, - txn: TwoPcTransaction, + txn: &TwoPcTransaction, phase: TwoPcPhase, ) -> Result<(), Error> { let mut futures = Vec::new(); diff --git a/pgdog/src/frontend/client/query_engine/end_transaction.rs b/pgdog/src/frontend/client/query_engine/end_transaction.rs index affcc9053..f8b9630d2 100644 --- a/pgdog/src/frontend/client/query_engine/end_transaction.rs +++ b/pgdog/src/frontend/client/query_engine/end_transaction.rs @@ -110,13 +110,17 @@ impl QueryEngine { // If interrupted here, the transaction must be rolled back. let _guard_phase_1 = self.two_pc.phase_one(&identifier).await?; - self.backend.two_pc(transaction, TwoPcPhase::Phase1).await?; + self.backend + .two_pc(&transaction, TwoPcPhase::Phase1) + .await?; debug!("[2pc] phase 1 complete"); // If interrupted here, the transaction must be committed. let _guard_phase_2 = self.two_pc.phase_two(&identifier).await?; - self.backend.two_pc(transaction, TwoPcPhase::Phase2).await?; + self.backend + .two_pc(&transaction, TwoPcPhase::Phase2) + .await?; debug!("[2pc] phase 2 complete"); 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..86bd75627 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs @@ -208,7 +208,7 @@ impl Manager { let prior = { let mut guard = self.inner.lock(); let prior = guard.transactions.get(&transaction).cloned(); - let entry = guard.transactions.entry(transaction).or_default(); + let entry = guard.transactions.entry(transaction.clone()).or_default(); entry.identifier = identifier.clone(); entry.phase = phase; prior @@ -218,13 +218,14 @@ impl Manager { let result = match phase { TwoPcPhase::Phase1 => { wal.append_begin( - transaction, + transaction.clone(), identifier.user.clone(), identifier.database.clone(), + transaction.to_string(), ) .await } - TwoPcPhase::Phase2 => wal.append_committing(transaction).await, + TwoPcPhase::Phase2 => wal.append_committing(transaction.clone()).await, TwoPcPhase::Rollback => { unreachable!("rollback is not a state transition; it's the cleanup direction") } @@ -233,7 +234,7 @@ impl Manager { let mut guard = self.inner.lock(); match prior { Some(prior) => { - guard.transactions.insert(transaction, prior); + guard.transactions.insert(transaction.clone(), prior); } None => { guard.transactions.remove(&transaction); @@ -269,7 +270,7 @@ impl Manager { let mut guard = self.inner.lock(); guard .transactions - .insert(transaction, TransactionInfo { phase, identifier }); + .insert(transaction.clone(), TransactionInfo { phase, identifier }); guard.queue.push_back(transaction); } self.stats.incr_recovered(); @@ -284,7 +285,7 @@ impl Manager { .contains_key(&guard.transaction); if exists { - self.inner.lock().queue.push_back(guard.transaction); + self.inner.lock().queue.push_back(guard.transaction.clone()); self.notify.notify.notify_one(); } } @@ -309,7 +310,7 @@ impl Manager { r#"[2pc] cleaning up transaction "{}""#, transaction.to_string() ); - match manager.cleanup_phase(transaction).await { + match manager.cleanup_phase(&transaction).await { Err(err) => { error!( r#"[2pc] error cleaning up "{}" transaction: {}"#, @@ -338,15 +339,15 @@ impl Manager { async fn remove(&self, transaction: &TwoPcTransaction) { self.inner.lock().transactions.remove(transaction); if let Some(wal) = self.wal.load_full() - && let Err(err) = wal.append_end(*transaction).await + && let Err(err) = wal.append_end(transaction.clone()).await { warn!("[2pc] wal end record failed for {}: {}", transaction, err); } } /// Reconnect to cluster if available and rollback the two-phase transaction. - async fn cleanup_phase(&self, transaction: TwoPcTransaction) -> Result<(), Error> { - let state = match self.inner.lock().transactions.get(&transaction).cloned() { + async fn cleanup_phase(&self, transaction: &TwoPcTransaction) -> Result<(), Error> { + let state = match self.inner.lock().transactions.get(transaction).cloned() { Some(state) => state, _ => { return Ok(()); @@ -379,6 +380,9 @@ impl Manager { &Route::write(ShardWithPriority::new_override_transaction(Shard::All)), ) .await?; + // `transaction` carries its own gid: for a recovered transaction + // that's the exact value it was prepared with, so this resolves the + // orphan on Postgres even after a restart with a new instance_id. connection.two_pc(transaction, phase).await?; connection.disconnect(); diff --git a/pgdog/src/frontend/client/query_engine/two_pc/mod.rs b/pgdog/src/frontend/client/query_engine/two_pc/mod.rs index b0cba45aa..dfc0dae9e 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/mod.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/mod.rs @@ -46,11 +46,9 @@ impl Default for TwoPc { impl TwoPc { /// Get a unique name for the two-pc transaction. pub(super) fn transaction(&mut self) -> TwoPcTransaction { - if self.transaction.is_none() { - self.transaction = Some(TwoPcTransaction::new()); - } - - self.transaction.unwrap() + self.transaction + .get_or_insert_with(TwoPcTransaction::new) + .clone() } /// Start phase one of two-phase commit. 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..a236ee9e8 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/statement.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/statement.rs @@ -22,7 +22,7 @@ impl TwoPcTransactionOnShard { /// Get the coordinator transaction. pub(crate) fn transaction(&self) -> TwoPcTransaction { - self.transaction + self.transaction.clone() } } @@ -47,17 +47,22 @@ impl FromStr for TwoPcTransactionOnShard { /// Build `PREPARE TRANSACTION`, `COMMIT PREPARED`, or `ROLLBACK PREPARED` /// for a shard participant. +/// +/// The per-shard participant name is `_`, rendered +/// via [`TwoPcTransactionOnShard`]. For a transaction recovered from the +/// WAL the coordinator gid is carried verbatim by [`TwoPcTransaction`], so +/// this produces the exact name Postgres holds even after a restart with a +/// different `instance_id`. pub(crate) fn phase_control( - transaction: TwoPcTransaction, + transaction: &TwoPcTransaction, shard: usize, phase: TwoPcPhase, ) -> String { - let txn = TwoPcTransactionOnShard::new(transaction, shard); - + let participant = TwoPcTransactionOnShard::new(transaction.clone(), shard).to_string(); match phase { - TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{txn}'"), - TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{txn}'"), - TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{txn}'"), + TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{participant}'"), + TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{participant}'"), + TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{participant}'"), } } @@ -70,11 +75,11 @@ mod test { let transaction = TwoPcTransaction::new(); assert_eq!( - TwoPcTransactionOnShard::new(transaction, 0).to_string(), + TwoPcTransactionOnShard::new(transaction.clone(), 0).to_string(), format!("{transaction}_0") ); assert_eq!( - TwoPcTransactionOnShard::new(transaction, 3).to_string(), + TwoPcTransactionOnShard::new(transaction.clone(), 3).to_string(), format!("{transaction}_3") ); } @@ -106,16 +111,32 @@ mod test { let transaction = TwoPcTransaction::new(); assert_eq!( - phase_control(transaction, 1, TwoPcPhase::Phase1), + phase_control(&transaction, 1, TwoPcPhase::Phase1), format!("PREPARE TRANSACTION '{transaction}_1'") ); assert_eq!( - phase_control(transaction, 1, TwoPcPhase::Phase2), + phase_control(&transaction, 1, TwoPcPhase::Phase2), format!("COMMIT PREPARED '{transaction}_1'") ); assert_eq!( - phase_control(transaction, 1, TwoPcPhase::Rollback), + phase_control(&transaction, 1, TwoPcPhase::Rollback), format!("ROLLBACK PREPARED '{transaction}_1'") ); } + + #[test] + fn phase_control_uses_recovered_gid_verbatim() { + // A transaction recovered from the WAL carries a gid that no longer + // matches this process's rendering (different instance_id). The + // control statement must use it exactly as stored, per-shard suffix + // appended. This is the invariant WAL recovery relies on across + // restarts. + let recovered = "__pgdog_2pc_oldnode_42" + .parse::() + .expect("valid recovered gid"); + assert_eq!( + phase_control(&recovered, 3, TwoPcPhase::Rollback), + "ROLLBACK PREPARED '__pgdog_2pc_oldnode_42_3'" + ); + } } 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..1ead99921 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/test.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/test.rs @@ -38,7 +38,7 @@ async fn test_cleanup_transaction_phase_one() { let info = Manager::get().transaction(&transaction).unwrap(); assert_eq!(info.phase, TwoPcPhase::Phase1); - conn.two_pc(transaction, TwoPcPhase::Phase1).await.unwrap(); + conn.two_pc(&transaction, TwoPcPhase::Phase1).await.unwrap(); let two_pc = conn .execute("SELECT * FROM pg_prepared_xacts") @@ -108,7 +108,7 @@ async fn test_cleanup_transaction_phase_two() { let info = Manager::get().transaction(&transaction).unwrap(); assert_eq!(info.phase, TwoPcPhase::Phase1); - conn.two_pc(transaction, TwoPcPhase::Phase1).await.unwrap(); + conn.two_pc(&transaction, TwoPcPhase::Phase1).await.unwrap(); let txns = conn .execute("SELECT * FROM pg_prepared_xacts") 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..a1505362b 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/transaction.rs @@ -1,11 +1,34 @@ use rand::{Rng, rng}; -use serde::{Deserialize, Serialize}; -use std::{fmt::Display, str::FromStr}; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use std::sync::Arc; +use std::{ + fmt::Display, + hash::{Hash, Hasher}, + str::FromStr, +}; use crate::util::{deployment_id, instance_id}; -#[derive(Debug, Clone, Copy, PartialEq, Hash, Eq, Serialize, Deserialize)] -pub struct TwoPcTransaction(usize); +/// Coordinator identifier for a two-phase commit transaction. +/// +/// A live transaction is just a random `id`; its gid string is rendered on +/// demand from this process's `instance_id`/`deployment_id`. A restarted +/// PgDog generates a fresh `instance_id`, so a transaction rebuilt during WAL +/// recovery carries the original gid verbatim in `gid` instead of +/// re-rendering it (which would no longer match the name Postgres holds in +/// `pg_prepared_xacts`). +/// +/// Identity (`Hash`/`Eq`) is the `id` only: the gid embeds it as its trailing +/// component, so a recovered transaction and its live counterpart compare +/// equal and collate in the same map slot. +#[derive(Debug, Clone)] +pub struct TwoPcTransaction { + id: usize, + /// Full coordinator gid, set only when it must be preserved verbatim + /// (a transaction recovered from the WAL). `None` for transactions + /// created in this process, where `Display` renders the gid live. + gid: Option>, +} static PREFIX: &str = "__pgdog_2pc_"; @@ -13,7 +36,18 @@ impl TwoPcTransaction { pub(crate) fn new() -> Self { // Transactions have random identifiers, // so multiple instances of PgDog don't create an identical transaction. - Self(rng().random_range(0..usize::MAX)) + Self { + id: rng().random_range(0..usize::MAX), + gid: None, + } + } + + /// Attach the exact gid this transaction was prepared with, so `Display` + /// reproduces it verbatim regardless of the current process's + /// `instance_id`. Used by WAL recovery. + pub(crate) fn with_gid(mut self, gid: impl Into>) -> Self { + self.gid = Some(gid.into()); + self } /// A prefix to identify two-phase commit transactions generated @@ -33,7 +67,43 @@ impl TwoPcTransaction { impl Display for TwoPcTransaction { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}{}", Self::global_prefix(), self.0) + match &self.gid { + Some(gid) => f.write_str(gid), + None => write!(f, "{}{}", Self::global_prefix(), self.id), + } + } +} + +impl PartialEq for TwoPcTransaction { + fn eq(&self, other: &Self) -> bool { + self.id == other.id + } +} + +impl Eq for TwoPcTransaction {} + +impl Hash for TwoPcTransaction { + fn hash(&self, state: &mut H) { + self.id.hash(state); + } +} + +/// On disk and on the wire a transaction is just its `id`; the gid is never +/// serialized. Recovery reattaches it from [`super::wal::BeginPayload::gid`]. +/// Keeping the representation a bare integer preserves compatibility with WAL +/// segments written before the gid was persisted. +impl Serialize for TwoPcTransaction { + fn serialize(&self, serializer: S) -> Result { + self.id.serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for TwoPcTransaction { + fn deserialize>(deserializer: D) -> Result { + Ok(Self { + id: usize::deserialize(deserializer)?, + gid: None, + }) } } @@ -44,7 +114,12 @@ impl FromStr for TwoPcTransaction { let id = s.rsplit("_").next().map(|id| id.parse()); if let Some(Ok(id)) = id { - Ok(Self(id)) + Ok(Self { + id, + // Preserve the parsed name verbatim: it may carry another + // process's instance_id that this one cannot reproduce. + gid: Some(Arc::from(s)), + }) } else { Err(()) } @@ -57,18 +132,33 @@ mod test { use super::*; + fn with_id(id: usize) -> TwoPcTransaction { + TwoPcTransaction { id, gid: None } + } + #[test] fn test_2pc_transaction_id() { let transaction = TwoPcTransaction::new(); assert!(transaction.to_string().contains("__pgdog_2pc_")); let reverse = TwoPcTransaction::from_str(transaction.to_string().as_str()).unwrap(); - assert_eq!(reverse.0, transaction.0); + assert_eq!(reverse.id, transaction.id); + } + + #[test] + fn recovered_gid_is_rendered_verbatim() { + // A gid from another process (different instance_id) must round-trip + // through Display unchanged, not be re-rendered with this process's + // prefix. + let stored = "__pgdog_2pc_oldnode_42"; + let txn = stored.parse::().unwrap(); + assert_eq!(txn.to_string(), stored); + assert_eq!(txn.id, 42); } #[test] fn test_instance_id() { for id in [1024, 11111111, usize::MAX, usize::MIN] { - let transaction = TwoPcTransaction(id); + let transaction = with_id(id); let instance_id = instance_id(); // It's a singleton. assert_eq!( format!("__pgdog_2pc_{instance_id}_{id}"), @@ -80,7 +170,7 @@ mod test { #[test] fn test_deployment_id() { let _guard = set_env_var("DEPLOYMENT_ID", "1"); - let txn = TwoPcTransaction(1678); + let txn = with_id(1678); let instance_id = instance_id(); // It's a singleton. assert_eq!(format!("__pgdog_2pc_1_{instance_id}_1678"), txn.to_string()); } diff --git a/pgdog/src/frontend/client/query_engine/two_pc/wal/mod.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/mod.rs index 6ec8abca8..a34b5383b 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/mod.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/mod.rs @@ -1,6 +1,6 @@ //! Two-phase commit write-ahead log. //! -//! See [`record`] for the on-disk record format. +//! See the `record` module for the on-disk record format. mod error; mod record; 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..9940716d9 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,14 @@ pub struct BeginPayload { pub txn: TwoPcTransaction, pub user: String, pub database: String, + /// The full coordinator gid (`TwoPcTransaction`'s `Display`) captured + /// when the transaction was created. It embeds a per-process + /// `instance_id` that a restarted PgDog does not reproduce, so + /// recovery must drive `COMMIT PREPARED` / `ROLLBACK PREPARED` with + /// this stored value rather than re-rendering the `txn`. Empty when + /// read back from a segment written before this field existed. + #[serde(default)] + pub gid: String, } /// Payload for records that carry only a transaction id @@ -118,6 +126,10 @@ pub struct CheckpointEntry { pub database: String, /// `true` iff a [`Record::Committing`] had been fsynced for this txn. pub decided: bool, + /// Coordinator gid carried forward from the txn's [`BeginPayload`]. + /// See [`BeginPayload::gid`]. + #[serde(default)] + pub gid: String, } /// A successfully decoded record and the number of bytes it consumed. @@ -221,12 +233,32 @@ mod tests { assert_eq!(d.consumed, buf.len()); } + /// Frame `tag + payload` exactly like [`Record::encode`] so a payload + /// serialized from an older struct can be fed through + /// [`Record::decode`]. Used to simulate segments written before the + /// `gid` field existed. + fn frame(tag: Tag, payload: &[u8]) -> Vec { + let body_len = (1 + payload.len()) as u32; + let mut crc = crc32c::crc32c(&[tag as u8]); + crc = crc32c::crc32c_append(crc, payload); + + let mut buf = Vec::new(); + buf.put_u32_le(body_len); + buf.put_u32_le(crc); + buf.put_u8(tag as u8); + buf.put_slice(payload); + buf + } + #[test] fn round_trip_begin() { + let txn = TwoPcTransaction::new(); + let gid = txn.to_string(); round_trip(Record::Begin(BeginPayload { - txn: TwoPcTransaction::new(), + txn, user: "alice".into(), database: "shop".into(), + gid, })); } @@ -253,12 +285,14 @@ mod tests { user: "u1".into(), database: "d1".into(), decided: false, + gid: "__pgdog_2pc_n1_1".into(), }, CheckpointEntry { txn: TwoPcTransaction::new(), user: "u2".into(), database: "d2".into(), decided: true, + gid: "__pgdog_2pc_n1_2".into(), }, ], })); @@ -313,6 +347,7 @@ mod tests { txn: TwoPcTransaction::new(), user: "u".into(), database: "d".into(), + gid: "__pgdog_2pc_n_1".into(), }); let b = Record::Committing(TxnPayload { txn: TwoPcTransaction::new(), @@ -326,4 +361,81 @@ mod tests { assert_eq!(d2.record, b); assert_eq!(d1.consumed + d2.consumed, buf.len()); } + + #[test] + fn begin_without_gid_field_decodes_to_empty_gid() { + // A segment written before `gid` existed carries a Begin payload + // with only txn/user/database. `#[serde(default)]` must fill an + // empty gid so recovery can detect the record predates persistence + // and fall back to re-rendering the transaction. + #[derive(Serialize)] + struct OldBeginPayload { + txn: TwoPcTransaction, + user: String, + database: String, + } + + let txn = TwoPcTransaction::new(); + let payload = rmp_serde::to_vec_named(&OldBeginPayload { + txn: txn.clone(), + user: "alice".into(), + database: "shop".into(), + }) + .unwrap(); + let buf = frame(Tag::Begin, &payload); + + let decoded = Record::decode(&buf).unwrap().unwrap(); + assert_eq!(decoded.consumed, buf.len()); + match decoded.record { + Record::Begin(p) => { + assert_eq!(p.txn, txn); + assert_eq!(p.user, "alice"); + assert_eq!(p.database, "shop"); + assert_eq!(p.gid, ""); + } + other => panic!("expected Begin, got {:?}", other), + } + } + + #[test] + fn checkpoint_entry_without_gid_field_decodes_to_empty_gid() { + // Checkpoints written before `gid` existed carry entries without + // the field; each must default to an empty gid so recovery treats + // them as pre-persistence records. + #[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: txn.clone(), + user: "u".into(), + database: "d".into(), + decided: true, + }], + }) + .unwrap(); + let buf = frame(Tag::Checkpoint, &payload); + + let decoded = Record::decode(&buf).unwrap().unwrap(); + match decoded.record { + Record::Checkpoint(p) => { + assert_eq!(p.active.len(), 1); + let entry = &p.active[0]; + assert_eq!(entry.txn, txn); + assert!(entry.decided); + assert_eq!(entry.gid, ""); + } + other => panic!("expected Checkpoint, got {:?}", other), + } + } } 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..b15e3e2f6 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,10 @@ struct Entry { user: String, database: String, decided: bool, + /// Coordinator gid the transaction was prepared with. See + /// [`super::record::BeginPayload::gid`]. Recovery drives + /// `COMMIT PREPARED` / `ROLLBACK PREPARED` with this exact value. + gid: String, } /// What [`recover_transactions`] hands back to the WAL setup path. @@ -123,15 +127,25 @@ pub(super) async fn recover_transactions( (false, true) => continue, }; snapshot.insert( - txn, + txn.clone(), CheckpointEntry { - txn, + txn: txn.clone(), user: entry.user.clone(), database: entry.database.clone(), decided: entry.decided, + gid: entry.gid.clone(), }, ); - manager.restore_transaction(txn, entry.user, entry.database, phase); + // Attach the exact gid the xact was prepared with so cleanup drives + // COMMIT/ROLLBACK PREPARED with it verbatim. An empty gid means the + // record predates gid persistence; leave the transaction to render + // from live process state (no worse than before persistence). + let restored = if entry.gid.is_empty() { + txn + } else { + txn.with_gid(entry.gid) + }; + manager.restore_transaction(restored, entry.user, entry.database, phase); } if corruption { tracing::warn!( @@ -184,6 +198,7 @@ fn apply(working: &mut HashMap, record: Record) { user: p.user, database: p.database, decided: false, + gid: p.gid, }, ); } @@ -202,6 +217,7 @@ fn apply(working: &mut HashMap, record: Record) { user, database, decided, + gid, .. } in p.active { @@ -211,6 +227,7 @@ fn apply(working: &mut HashMap, record: Record) { user, database, decided, + gid, }, ); } @@ -238,14 +255,54 @@ mod tests { txn: txn(1), user: "alice".into(), database: "shop".into(), + gid: txn(1).to_string(), }), ); let entry = w.get(&txn(1)).unwrap(); assert_eq!(entry.user, "alice"); assert_eq!(entry.database, "shop"); + assert_eq!(entry.gid, txn(1).to_string()); assert!(!entry.decided); } + #[test] + fn begin_without_gid_keeps_empty_gid() { + // A Begin decoded from a pre-persistence segment has an empty gid. + // Recovery must preserve it: `restore_transaction` maps an empty + // gid to the re-render fallback, so losing it here would send a + // wrong gid to Postgres and miss the orphan. + let mut w: HashMap = HashMap::default(); + apply( + &mut w, + Record::Begin(BeginPayload { + txn: txn(1), + user: "u".into(), + database: "d".into(), + gid: String::new(), + }), + ); + assert_eq!(w.get(&txn(1)).unwrap().gid, ""); + } + + #[test] + fn checkpoint_entry_without_gid_keeps_empty_gid() { + // Same fallback signal, carried through a checkpoint entry. + let mut w: HashMap = HashMap::default(); + apply( + &mut w, + Record::Checkpoint(CheckpointPayload { + active: vec![CheckpointEntry { + txn: txn(1), + user: "u".into(), + database: "d".into(), + decided: true, + gid: String::new(), + }], + }), + ); + assert_eq!(w.get(&txn(1)).unwrap().gid, ""); + } + #[test] fn committing_marks_decided() { let mut w: HashMap = HashMap::default(); @@ -255,6 +312,7 @@ mod tests { txn: txn(1), user: "u".into(), database: "d".into(), + gid: txn(1).to_string(), }), ); apply(&mut w, Record::Committing(TxnPayload { txn: txn(1) })); @@ -277,6 +335,7 @@ mod tests { txn: txn(1), user: "u".into(), database: "d".into(), + gid: txn(1).to_string(), }), ); apply(&mut w, Record::End(TxnPayload { txn: txn(1) })); @@ -292,6 +351,7 @@ mod tests { txn: txn(1), user: "u1".into(), database: "d1".into(), + gid: txn(1).to_string(), }), ); apply( @@ -300,6 +360,7 @@ mod tests { txn: txn(2), user: "u2".into(), database: "d2".into(), + gid: txn(2).to_string(), }), ); apply( @@ -310,12 +371,14 @@ mod tests { user: "u99".into(), database: "d99".into(), decided: true, + gid: txn(99).to_string(), }], }), ); assert_eq!(w.len(), 1); let entry = w.get(&txn(99)).unwrap(); assert_eq!(entry.user, "u99"); + assert_eq!(entry.gid, txn(99).to_string()); 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..4b2098938 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(), + gid: "__pgdog_2pc_n_1".into(), }); let r2 = Record::End(TxnPayload { txn: TwoPcTransaction::new(), @@ -442,6 +443,7 @@ mod tests { txn: TwoPcTransaction::new(), user: "u".into(), database: "d".into(), + gid: "__pgdog_2pc_n_1".into(), }) .encode(&mut buf) .unwrap(); @@ -489,6 +491,7 @@ mod tests { txn: TwoPcTransaction::new(), user: "u".into(), database: "d".into(), + gid: "__pgdog_2pc_n_1".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(), + gid: "__pgdog_2pc_n_1".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..0d9f9ba1a 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 @@ -188,11 +188,13 @@ impl Wal { txn: TwoPcTransaction, user: String, database: String, + gid: String, ) -> Result> { self.append(Record::Begin(BeginPayload { txn, user, database, + gid, })) .await } @@ -449,22 +451,26 @@ fn apply_to_snapshot( match record { Record::Begin(p) => { let prior = snapshot.insert( - p.txn, + p.txn.clone(), CheckpointEntry { - txn: p.txn, + txn: p.txn.clone(), user: p.user.clone(), database: p.database.clone(), decided: false, + gid: p.gid.clone(), }, ); - undo.push(Undo { txn: p.txn, prior }); + undo.push(Undo { + txn: p.txn.clone(), + prior, + }); } Record::Committing(p) => { if let Some(entry) = snapshot.get_mut(&p.txn) { let prior = entry.clone(); entry.decided = true; undo.push(Undo { - txn: p.txn, + txn: p.txn.clone(), prior: Some(prior), }); } @@ -472,7 +478,7 @@ fn apply_to_snapshot( Record::End(p) => { if let Some(prior) = snapshot.remove(&p.txn) { undo.push(Undo { - txn: p.txn, + txn: p.txn.clone(), prior: Some(prior), }); } diff --git a/pgdog/src/util.rs b/pgdog/src/util.rs index 8ea220448..6c7e3356b 100644 --- a/pgdog/src/util.rs +++ b/pgdog/src/util.rs @@ -144,8 +144,21 @@ pub fn instance_id() -> &'static str { /// - /// pub fn node_id() -> Result { - // split always returns at least one element. - instance_id().split("-").last().unwrap().parse() + parse_node_id(instance_id()) +} + +/// Parse the numeric node id out of an instance id of the form +/// `-`. Kept separate from +/// [`node_id`] so it can be tested with fixed inputs: `node_id` derives +/// from the process-global `INSTANCE_ID`, whose random hex form parses as +/// a valid number often enough to make direct tests flaky. +fn parse_node_id(instance_id: &str) -> Result { + // rsplit always yields at least one element, so next() is never None. + instance_id + .rsplit('-') + .next() + .unwrap_or(instance_id) + .parse() } static DEPLOYMENT_ID: Lazy> = Lazy::new(|| env::var("DEPLOYMENT_ID").ok()); @@ -408,8 +421,12 @@ mod test { #[test] fn test_node_id_error() { - let _guard = remove_env_var("NODE_ID"); - assert!(node_id().is_err()); + // Test the parser directly with a fixed non-numeric trailing + // segment. Going through node_id() would read the random global + // INSTANCE_ID, which parses as a valid number ~2% of the time + // (all-digit hex) and makes this assertion flaky. + assert!(parse_node_id("host-abc").is_err()); + assert!(parse_node_id("abcdef12").is_err()); } #[test]