From eaf8a4bcd3ee239078a3896c99fb02b525fe3bf5 Mon Sep 17 00:00:00 2001 From: Cesar Rodas Date: Fri, 24 Jul 2026 01:07:17 -0300 Subject: [PATCH] Persist coordinator gid in the WAL for crash recovery Attempt to fix #1175 After a SIGTERM and reboot, prepared 2PC transactions were left orphaned on the shards even with the WAL enabled. Recovery replayed the in-flight transactions but drove COMMIT PREPARED / ROLLBACK PREPARED against a gid that no longer matched what Postgres held, so the statements failed and the monitor retried them forever. The gid a transaction is prepared with embeds a per-process instance_id, which is a fresh random value on every start unless NODE_ID is set. The WAL only persisted the inner transaction id, so a restarted PgDog reconstructed the gid with a different instance_id prefix and missed the orphan. The gid used at recovery must be byte-identical to the gid used at PREPARE, so it has to come from the WAL, not from live process state. Store the full coordinator gid in the Begin and Checkpoint records (serde default keeps old segments readable) and thread it through recovery so cleanup resolves each prepared xact with the exact gid it was created with. Live transactions still render the gid from process state, which is correct in the process that created them. An empty stored gid falls back to the previous behavior. The crash-safety WAL helper now records the full gid too, so the integration spec exercises the real recovery path. --- pgdog/src/backend/pool/connection/binding.rs | 4 +- .../backend/pool/connection/binding_test.rs | 50 +++++---- .../src/backend/replication/logical/error.rs | 2 +- .../replication/logical/subscriber/copy.rs | 12 +-- .../client/query_engine/end_transaction.rs | 4 +- .../client/query_engine/two_pc/manager.rs | 16 +-- .../client/query_engine/two_pc/mod.rs | 8 +- .../client/query_engine/two_pc/statement.rs | 31 ++++-- .../client/query_engine/two_pc/test.rs | 4 +- .../client/query_engine/two_pc/transaction.rs | 101 ++++++++++++++++-- .../two_pc/wal/record/identity.rs | 69 +++++++++++- .../query_engine/two_pc/wal/record/phase.rs | 4 +- .../query_engine/two_pc/wal/record/records.rs | 7 +- .../query_engine/two_pc/wal/record/remove.rs | 4 +- .../query_engine/two_pc/wal/tests/client.rs | 8 +- pgdog/src/util.rs | 25 ++++- 16 files changed, 267 insertions(+), 82 deletions(-) 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]