diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index 6db7867e9..dc26acdc9 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -1,7 +1,10 @@ //! Binding between frontend client and a connection on the backend. use crate::{ - frontend::{ClientRequest, client::query_engine::TwoPcPhase}, + frontend::{ + ClientRequest, + client::query_engine::{TwoPcPhase, two_pc::statement::phase_control}, + }, net::{FrontendPid, ProtocolMessage, Query, parameter::Parameters}, state::State, }; @@ -374,14 +377,7 @@ impl Binding { let mut futures = Vec::new(); for (shard, server) in servers.iter_mut().enumerate() { - let shard_name = format!("{}_{}", name, shard); - - let query = match phase { - TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{}'", shard_name), - TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{}'", shard_name), - TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{}'", shard_name), - }; - + let query = phase_control(name, shard, phase); futures.push(server.execute(query)); } diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index bc6bead50..3f52c5d19 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -5,6 +5,7 @@ use derive_more::{Display, Error}; use crate::{ backend::replication::publisher::PublicationTable, + frontend::client::query_engine::two_pc::TwoPcTransaction, net::{CommandComplete, ErrorResponse}, }; @@ -188,6 +189,9 @@ pub enum Error { #[error("binary format mismatch (likely int -> bigint), use text copy instead: {0}")] BinaryFormatMismatch(Box), + #[error("frontend: {0}")] + Frontend(#[from] crate::frontend::Error), + #[error("command complete has no rows: {0}")] CommandCompleteNoRows(CommandComplete), @@ -216,6 +220,13 @@ pub enum Error { oid: pgdog_postgres_types::Oid, op: &'static str, }, + + #[error("2pc commit failed for {transaction}: {source}")] + TwoPcCleanupPending { + transaction: TwoPcTransaction, + #[source] + source: Box, + }, } impl From for Error { @@ -237,9 +248,18 @@ impl From for Error { } 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), + _ => None, + } + } + /// Whether the table copy should be retried after this error. pub fn is_retryable(&self) -> bool { match self { + Self::TwoPcCleanupPending { source, .. } => source.is_retryable(), Self::Net(inner) => inner.is_retryable(), Self::Pool(inner) => inner.is_retryable(), Self::Backend(inner) => inner.is_retryable(), diff --git a/pgdog/src/backend/replication/logical/publisher/parallel_sync.rs b/pgdog/src/backend/replication/logical/publisher/parallel_sync.rs index 1d0cf9d39..8b8f90b21 100644 --- a/pgdog/src/backend/replication/logical/publisher/parallel_sync.rs +++ b/pgdog/src/backend/replication/logical/publisher/parallel_sync.rs @@ -23,6 +23,7 @@ use crate::backend::{ pool::{Address, Request}, replication::{publisher::Table, status::TableCopy}, }; +use crate::frontend::client::query_engine::two_pc::Manager; use crate::net::messages::Protocol; use crate::util::escape_identifier; @@ -103,6 +104,12 @@ impl ParallelSync { sleep(backoff).await; + if self.dest.two_pc_enabled() + && let Some(txn) = err.two_pc_cleanup_transaction() + { + Manager::get().wait_until_cleaned_up(txn).await; + } + // Not idempotent (no truncate): if a prior attempt left rows on a // reachable shard, the re-copy can only collide, so stop instead of // retrying. A shard we cannot probe is logged (not blocked on), since we diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index 0630c1b44..4d7f3e707 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -1,10 +1,16 @@ //! Shard COPY stream from one source //! between N shards. +use futures::future::join_all; use pg_query::{NodeEnum, parse_raw}; use pgdog_config::QueryParserEngine; use tracing::debug; +use crate::frontend::client::query_engine::TwoPcPhase; +use crate::frontend::client::query_engine::two_pc::{ + Manager, TwoPcTransaction, statement::phase_control, +}; + use crate::{ backend::{Cluster, ConnectReason, replication::subscriber::ParallelConnection}, config::Role, @@ -232,20 +238,78 @@ impl CopySubscriber { // scope). Shards not yet committed roll back on connection close. The // destination_has_rows() guard in parallel_sync.rs prevents a doomed retry if this // window is ever hit. - for (shard, server) in self.connections.iter_mut().enumerate() { - if let Err(error) = Self::send_and_confirm(server, Query::new("COMMIT").into()).await { - tracing::error!( - "COMMIT failed on destination shard {shard} during copy_done: {error}; \ - shards committed before it stay committed, the rest roll back on \ - connection close" - ); - return Err(error); + if self.cluster.two_pc_enabled() { + self.commit_two_pc().await?; + } else { + for (shard, server) in self.connections.iter_mut().enumerate() { + if let Err(error) = + Self::send_and_confirm(server, Query::new("COMMIT").into()).await + { + tracing::error!( + "COMMIT failed on destination shard {shard} during copy_done: {error}; \ + shards committed before it stay committed, the rest roll back on \ + connection close" + ); + return Err(error); + } } } Ok(()) } + async fn commit_two_pc(&mut self) -> Result<(), Error> { + let manager = Manager::get(); + let txn = TwoPcTransaction::new(); + let identifier = self.cluster.identifier(); + + async { + let _guard_phase_1 = manager + .transaction_state(&txn, &identifier, TwoPcPhase::Phase1) + .await?; + self.two_pc_on_shards(&txn, TwoPcPhase::Phase1).await?; + + let _guard_phase_2 = manager + .transaction_state(&txn, &identifier, TwoPcPhase::Phase2) + .await?; + self.two_pc_on_shards(&txn, TwoPcPhase::Phase2).await?; + + manager.done(&txn).await?; + Ok(()) + } + .await + .map_err(|source| Error::TwoPcCleanupPending { + transaction: txn, + source: Box::new(source), + }) + } + + async fn two_pc_on_shards( + &mut self, + txn: &TwoPcTransaction, + phase: TwoPcPhase, + ) -> Result<(), Error> { + let mut futures = Vec::new(); + + for (shard, server) in self.connections.iter_mut().enumerate() { + // Rollback is not issued here. If this path fails, the TwoPcGuards in + // commit_two_pc() are dropped without manager.done(), and the 2PC Manager + // cleanup task issues ROLLBACK PREPARED (Phase1) or COMMIT PREPARED (Phase2) + // via binding.rs using the same phase_control() helper. + let query = match phase { + TwoPcPhase::Rollback => unreachable!(), + phase => phase_control(&txn.to_string(), shard, phase), + }; + futures.push(Self::send_and_confirm(server, Query::new(query).into())); + } + + for result in join_all(futures).await { + result?; + } + + Ok(()) + } + /// Send data to subscriber, buffered. pub async fn copy_data(&mut self, data: CopyData) -> Result<(usize, usize), Error> { self.buffer.push(data); 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 ebefe574b..2c26b55c5 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs @@ -12,7 +12,11 @@ use std::{ }, time::Duration, }; -use tokio::{select, spawn, sync::Notify, time::interval}; +use tokio::{ + select, spawn, + sync::Notify, + time::{Instant, interval, sleep}, +}; use tracing::{debug, error, info, warn}; use crate::{ @@ -148,12 +152,39 @@ impl Manager { } /// Two-pc transaction finished. - pub(super) async fn done(&self, transaction: &TwoPcTransaction) -> Result<(), Error> { + pub(crate) async fn done(&self, transaction: &TwoPcTransaction) -> Result<(), Error> { self.remove(transaction).await; Ok(()) } + /// Block until the monitor has removed this transaction from the manager, + /// or until a fixed timeout elapses. + /// + /// No-op if the transaction was never registered or is already gone. + pub async fn wait_until_cleaned_up(&self, transaction: TwoPcTransaction) { + const WAIT_TIMEOUT: Duration = Duration::from_secs(10); + + let deadline = Instant::now() + WAIT_TIMEOUT; + loop { + if !self.inner.lock().transactions.contains_key(&transaction) { + return; + } + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + warn!( + "[2pc] timed out waiting for transaction {} cleanup; monitor will retry", + transaction + ); + return; + } + select! { + _ = self.notify.notify.notified() => {} + _ = sleep(remaining) => {} + } + } + } + /// Record a phase transition for a 2PC transaction. Updates the /// in-memory state first, then writes the corresponding WAL record /// (Begin for Phase1, Committing for Phase2). The inner-first @@ -166,7 +197,7 @@ impl Manager { /// caller to issue PREPARE / COMMIT PREPARED to backends without a /// durable record, which is exactly the orphan-prepared-xact case /// the WAL exists to prevent. - pub(super) async fn transaction_state( + pub(crate) async fn transaction_state( &self, transaction: &TwoPcTransaction, identifier: &Arc, 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 b55ac017e..1024d12ab 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/mod.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/mod.rs @@ -9,6 +9,7 @@ pub mod guard; pub mod manager; pub mod phase; pub mod server_transactions; +pub mod statement; pub mod stats; pub mod transaction; pub mod wal; diff --git a/pgdog/src/frontend/client/query_engine/two_pc/statement.rs b/pgdog/src/frontend/client/query_engine/two_pc/statement.rs new file mode 100644 index 000000000..4456db029 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/two_pc/statement.rs @@ -0,0 +1,47 @@ +//! Per-shard two-phase commit transaction names and control statements. + +use super::TwoPcPhase; + +/// Prepared transaction name for a coordinator transaction on one shard. +pub fn shard_name(transaction: &str, shard: usize) -> String { + format!("{transaction}_{shard}") +} + +/// Build `PREPARE TRANSACTION`, `COMMIT PREPARED`, or `ROLLBACK PREPARED` +/// for a shard participant. +pub fn phase_control(transaction: &str, shard: usize, phase: TwoPcPhase) -> String { + let name = shard_name(transaction, shard); + + match phase { + TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{name}'"), + TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{name}'"), + TwoPcPhase::Rollback => format!("ROLLBACK PREPARED '{name}'"), + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn shard_name_appends_index() { + assert_eq!(shard_name("__pgdog_2pc_42", 0), "__pgdog_2pc_42_0"); + assert_eq!(shard_name("test_txn", 3), "test_txn_3"); + } + + #[test] + fn phase_control_statements() { + assert_eq!( + phase_control("test", 1, TwoPcPhase::Phase1), + "PREPARE TRANSACTION 'test_1'" + ); + assert_eq!( + phase_control("test", 1, TwoPcPhase::Phase2), + "COMMIT PREPARED 'test_1'" + ); + assert_eq!( + phase_control("test", 1, TwoPcPhase::Rollback), + "ROLLBACK PREPARED 'test_1'" + ); + } +} diff --git a/pgdog/src/main.rs b/pgdog/src/main.rs index 7b1e6982b..16ccd3ead 100644 --- a/pgdog/src/main.rs +++ b/pgdog/src/main.rs @@ -165,7 +165,11 @@ async fn pgdog(command: Option) -> Result<(), Box { if let Commands::DataSync { .. } = command { info!("🔄 entering data sync mode"); - if let Err(err) = cli::data_sync(command.clone()).await { + let result = cli::data_sync(command.clone()).await; + // Wait for the 2PC monitor to drain any in-flight cleanup + // before the process exits, even on error. + Manager::get().shutdown().await; + if let Err(err) = result { error!("{}", err); return Err(err); }