fix(2pc): clean up prepared transactions recorded in the WAL after a restart

integration/two_pc_crash_safety/wal_helper/Cargo.toml

@@ -10,5 +10,5 @@ path = "src/main.rs"
 
 [dependencies]
 bytes = "1"
-pgdog = { path = "../../../pgdog", default-features = false }
+pgdog = { path = "../../../pgdog" }
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }

integration/two_pc_crash_safety/wal_helper/src/main.rs

@@ -33,10 +33,13 @@ async fn main() {
     let mut segment = Segment::create(&dir, 0).await.expect("create segment");
 
     let mut buf = BytesMut::new();
+    // The prefix is left empty so recovery takes the numeric-ID
+    // fallback path, keeping that path covered end to end.
     Record::Begin(BeginPayload {
         txn,
         user,
         database,
+        prefix: String::new(),
     })
     .encode(&mut buf)
     .expect("encode begin");

pgdog/src/backend/pool/connection/binding.rs

@@ -5,14 +5,16 @@ use crate::{
         ClientRequest,
         client::query_engine::{
             TwoPcPhase,
-            two_pc::{TwoPcTransaction, statement::phase_control},
+            two_pc::{TwoPcTransaction, TwoPcTransactionOnShard, statement::phase_control},
         },
     },
-    net::{FrontendPid, ProtocolMessage, Query, parameter::Parameters},
+    net::{DataRow, FrontendPid, ProtocolMessage, Query, parameter::Parameters},
     state::State,
+    util::is_safe_identifier,
 };
 
 use futures::future::join_all;
+use tracing::warn;
 
 use super::*;
 
@@ -364,20 +366,25 @@ impl Binding {
 
     pub(crate) async fn two_pc_on_guards(
         servers: &mut [Guard],
+        state: &MultiShard,
         transaction: TwoPcTransaction,
         phase: TwoPcPhase,
     ) -> Result<(), Error> {
         let skip_missing = matches!(phase, TwoPcPhase::Phase2 | TwoPcPhase::Rollback);
 
         let mut futures = Vec::new();
-        for (shard, server) in servers.iter_mut().enumerate() {
+        for (position, server) in servers.iter_mut().enumerate() {
+            // Map positional index to actual shard number.
+            // When only a subset of shards is connected (Shard::Multi binding),
+            // positional indices don't match actual shard numbers.
+            let shard = state.shard_index(position);
             let query = phase_control(transaction, shard, phase);
             futures.push(server.execute(query));
         }
 
         let results = join_all(futures).await;
 
-        for (shard, result) in results.into_iter().enumerate() {
+        for (position, result) in results.into_iter().enumerate() {
             match result {
                 Err(Error::ExecutionError(err)) => {
                     if !(skip_missing && err.code == "42704") {
@@ -387,7 +394,7 @@ impl Binding {
                 Err(err) => return Err(err),
                 Ok(_) => {
                     if phase == TwoPcPhase::Phase2 {
-                        servers[shard].stats_mut().transaction_2pc();
+                        servers[position].stats_mut().transaction_2pc();
                     }
                 }
             }
@@ -401,10 +408,116 @@ impl Binding {
         &mut self,
         transaction: TwoPcTransaction,
         phase: TwoPcPhase,
+    ) -> Result<(), Error> {
+        match self {
+            Binding::MultiShard(servers, state) => {
+                Self::two_pc_on_guards(servers, state, transaction, phase).await
+            }
+
+            _ => Err(Error::TwoPcMultiShardOnly),
+        }
+    }
+
+    /// Resolve a two-phase transaction on every shard during cleanup
+    /// and crash recovery.
+    ///
+    /// `prefix` is the coordinator GID prefix recorded when the
+    /// transaction was created; combined with the transaction's numeric
+    /// ID it names the exact prepared transaction on each shard, and the
+    /// phase statement runs against the GIDs found in
+    /// `pg_prepared_xacts`. An empty `prefix` means the transaction was
+    /// restored from a WAL record that did not store it; matching then
+    /// falls back to the durable numeric transaction ID and shard index.
+    /// A shard with no matching GID is already resolved and is skipped.
+    pub(crate) async fn two_pc_cleanup(
+        &mut self,
+        transaction: TwoPcTransaction,
+        prefix: &str,
+        phase: TwoPcPhase,
     ) -> Result<(), Error> {
         match self {
             Binding::MultiShard(servers, _) => {
-                Self::two_pc_on_guards(servers, transaction, phase).await
+                let futures = servers
+                    .iter_mut()
+                    .enumerate()
+                    .map(|(shard, server)| async move {
+                        let target = TwoPcTransactionOnShard::new(transaction, shard);
+                        // A GID rendered from the recorded prefix is
+                        // checked against the identifier alphabet before
+                        // it is embedded in a quoted literal; a prefix
+                        // outside the alphabet falls back to numeric-ID
+                        // matching, which performs the same check.
+                        let expected = if prefix.is_empty() {
+                            None
+                        } else {
+                            let gid = target.gid(prefix);
+                            if is_safe_identifier(&gid) {
+                                Some(gid)
+                            } else {
+                                warn!(
+                                    "[2pc] recorded gid {:?} contains characters \
+                                     PgDog never generates; matching by numeric ID",
+                                    gid
+                                );
+                                None
+                            }
+                        };
+                        let rows: Vec<DataRow> = server
+                            .fetch_all(
+                                "SELECT gid FROM pg_prepared_xacts WHERE database = current_database()",
+                            )
+                            .await?;
+
+                        for row in rows {
+                            let Some(gid) = row.get_text(0) else {
+                                continue;
+                            };
+                            let matched = match &expected {
+                                Some(expected) => &gid == expected,
+                                None => target.matches_gid(&gid),
+                            };
+                            if !matched {
+                                continue;
+                            }
+                            // Both match paths guarantee the gid contains
+                            // no characters that need quoting.
+                            let statement = match phase {
+                                TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{gid}'"),
+                                TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{gid}'"),
+                                TwoPcPhase::Phase1 => {
+                                    unreachable!("cleanup resolves transactions; it never prepares")
+                                }
+                            };
+                            match server.execute(&statement).await {
+                                Ok(_) => {
+                                    if phase == TwoPcPhase::Phase2 {
+                                        server.stats_mut().transaction_2pc();
+                                    }
+                                }
+                                // Insufficient privilege: the prepared transaction
+                                // is owned by a role the configured user can't act
+                                // for. Retrying can't succeed until an operator
+                                // intervenes, so leave the transaction to them
+                                // rather than retrying forever.
+                                Err(Error::ExecutionError(err)) if err.code == "42501" => {
+                                    warn!(
+                                        "[2pc] insufficient privilege to run {}; \
+                                         resolve the prepared transaction manually",
+                                        statement
+                                    );
+                                }
+                                Err(err) => return Err(err),
+                            }
+                        }
+
+                        Ok::<(), Error>(())
+                    });
+
+                for result in join_all(futures).await {
+                    result?;
+                }
+
+                Ok(())
             }
 
             _ => Err(Error::TwoPcMultiShardOnly),

pgdog/src/frontend/client/query_engine/two_pc/manager.rs

@@ -205,12 +205,14 @@ impl Manager {
         identifier: &Arc<User>,
         phase: TwoPcPhase,
     ) -> Result<TwoPcGuard, Error> {
+        let prefix = TwoPcTransaction::global_prefix();
         let prior = {
             let mut guard = self.inner.lock();
             let prior = guard.transactions.get(&transaction).cloned();
             let entry = guard.transactions.entry(transaction).or_default();
             entry.identifier = identifier.clone();
             entry.phase = phase;
+            entry.prefix = prefix.clone();
             prior
         };
 
@@ -221,6 +223,7 @@ impl Manager {
                         transaction,
                         identifier.user.clone(),
                         identifier.database.clone(),
+                        prefix,
                     )
                     .await
                 }
@@ -262,14 +265,20 @@ impl Manager {
         transaction: TwoPcTransaction,
         user: String,
         database: String,
+        prefix: String,
         phase: TwoPcPhase,
     ) {
         let identifier = Arc::new(User { user, database });
         {
             let mut guard = self.inner.lock();
-            guard
-                .transactions
-                .insert(transaction, TransactionInfo { phase, identifier });
+            guard.transactions.insert(
+                transaction,
+                TransactionInfo {
+                    phase,
+                    identifier,
+                    prefix,
+                },
+            );
             guard.queue.push_back(transaction);
         }
         self.stats.incr_recovered();
@@ -345,6 +354,13 @@ impl Manager {
     }
 
     /// Reconnect to cluster if available and rollback the two-phase transaction.
+    ///
+    /// The transaction may have been created by an earlier PgDog
+    /// process; the GID prefix recorded when it was created names the
+    /// exact prepared transactions on the shards. Cleanup scans each
+    /// shard's prepared transactions and acts on the GIDs that match,
+    /// falling back to the durable numeric transaction ID for WAL
+    /// records that did not store the prefix.
     async fn cleanup_phase(&self, transaction: TwoPcTransaction) -> Result<(), Error> {
         let state = match self.inner.lock().transactions.get(&transaction).cloned() {
             Some(state) => state,
@@ -379,7 +395,9 @@ impl Manager {
                 &Route::write(ShardWithPriority::new_override_transaction(Shard::All)),
             )
             .await?;
-        connection.two_pc(transaction, phase).await?;
+        connection
+            .two_pc_cleanup(transaction, &state.prefix, phase)
+            .await?;
         connection.disconnect();
 
         Ok(())
@@ -412,6 +430,10 @@ impl Manager {
 pub struct TransactionInfo {
     pub phase: TwoPcPhase,
     pub identifier: Arc<User>,
+    /// Coordinator GID prefix the transaction's prepared names were
+    /// rendered with. Empty when restored from a WAL record written by
+    /// a version that did not store it.
+    pub prefix: String,
 }
 
 #[derive(Default, Debug)]

pgdog/src/frontend/client/query_engine/two_pc/statement.rs

@@ -2,7 +2,10 @@
 
 use std::{fmt::Display, str::FromStr};
 
+use tracing::warn;
+
 use crate::frontend::client::query_engine::two_pc::TwoPcTransaction;
+use crate::util::is_safe_identifier;
 
 use super::TwoPcPhase;
 
@@ -24,6 +27,47 @@ impl TwoPcTransactionOnShard {
     pub(crate) fn transaction(&self) -> TwoPcTransaction {
         self.transaction
     }
+
+    /// The GID this transaction is prepared under on this shard when
+    /// rendered with `prefix`, the coordinator GID prefix recorded at
+    /// transaction creation.
+    pub(crate) fn gid(&self, prefix: &str) -> String {
+        format!("{}{}_{}", prefix, self.transaction.number(), self.shard)
+    }
+
+    /// Whether `gid`, as listed in `pg_prepared_xacts`, refers to this
+    /// transaction on this shard. GID prefixes embed the identity of the
+    /// PgDog process that created the transaction, which can change
+    /// across restarts, so matching uses the durable numeric transaction
+    /// ID and the shard index. This is the fallback for transactions
+    /// restored from WAL records that did not store the coordinator GID
+    /// prefix, and for recorded prefixes that fail the identifier
+    /// alphabet check; with a recorded prefix, cleanup matches the exact
+    /// GID from [`Self::gid`] instead.
+    ///
+    /// A matched GID is guaranteed to contain only [`is_safe_identifier`]
+    /// characters, so it can be embedded verbatim in a quoted SQL literal.
+    /// PgDog only generates such GIDs (`NODE_ID` and `DEPLOYMENT_ID` are
+    /// validated at startup); a GID that matches the numeric ID but not
+    /// the alphabet was not created by PgDog and is refused with a
+    /// warning.
+    pub(crate) fn matches_gid(&self, gid: &str) -> bool {
+        let matches = match gid.parse::<Self>() {
+            Ok(parsed) => parsed.transaction == self.transaction && parsed.shard == self.shard,
+            Err(()) => false,
+        };
+
+        if matches && !is_safe_identifier(gid) {
+            warn!(
+                "[2pc] prepared transaction {:?} matches transaction {} on shard {} \
+                 but contains characters PgDog never generates; refusing to resolve it",
+                gid, self.transaction, self.shard
+            );
+            return false;
+        }
+
+        matches
+    }
 }
 
 impl Display for TwoPcTransactionOnShard {
@@ -101,6 +145,47 @@ mod test {
         );
     }
 
+    #[test]
+    fn matches_gid_ignores_prefix() {
+        let transaction: TwoPcTransaction = "__pgdog_2pc_123".parse().unwrap();
+        let target = TwoPcTransactionOnShard::new(transaction, 1);
+
+        // Rendered by this process.
+        assert!(target.matches_gid(&target.to_string()));
+        // Rendered by a process with a different instance ID.
+        assert!(target.matches_gid("__pgdog_2pc_deadbeef_123_1"));
+        // Rendered by a process with a deployment ID.
+        assert!(target.matches_gid("__pgdog_2pc_prod_deadbeef_123_1"));
+
+        // Wrong shard.
+        assert!(!target.matches_gid("__pgdog_2pc_deadbeef_123_0"));
+        // Wrong transaction ID.
+        assert!(!target.matches_gid("__pgdog_2pc_deadbeef_9123_1"));
+        // Prepared transactions from other applications.
+        assert!(!target.matches_gid("app_txn_123_1"));
+        assert!(!target.matches_gid("123_1"));
+        assert!(!target.matches_gid(""));
+        // Matching ID but characters PgDog never generates: such GIDs
+        // are refused so they can never reach a quoted SQL literal.
+        assert!(!target.matches_gid("__pgdog_2pc_it's_123_1"));
+        assert!(!target.matches_gid("__pgdog_2pc_a\\'b_123_1"));
+        assert!(!target.matches_gid("__pgdog_2pc_a b_123_1"));
+    }
+
+    #[test]
+    fn gid_renders_recorded_prefix() {
+        let transaction: TwoPcTransaction = "__pgdog_2pc_123".parse().unwrap();
+        let target = TwoPcTransactionOnShard::new(transaction, 1);
+
+        assert_eq!(
+            target.gid("__pgdog_2pc_prod_deadbeef_"),
+            "__pgdog_2pc_prod_deadbeef_123_1"
+        );
+        // The recorded prefix names one exact GID; the same number
+        // under a different prefix is a different transaction.
+        assert_ne!(target.gid("__pgdog_2pc_a_"), target.gid("__pgdog_2pc_b_"));
+    }
+
     #[test]
     fn phase_control_statements() {
         let transaction = TwoPcTransaction::new();