From a263534b18a9c07c8f8f5accd2d8ebc610d64f80 Mon Sep 17 00:00:00 2001 From: Cesar Rodas Date: Fri, 24 Jul 2026 01:07:17 -0300 Subject: [PATCH 1/3] 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. --- .../wal_helper/src/main.rs | 5 +- pgdog/src/backend/pool/connection/binding.rs | 23 +++- .../client/query_engine/two_pc/manager.rs | 31 ++++- .../client/query_engine/two_pc/statement.rs | 61 +++++++++- .../client/query_engine/two_pc/wal/mod.rs | 2 +- .../client/query_engine/two_pc/wal/record.rs | 113 +++++++++++++++++- .../query_engine/two_pc/wal/recovery.rs | 56 ++++++++- .../client/query_engine/two_pc/wal/segment.rs | 4 + .../client/query_engine/two_pc/wal/writer.rs | 3 + 9 files changed, 279 insertions(+), 19 deletions(-) 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..1a04e2a30 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(); @@ -37,6 +37,7 @@ async fn main() { txn, 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..a46bf9fe9 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -5,7 +5,7 @@ use crate::{ ClientRequest, client::query_engine::{ TwoPcPhase, - two_pc::{TwoPcTransaction, statement::phase_control}, + two_pc::{TwoPcTransaction, statement::phase_control_gid}, }, }, net::{FrontendPid, ProtocolMessage, Query, parameter::Parameters}, @@ -364,14 +364,14 @@ impl Binding { pub(crate) async fn two_pc_on_guards( servers: &mut [Guard], - transaction: TwoPcTransaction, + coordinator_gid: &str, phase: TwoPcPhase, ) -> Result<(), Error> { let skip_missing = matches!(phase, TwoPcPhase::Phase2 | TwoPcPhase::Rollback); let mut futures = Vec::new(); for (shard, server) in servers.iter_mut().enumerate() { - let query = phase_control(transaction, shard, phase); + let query = phase_control_gid(coordinator_gid, shard, phase); futures.push(server.execute(query)); } @@ -396,15 +396,28 @@ impl Binding { Ok(()) } - /// Execute two-phase commit transaction control statements. + /// Execute two-phase commit transaction control statements for a + /// live [`TwoPcTransaction`]. The coordinator gid is rendered from the + /// transaction, which is correct in the process that created it. pub(crate) async fn two_pc( &mut self, transaction: TwoPcTransaction, phase: TwoPcPhase, + ) -> Result<(), Error> { + self.two_pc_gid(&transaction.to_string(), phase).await + } + + /// Like [`Self::two_pc`] but drives the statements with an explicit + /// coordinator gid. WAL recovery uses this to resolve a prepared xact + /// with the exact gid it was created with (see [`phase_control_gid`]). + pub(crate) async fn two_pc_gid( + &mut self, + coordinator_gid: &str, + phase: TwoPcPhase, ) -> Result<(), Error> { match self { Binding::MultiShard(servers, _) => { - Self::two_pc_on_guards(servers, transaction, phase).await + Self::two_pc_on_guards(servers, coordinator_gid, phase).await } _ => Err(Error::TwoPcMultiShardOnly), diff --git a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs index ff22ffb23..8d5cd1c3b 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs @@ -221,6 +221,7 @@ impl Manager { transaction, identifier.user.clone(), identifier.database.clone(), + transaction.to_string(), ) .await } @@ -262,14 +263,23 @@ impl Manager { transaction: TwoPcTransaction, user: String, database: String, + gid: String, phase: TwoPcPhase, ) { let identifier = Arc::new(User { user, database }); + // An empty gid means the record predates gid persistence; fall + // back to re-rendering the transaction (no worse than before). + let recovered_gid = if gid.is_empty() { None } else { Some(gid) }; { let mut guard = self.inner.lock(); - guard - .transactions - .insert(transaction, TransactionInfo { phase, identifier }); + guard.transactions.insert( + transaction, + TransactionInfo { + phase, + identifier, + recovered_gid, + }, + ); guard.queue.push_back(transaction); } self.stats.incr_recovered(); @@ -379,7 +389,13 @@ impl Manager { &Route::write(ShardWithPriority::new_override_transaction(Shard::All)), ) .await?; - connection.two_pc(transaction, phase).await?; + // Recovered transactions must be resolved with the gid they were + // prepared with; re-rendering `transaction` here would embed this + // process's fresh instance_id and miss the orphan on Postgres. + match &state.recovered_gid { + Some(gid) => connection.two_pc_gid(gid, phase).await?, + None => connection.two_pc(transaction, phase).await?, + } connection.disconnect(); Ok(()) @@ -412,6 +428,13 @@ impl Manager { pub struct TransactionInfo { pub phase: TwoPcPhase, pub identifier: Arc, + /// Set only for transactions rebuilt by WAL recovery: the exact + /// coordinator gid the xact was prepared with. A restarted PgDog + /// generates a fresh `instance_id`, so `TwoPcTransaction`'s `Display` + /// no longer reproduces the original gid; cleanup must use this + /// stored value. `None` for transactions created in this process, + /// where `Display` is still correct. + pub recovered_gid: Option, } #[derive(Default, Debug)] diff --git a/pgdog/src/frontend/client/query_engine/two_pc/statement.rs b/pgdog/src/frontend/client/query_engine/two_pc/statement.rs index 9c59824e4..c1172866a 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/statement.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/statement.rs @@ -46,18 +46,39 @@ impl FromStr for TwoPcTransactionOnShard { } /// Build `PREPARE TRANSACTION`, `COMMIT PREPARED`, or `ROLLBACK PREPARED` -/// for a shard participant. +/// for a shard participant, deriving the coordinator gid from a live +/// [`TwoPcTransaction`]. pub(crate) fn phase_control( transaction: TwoPcTransaction, shard: usize, phase: TwoPcPhase, ) -> String { - let txn = TwoPcTransactionOnShard::new(transaction, shard); + control_statement( + &TwoPcTransactionOnShard::new(transaction, shard).to_string(), + phase, + ) +} + +/// Same as [`phase_control`] but takes the coordinator gid as an already +/// rendered string. +/// +/// WAL recovery must use this: a prepared xact's gid embeds the +/// `instance_id` of the process that created it, which a restarted PgDog +/// does not reproduce. The gid therefore has to come from the WAL, not +/// from re-rendering a [`TwoPcTransaction`] against live process state. +/// The per-shard participant name is `_`, matching +/// [`TwoPcTransactionOnShard`]'s `Display`. +pub(crate) fn phase_control_gid(coordinator_gid: &str, shard: usize, phase: TwoPcPhase) -> String { + control_statement(&format!("{coordinator_gid}_{shard}"), phase) +} +/// Render the transaction-control statement for an already-resolved +/// per-shard participant name. +fn control_statement(participant_gid: &str, phase: TwoPcPhase) -> 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_gid}'"), + TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{participant_gid}'"), + TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{participant_gid}'"), } } @@ -118,4 +139,34 @@ mod test { format!("ROLLBACK PREPARED '{transaction}_1'") ); } + + #[test] + fn phase_control_gid_matches_live_rendering() { + // For a gid captured while the transaction is live, the gid-based + // path must produce byte-identical statements to the live path. + // This is the invariant WAL recovery relies on across restarts. + let transaction = TwoPcTransaction::new(); + let gid = transaction.to_string(); + + for phase in [TwoPcPhase::Phase1, TwoPcPhase::Phase2, TwoPcPhase::Rollback] { + for shard in [0usize, 1, 7] { + assert_eq!( + phase_control_gid(&gid, shard, phase), + phase_control(transaction, shard, phase), + ); + } + } + } + + #[test] + fn phase_control_gid_uses_stored_gid_verbatim() { + // A gid recovered from the WAL that no longer matches this + // process's rendering (different instance_id) must still be used + // exactly as stored. + let stored = "__pgdog_2pc_oldnode_42"; + assert_eq!( + phase_control_gid(stored, 3, TwoPcPhase::Rollback), + "ROLLBACK PREPARED '__pgdog_2pc_oldnode_42_3'" + ); + } } 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..69f159d5a 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,31 @@ 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(); round_trip(Record::Begin(BeginPayload { - txn: TwoPcTransaction::new(), + txn, user: "alice".into(), database: "shop".into(), + gid: txn.to_string(), })); } @@ -253,12 +284,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 +346,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 +360,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, + 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, + 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..6d2c4c22f 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. @@ -129,9 +133,10 @@ pub(super) async fn recover_transactions( user: entry.user.clone(), database: entry.database.clone(), decided: entry.decided, + gid: entry.gid.clone(), }, ); - manager.restore_transaction(txn, entry.user, entry.database, phase); + manager.restore_transaction(txn, entry.user, entry.database, entry.gid, phase); } if corruption { tracing::warn!( @@ -184,6 +189,7 @@ fn apply(working: &mut HashMap, record: Record) { user: p.user, database: p.database, decided: false, + gid: p.gid, }, ); } @@ -202,6 +208,7 @@ fn apply(working: &mut HashMap, record: Record) { user, database, decided, + gid, .. } in p.active { @@ -211,6 +218,7 @@ fn apply(working: &mut HashMap, record: Record) { user, database, decided, + gid, }, ); } @@ -238,14 +246,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 +303,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 +326,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 +342,7 @@ mod tests { txn: txn(1), user: "u1".into(), database: "d1".into(), + gid: txn(1).to_string(), }), ); apply( @@ -300,6 +351,7 @@ mod tests { txn: txn(2), user: "u2".into(), database: "d2".into(), + gid: txn(2).to_string(), }), ); apply( @@ -310,12 +362,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..1b3d70b1d 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 } @@ -455,6 +457,7 @@ fn apply_to_snapshot( user: p.user.clone(), database: p.database.clone(), decided: false, + gid: p.gid.clone(), }, ); undo.push(Undo { txn: p.txn, prior }); From 98f1df0cb3b64854993fbf43053615ca4398f9f6 Mon Sep 17 00:00:00 2001 From: Cesar Rodas Date: Mon, 27 Jul 2026 15:50:51 -0300 Subject: [PATCH 2/3] Deflake node_id parsing test test_node_id_error asserted node_id() errors when NODE_ID is unset, but node_id() reads the process-global random hex INSTANCE_ID. Roughly 2.3% of the time (all-digit hex) that string parses as a valid u64, so the assertion failed intermittently in CI. Extract the parse into a pure parse_node_id() helper and test it with fixed non-numeric inputs, decoupling the test from the random global. node_id() returns exactly what it did before for every input. --- pgdog/src/util.rs | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) 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] From 953a46236b64ffdba1aca3885e2806435eff0d8f Mon Sep 17 00:00:00 2001 From: Cesar Rodas Date: Mon, 27 Jul 2026 18:13:01 -0300 Subject: [PATCH 3/3] Make TwoPcTransaction carry its own gid Address review feedback on the WAL gid recovery change: the cleanup path passed the coordinator gid around as a bare string, which is easy to misuse and gives no type safety. Make TwoPcTransaction self-describing instead. It gains an optional gid that Display renders verbatim when set (a transaction rebuilt by WAL recovery, whose gid embeds another process's instance_id that this one cannot reproduce) and renders from live process state when absent. Serialization stays a bare integer via a custom impl, so the on-disk WAL format is unchanged and old segments still decode. The driver functions now take a typed TwoPcTransaction end to end. This removes phase_control_gid, Binding::two_pc_gid, and the separate recovered_gid on TransactionInfo. TwoPcTransaction is no longer Copy because it holds an Arc; call sites clone it explicitly. --- .../wal_helper/src/main.rs | 2 +- pgdog/src/backend/pool/connection/binding.rs | 30 ++--- .../backend/pool/connection/binding_test.rs | 42 +++---- .../src/backend/replication/logical/error.rs | 2 +- .../replication/logical/subscriber/copy.rs | 10 +- .../client/query_engine/end_transaction.rs | 8 +- .../client/query_engine/two_pc/manager.rs | 51 +++----- .../client/query_engine/two_pc/mod.rs | 8 +- .../client/query_engine/two_pc/statement.rs | 86 +++++--------- .../client/query_engine/two_pc/test.rs | 4 +- .../client/query_engine/two_pc/transaction.rs | 110 ++++++++++++++++-- .../client/query_engine/two_pc/wal/record.rs | 7 +- .../query_engine/two_pc/wal/recovery.rs | 15 ++- .../client/query_engine/two_pc/wal/writer.rs | 13 ++- 14 files changed, 218 insertions(+), 170 deletions(-) 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 1a04e2a30..f62cdbf57 100644 --- a/integration/two_pc_crash_safety/wal_helper/src/main.rs +++ b/integration/two_pc_crash_safety/wal_helper/src/main.rs @@ -34,7 +34,7 @@ async fn main() { let mut buf = BytesMut::new(); Record::Begin(BeginPayload { - txn, + txn: txn.clone(), user, database, gid, diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index a46bf9fe9..7a77b7e4a 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -5,7 +5,7 @@ use crate::{ ClientRequest, client::query_engine::{ TwoPcPhase, - two_pc::{TwoPcTransaction, statement::phase_control_gid}, + two_pc::{TwoPcTransaction, statement::phase_control}, }, }, net::{FrontendPid, ProtocolMessage, Query, parameter::Parameters}, @@ -364,14 +364,14 @@ impl Binding { pub(crate) async fn two_pc_on_guards( servers: &mut [Guard], - coordinator_gid: &str, + transaction: &TwoPcTransaction, phase: TwoPcPhase, ) -> Result<(), Error> { let skip_missing = matches!(phase, TwoPcPhase::Phase2 | TwoPcPhase::Rollback); let mut futures = Vec::new(); for (shard, server) in servers.iter_mut().enumerate() { - let query = phase_control_gid(coordinator_gid, shard, phase); + let query = phase_control(transaction, shard, phase); futures.push(server.execute(query)); } @@ -396,28 +396,20 @@ impl Binding { Ok(()) } - /// Execute two-phase commit transaction control statements for a - /// live [`TwoPcTransaction`]. The coordinator gid is rendered from the - /// transaction, which is correct in the process that created it. + /// 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, - phase: TwoPcPhase, - ) -> Result<(), Error> { - self.two_pc_gid(&transaction.to_string(), phase).await - } - - /// Like [`Self::two_pc`] but drives the statements with an explicit - /// coordinator gid. WAL recovery uses this to resolve a prepared xact - /// with the exact gid it was created with (see [`phase_control_gid`]). - pub(crate) async fn two_pc_gid( - &mut self, - coordinator_gid: &str, + transaction: &TwoPcTransaction, phase: TwoPcPhase, ) -> Result<(), Error> { match self { Binding::MultiShard(servers, _) => { - Self::two_pc_on_guards(servers, coordinator_gid, phase).await + Self::two_pc_on_guards(servers, transaction, phase).await } _ => Err(Error::TwoPcMultiShardOnly), 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 8d5cd1c3b..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,14 +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") } @@ -234,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); @@ -263,23 +263,14 @@ impl Manager { transaction: TwoPcTransaction, user: String, database: String, - gid: String, phase: TwoPcPhase, ) { let identifier = Arc::new(User { user, database }); - // An empty gid means the record predates gid persistence; fall - // back to re-rendering the transaction (no worse than before). - let recovered_gid = if gid.is_empty() { None } else { Some(gid) }; { let mut guard = self.inner.lock(); - guard.transactions.insert( - transaction, - TransactionInfo { - phase, - identifier, - recovered_gid, - }, - ); + guard + .transactions + .insert(transaction.clone(), TransactionInfo { phase, identifier }); guard.queue.push_back(transaction); } self.stats.incr_recovered(); @@ -294,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(); } } @@ -319,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: {}"#, @@ -348,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(()); @@ -389,13 +380,10 @@ impl Manager { &Route::write(ShardWithPriority::new_override_transaction(Shard::All)), ) .await?; - // Recovered transactions must be resolved with the gid they were - // prepared with; re-rendering `transaction` here would embed this - // process's fresh instance_id and miss the orphan on Postgres. - match &state.recovered_gid { - Some(gid) => connection.two_pc_gid(gid, phase).await?, - None => connection.two_pc(transaction, phase).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(); Ok(()) @@ -428,13 +416,6 @@ impl Manager { pub struct TransactionInfo { pub phase: TwoPcPhase, pub identifier: Arc, - /// Set only for transactions rebuilt by WAL recovery: the exact - /// coordinator gid the xact was prepared with. A restarted PgDog - /// generates a fresh `instance_id`, so `TwoPcTransaction`'s `Display` - /// no longer reproduces the original gid; cleanup must use this - /// stored value. `None` for transactions created in this process, - /// where `Display` is still correct. - pub recovered_gid: Option, } #[derive(Default, Debug)] 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 c1172866a..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() } } @@ -46,39 +46,23 @@ impl FromStr for TwoPcTransactionOnShard { } /// Build `PREPARE TRANSACTION`, `COMMIT PREPARED`, or `ROLLBACK PREPARED` -/// for a shard participant, deriving the coordinator gid from a live -/// [`TwoPcTransaction`]. +/// 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 { - control_statement( - &TwoPcTransactionOnShard::new(transaction, shard).to_string(), - phase, - ) -} - -/// Same as [`phase_control`] but takes the coordinator gid as an already -/// rendered string. -/// -/// WAL recovery must use this: a prepared xact's gid embeds the -/// `instance_id` of the process that created it, which a restarted PgDog -/// does not reproduce. The gid therefore has to come from the WAL, not -/// from re-rendering a [`TwoPcTransaction`] against live process state. -/// The per-shard participant name is `_`, matching -/// [`TwoPcTransactionOnShard`]'s `Display`. -pub(crate) fn phase_control_gid(coordinator_gid: &str, shard: usize, phase: TwoPcPhase) -> String { - control_statement(&format!("{coordinator_gid}_{shard}"), phase) -} - -/// Render the transaction-control statement for an already-resolved -/// per-shard participant name. -fn control_statement(participant_gid: &str, phase: TwoPcPhase) -> String { + let participant = TwoPcTransactionOnShard::new(transaction.clone(), shard).to_string(); match phase { - TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{participant_gid}'"), - TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{participant_gid}'"), - TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{participant_gid}'"), + TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{participant}'"), + TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{participant}'"), + TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{participant}'"), } } @@ -91,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") ); } @@ -127,45 +111,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_gid_matches_live_rendering() { - // For a gid captured while the transaction is live, the gid-based - // path must produce byte-identical statements to the live path. - // This is the invariant WAL recovery relies on across restarts. - let transaction = TwoPcTransaction::new(); - let gid = transaction.to_string(); - - for phase in [TwoPcPhase::Phase1, TwoPcPhase::Phase2, TwoPcPhase::Rollback] { - for shard in [0usize, 1, 7] { - assert_eq!( - phase_control_gid(&gid, shard, phase), - phase_control(transaction, shard, phase), - ); - } - } - } - - #[test] - fn phase_control_gid_uses_stored_gid_verbatim() { - // A gid recovered from the WAL that no longer matches this - // process's rendering (different instance_id) must still be used - // exactly as stored. - let stored = "__pgdog_2pc_oldnode_42"; + 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_gid(stored, 3, TwoPcPhase::Rollback), + 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/record.rs b/pgdog/src/frontend/client/query_engine/two_pc/wal/record.rs index 69f159d5a..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 @@ -253,11 +253,12 @@ mod tests { #[test] fn round_trip_begin() { let txn = TwoPcTransaction::new(); + let gid = txn.to_string(); round_trip(Record::Begin(BeginPayload { txn, user: "alice".into(), database: "shop".into(), - gid: txn.to_string(), + gid, })); } @@ -376,7 +377,7 @@ mod tests { let txn = TwoPcTransaction::new(); let payload = rmp_serde::to_vec_named(&OldBeginPayload { - txn, + txn: txn.clone(), user: "alice".into(), database: "shop".into(), }) @@ -416,7 +417,7 @@ mod tests { let txn = TwoPcTransaction::new(); let payload = rmp_serde::to_vec_named(&OldCheckpointPayload { active: vec![OldCheckpointEntry { - txn, + txn: txn.clone(), user: "u".into(), database: "d".into(), decided: true, 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 6d2c4c22f..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 @@ -127,16 +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, entry.gid, 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!( 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 1b3d70b1d..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 @@ -451,23 +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), }); } @@ -475,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), }); }