diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index 1acf1b045..97e1fd236 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -365,7 +365,7 @@ impl Binding { pub(crate) async fn two_pc_on_guards( servers: &mut [Guard], - transaction: TwoPcTransaction, + transaction: &TwoPcTransaction, phase: TwoPcPhase, ignore_missing: bool, ) -> Result<(), Error> { @@ -399,7 +399,7 @@ impl Binding { /// Execute two-phase commit transaction control statements. pub(crate) async fn two_pc( &mut self, - transaction: TwoPcTransaction, + transaction: &TwoPcTransaction, phase: TwoPcPhase, ignore_missing: bool, ) -> Result<(), Error> { diff --git a/pgdog/src/backend/pool/connection/binding_test.rs b/pgdog/src/backend/pool/connection/binding_test.rs index b9f541a78..96f91cbac 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, false) + .two_pc(&TwoPcTransaction::new(), TwoPcPhase::Phase1, false) .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, false) + .two_pc(&TwoPcTransaction::new(), TwoPcPhase::Phase1, false) .await; // Should fail with TwoPcMultiShardOnly error @@ -113,7 +113,9 @@ mod tests { let transaction = TwoPcTransaction::new(); // Test Phase1 - PREPARE TRANSACTION - let result = binding.two_pc(transaction, TwoPcPhase::Phase1, false).await; + let result = binding + .two_pc(&transaction, TwoPcPhase::Phase1, false) + .await; // Should succeed if let Err(ref error) = result { @@ -123,7 +125,7 @@ mod tests { // Cleanup: Rollback the prepared transaction to avoid leaving dangling transactions let _cleanup = binding - .two_pc(transaction, TwoPcPhase::Rollback, false) + .two_pc(&transaction, TwoPcPhase::Rollback, false) .await; } @@ -134,12 +136,14 @@ mod tests { // First prepare the transaction binding - .two_pc(transaction, TwoPcPhase::Phase1, false) + .two_pc(&transaction, TwoPcPhase::Phase1, false) .await .expect("Phase1 should succeed"); // Then commit it - let result = binding.two_pc(transaction, TwoPcPhase::Phase2, false).await; + let result = binding + .two_pc(&transaction, TwoPcPhase::Phase2, false) + .await; assert!(result.is_ok()); } @@ -150,13 +154,13 @@ mod tests { // First prepare the transaction binding - .two_pc(transaction, TwoPcPhase::Phase1, false) + .two_pc(&transaction, TwoPcPhase::Phase1, false) .await .expect("Phase1 should succeed"); // Then rollback let result = binding - .two_pc(transaction, TwoPcPhase::Rollback, false) + .two_pc(&transaction, TwoPcPhase::Rollback, false) .await; assert!(result.is_ok()); } @@ -168,17 +172,17 @@ mod tests { // First prepare the transaction binding - .two_pc(transaction, TwoPcPhase::Phase1, false) + .two_pc(&transaction, TwoPcPhase::Phase1, false) .await .expect("Phase1 should succeed"); // Then commit it binding - .two_pc(transaction, TwoPcPhase::Phase2, true) + .two_pc(&transaction, TwoPcPhase::Phase2, true) .await .expect("Phase2 should succeed"); - let result = binding.two_pc(transaction, TwoPcPhase::Phase2, true).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase2, true).await; assert!( result.is_ok(), "Committing non-existent prepared transaction should be skipped" @@ -192,19 +196,19 @@ mod tests { // First prepare the transaction binding - .two_pc(transaction, TwoPcPhase::Phase1, true) + .two_pc(&transaction, TwoPcPhase::Phase1, true) .await .expect("Phase1 should succeed"); // Then rollback it binding - .two_pc(transaction, TwoPcPhase::Rollback, true) + .two_pc(&transaction, TwoPcPhase::Rollback, true) .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, true) + .two_pc(&transaction, TwoPcPhase::Rollback, true) .await; assert!( result.is_ok(), @@ -221,17 +225,19 @@ mod tests { let transaction = TwoPcTransaction::new(); // 1. Prepare transaction - let result = binding.two_pc(transaction, TwoPcPhase::Phase1, false).await; + let result = binding + .two_pc(&transaction, TwoPcPhase::Phase1, false) + .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, true).await; + let _result = binding.two_pc(&transaction, TwoPcPhase::Phase1, true).await; // 3. Commit the prepared transaction - let result = binding.two_pc(transaction, TwoPcPhase::Phase2, true).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase2, true).await; assert!(result.is_ok(), "Phase2 commit should succeed"); - let result = binding.two_pc(transaction, TwoPcPhase::Phase2, true).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase2, true).await; assert!( result.is_ok(), "Committing non-existent transaction should be skipped" @@ -244,17 +250,19 @@ mod tests { let transaction = TwoPcTransaction::new(); // 1. Prepare transaction - let result = binding.two_pc(transaction, TwoPcPhase::Phase1, false).await; + let result = binding + .two_pc(&transaction, TwoPcPhase::Phase1, false) + .await; assert!(result.is_ok(), "Phase1 preparation should succeed"); // 2. Rollback the prepared transaction let result = binding - .two_pc(transaction, TwoPcPhase::Rollback, true) + .two_pc(&transaction, TwoPcPhase::Rollback, true) .await; assert!(result.is_ok(), "Rollback should succeed"); // 3. Try to commit after rollback - let result = binding.two_pc(transaction, TwoPcPhase::Phase2, true).await; + let result = binding.two_pc(&transaction, TwoPcPhase::Phase2, true).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 92e30b3ac..75d20ae33 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -296,16 +296,16 @@ 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?; + manager.done(txn.clone()).await?; Ok(()) } .await @@ -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 d083dffe9..2d72d74b8 100644 --- a/pgdog/src/frontend/client/query_engine/end_transaction.rs +++ b/pgdog/src/frontend/client/query_engine/end_transaction.rs @@ -111,7 +111,7 @@ 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, false) + .two_pc(&transaction, TwoPcPhase::Phase1, false) .await?; debug!("[2pc] phase 1 complete"); @@ -119,7 +119,7 @@ impl QueryEngine { // 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, false) + .two_pc(&transaction, TwoPcPhase::Phase2, false) .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 4ac576b98..0b9bd97cf 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs @@ -119,7 +119,7 @@ impl Manager { /// Two-pc transaction finished. pub(crate) async fn done(&self, transaction: TwoPcTransaction) -> Result<(), Error> { - if self.remove(transaction).is_some() + if self.remove(transaction.clone()).is_some() && let Some(wal) = self.wal.load_full() { wal.add(TwoPcRecordRemove { transaction }).await?; @@ -174,17 +174,17 @@ impl Manager { TwoPcPhase::Rollback, "rollback is derived during recovery and is never written to the WAL" ); - self.set_transaction_state(transaction, identifier, phase); + self.set_transaction_state(transaction.clone(), identifier, phase); if let Some(wal) = self.wal.load_full() { if phase == TwoPcPhase::Phase1 { wal.add(TwoPcRecordIdentity { - transaction, + transaction: transaction.clone(), identifier: identifier.clone(), }) .await?; } else { - wal.add(TwoPcRecordPhase::new(transaction)).await?; + wal.add(TwoPcRecordPhase::new(transaction.clone())).await?; } } @@ -267,7 +267,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(); } } @@ -302,7 +302,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,10 +338,10 @@ impl Manager { } /// Reconnect to cluster if available and close the two-phase transaction. - async fn cleanup_phase(&self, transaction: TwoPcTransaction) -> Result<(), Error> { + async fn cleanup_phase(&self, transaction: &TwoPcTransaction) -> Result<(), Error> { let (state, in_recovery) = { let guard = self.inner.lock(); - let state = guard.transactions.get(&transaction).cloned(); + let state = guard.transactions.get(transaction).cloned(); if let Some(state) = state { (state, guard.in_recovery) 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 3a9de39fc..c61fa372f 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..9edb75b33 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() } } @@ -48,11 +48,11 @@ impl FromStr for TwoPcTransactionOnShard { /// Build `PREPARE TRANSACTION`, `COMMIT PREPARED`, or `ROLLBACK PREPARED` /// for a shard participant. pub(crate) fn phase_control( - transaction: TwoPcTransaction, + transaction: &TwoPcTransaction, shard: usize, phase: TwoPcPhase, ) -> String { - let txn = TwoPcTransactionOnShard::new(transaction, shard); + let txn = TwoPcTransactionOnShard::new(transaction.clone(), shard); match phase { TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{txn}'"), @@ -70,11 +70,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 +106,31 @@ 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, with the per-shard + // suffix appended. + 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 8a15e0db0..a23417bb1 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, false) + conn.two_pc(&transaction, TwoPcPhase::Phase1, false) .await .unwrap(); @@ -110,7 +110,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, false) + conn.two_pc(&transaction, TwoPcPhase::Phase1, false) .await .unwrap(); 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 2a9f39cce..1b0b82132 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,33 @@ use rand::{Rng, rng}; -use serde::{Deserialize, Serialize}; -use std::{fmt::Display, str::FromStr}; +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(pub(crate) 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 +35,29 @@ 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, + } + } + + /// Rebuild a transaction from the raw id stored in the WAL. The gid is + /// reattached separately via [`Self::with_gid`] when known. + pub(crate) fn from_id(id: usize) -> Self { + Self { id, gid: None } + } + + /// Raw id, as persisted in the WAL record. + pub(crate) fn id(&self) -> usize { + self.id + } + + /// 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 +77,24 @@ 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); } } @@ -44,7 +105,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 +123,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 +161,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/record/identity.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/record/identity.rs index 0adf45eba..2a2b83775 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/record/identity.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/record/identity.rs @@ -19,7 +19,14 @@ use crate::net::{Payload, c_string_buf}; /// | tid | u64 | 8 | /// | user | string | variable | /// | database | string | variable | +/// | gid | string | variable | /// +/// `gid` is the full coordinator gid the transaction was prepared with. It +/// embeds this process's `instance_id`, which a restarted PgDog does not +/// reproduce, so recovery must drive `COMMIT PREPARED` / `ROLLBACK PREPARED` +/// with this stored value rather than re-rendering it. Records written before +/// gid persistence have no trailing `gid`; the transaction then renders its +/// gid live, which is correct only in the process that created it. #[derive(Debug, Clone, PartialEq)] pub(crate) struct TwoPcRecordIdentity { pub(crate) transaction: TwoPcTransaction, @@ -29,9 +36,10 @@ pub(crate) struct TwoPcRecordIdentity { impl From for Record { fn from(value: TwoPcRecordIdentity) -> Self { let mut payload = Payload::raw(); - payload.put_u64(value.transaction.0 as u64); + payload.put_u64(value.transaction.id() as u64); payload.put_string(&value.identifier.user); payload.put_string(&value.identifier.database); + payload.put_string(&value.transaction.to_string()); Record { code: 'i', @@ -44,9 +52,17 @@ impl TryFrom for TwoPcRecordIdentity { type Error = (); fn try_from(mut value: Record) -> Result { - let transaction = TwoPcTransaction(value.data.get_u64() as usize); + let tid = value.data.get_u64() as usize; let user = c_string_buf(&mut value.data); let database = c_string_buf(&mut value.data); + // Records written before gid persistence stop here; c_string_buf + // returns "" on the exhausted buffer and we leave the gid unset. + let gid = c_string_buf(&mut value.data); + let transaction = if gid.is_empty() { + TwoPcTransaction::from_id(tid) + } else { + TwoPcTransaction::from_id(tid).with_gid(gid) + }; Ok(Self { transaction, @@ -54,3 +70,52 @@ impl TryFrom for TwoPcRecordIdentity { }) } } + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn identity_round_trips_gid() { + // Encoding then decoding an identity record must preserve the exact + // coordinator gid, so recovery drives COMMIT/ROLLBACK PREPARED with + // the name Postgres holds even after a restart with a new instance_id. + let transaction = TwoPcTransaction::new(); + let gid = transaction.to_string(); + let record = Record::from(TwoPcRecordIdentity { + transaction, + identifier: Arc::new(User { + user: "alice".into(), + database: "shop".into(), + }), + }); + + let decoded = TwoPcRecordIdentity::try_from(record).unwrap(); + assert_eq!(decoded.transaction.to_string(), gid); + assert_eq!(decoded.identifier.user, "alice"); + assert_eq!(decoded.identifier.database, "shop"); + } + + #[test] + fn identity_without_gid_renders_live() { + // A record written before gid persistence carries only tid/user/ + // database. Decoding must leave the gid unset so the transaction + // renders from live process state (no worse than before the field + // existed). + let tid = 4242usize; + let mut payload = Payload::raw(); + payload.put_u64(tid as u64); + payload.put_string("u"); + payload.put_string("d"); + let record = Record { + code: 'i', + data: payload.freeze(), + }; + + let decoded = TwoPcRecordIdentity::try_from(record).unwrap(); + assert_eq!( + decoded.transaction.to_string(), + TwoPcTransaction::from_id(tid).to_string() + ); + } +} diff --git a/pgdog/src/frontend/client/query_engine/two_pc/wal/record/phase.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/record/phase.rs index 8ea507214..1bec855ec 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/record/phase.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/record/phase.rs @@ -26,7 +26,7 @@ impl TwoPcRecordPhase { impl From for Record { fn from(value: TwoPcRecordPhase) -> Self { let mut payload = Payload::raw(); - payload.put_u64(value.transaction.0 as u64); + payload.put_u64(value.transaction.id() as u64); Self { code: '2', @@ -46,7 +46,7 @@ impl TryFrom for TwoPcRecordPhase { if value.code != '2' { return Err(()); } - let transaction = TwoPcTransaction(value.data.get_u64() as usize); + let transaction = TwoPcTransaction::from_id(value.data.get_u64() as usize); Ok(Self { transaction }) } diff --git a/pgdog/src/frontend/client/query_engine/two_pc/wal/record/records.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/record/records.rs index 5f725da17..60de55131 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/record/records.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/record/records.rs @@ -21,13 +21,14 @@ impl Records { pub(crate) fn replay(&self, manager: &Manager) { match self { Records::Identity(identity) => { - manager.set_transaction_identity(identity.transaction, &identity.identifier); + manager + .set_transaction_identity(identity.transaction.clone(), &identity.identifier); } Records::Phase(phase) => { - manager.set_transaction_phase(phase.transaction, TwoPcPhase::Phase2); + manager.set_transaction_phase(phase.transaction.clone(), TwoPcPhase::Phase2); } Records::Remove(remove) => { - manager.remove(remove.transaction); + manager.remove(remove.transaction.clone()); } } } diff --git a/pgdog/src/frontend/client/query_engine/two_pc/wal/record/remove.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/record/remove.rs index 69f28a414..5fbc5513f 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/record/remove.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/record/remove.rs @@ -15,7 +15,7 @@ impl TryFrom for TwoPcRecordRemove { let tid = value.data.get_u64() as usize; Ok(TwoPcRecordRemove { - transaction: TwoPcTransaction(tid), + transaction: TwoPcTransaction::from_id(tid), }) } } @@ -23,7 +23,7 @@ impl TryFrom for TwoPcRecordRemove { impl From for Record { fn from(value: TwoPcRecordRemove) -> Self { let mut payload = Payload::raw(); - payload.put_u64(value.transaction.0 as u64); + payload.put_u64(value.transaction.id() as u64); Record { code: 'r', diff --git a/pgdog/src/frontend/client/query_engine/two_pc/wal/tests/client.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/tests/client.rs index a69ca7382..37b0649bc 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/wal/tests/client.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/wal/tests/client.rs @@ -41,24 +41,24 @@ impl TwoPcTestClient { let _guard = self .manager - .transaction_state(txn, &self.cluster.identifier(), TwoPcPhase::Phase1) + .transaction_state(txn.clone(), &self.cluster.identifier(), TwoPcPhase::Phase1) .await .unwrap(); for (shard, conn) in conns.iter_mut().enumerate() { - conn.execute(phase_control(txn, shard, TwoPcPhase::Phase1)) + conn.execute(phase_control(&txn, shard, TwoPcPhase::Phase1)) .await .unwrap(); } let _guard = self .manager - .transaction_state(txn, &self.cluster.identifier(), TwoPcPhase::Phase2) + .transaction_state(txn.clone(), &self.cluster.identifier(), TwoPcPhase::Phase2) .await .unwrap(); for (shard, conn) in conns.iter_mut().enumerate() { - conn.execute(phase_control(txn, shard, TwoPcPhase::Phase2)) + conn.execute(phase_control(&txn, shard, TwoPcPhase::Phase2)) .await .unwrap(); } diff --git a/pgdog/src/util.rs b/pgdog/src/util.rs index 02b2f05c2..85864c568 100644 --- a/pgdog/src/util.rs +++ b/pgdog/src/util.rs @@ -146,8 +146,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()); @@ -528,8 +541,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]