From 3090989586b09262820f5e3f373ad876f18d91b3 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Thu, 18 Jun 2026 17:58:56 +0530 Subject: [PATCH 1/7] feat: 2pc logic for resharding --- .../src/backend/replication/logical/error.rs | 3 + .../replication/logical/subscriber/copy.rs | 69 ++++++++++++++++--- .../client/query_engine/two_pc/manager.rs | 4 +- 3 files changed, 66 insertions(+), 10 deletions(-) diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index bc6bead50..355ce8641 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -188,6 +188,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), diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index 0630c1b44..3f41e582b 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -1,10 +1,14 @@ //! 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}; + use crate::{ backend::{Cluster, ConnectReason, replication::subscriber::ParallelConnection}, config::Role, @@ -232,20 +236,69 @@ 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(); + + 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(()) + } + + 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() { + let query = match phase { + TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{txn}_{shard}'"), + TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{txn}_{shard}'"), + _ => unreachable!(), + }; + 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..e2ea72dcc 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs @@ -148,7 +148,7 @@ 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(()) @@ -166,7 +166,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, From e6add871ad5512d7b6ea2474612a06ca7dddea43 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Sun, 21 Jun 2026 17:55:11 +0530 Subject: [PATCH 2/7] fix: comment for rollback --- pgdog/src/backend/replication/logical/subscriber/copy.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index 3f41e582b..eababcbfb 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -287,7 +287,10 @@ impl CopySubscriber { let query = match phase { TwoPcPhase::Phase1 => format!("PREPARE TRANSACTION '{txn}_{shard}'"), TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{txn}_{shard}'"), - _ => unreachable!(), + // 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). + TwoPcPhase::Rollback => unreachable!(), }; futures.push(Self::send_and_confirm(server, Query::new(query).into())); } From d8e60e0a6b8c53e8ffa81d950de39cdd4c84b9b2 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Tue, 23 Jun 2026 19:51:25 +0530 Subject: [PATCH 3/7] feat: transaction naming moved to another file --- pgdog/src/backend/pool/connection/binding.rs | 14 ++---- .../replication/logical/subscriber/copy.rs | 14 +++--- .../client/query_engine/two_pc/mod.rs | 1 + .../client/query_engine/two_pc/statement.rs | 47 +++++++++++++++++++ 4 files changed, 61 insertions(+), 15 deletions(-) create mode 100644 pgdog/src/frontend/client/query_engine/two_pc/statement.rs 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/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index eababcbfb..649f5a50a 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -7,7 +7,9 @@ use pgdog_config::QueryParserEngine; use tracing::debug; use crate::frontend::client::query_engine::TwoPcPhase; -use crate::frontend::client::query_engine::two_pc::{Manager, TwoPcTransaction}; +use crate::frontend::client::query_engine::two_pc::{ + Manager, TwoPcTransaction, statement::phase_control, +}; use crate::{ backend::{Cluster, ConnectReason, replication::subscriber::ParallelConnection}, @@ -284,13 +286,13 @@ impl CopySubscriber { 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::Phase1 => format!("PREPARE TRANSACTION '{txn}_{shard}'"), - TwoPcPhase::Phase2 => format!("COMMIT PREPARED '{txn}_{shard}'"), - // 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). TwoPcPhase::Rollback => unreachable!(), + phase => phase_control(&txn.to_string(), shard, phase), }; futures.push(Self::send_and_confirm(server, Query::new(query).into())); } 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'" + ); + } +} From 0b20dcbafb64f4388dc62d27e6e4427ac567ba17 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Wed, 24 Jun 2026 02:05:59 +0530 Subject: [PATCH 4/7] feat: add cleanup wait before retry --- .../src/backend/replication/logical/error.rs | 17 +++++++++++ .../logical/publisher/parallel_sync.rs | 7 +++++ .../replication/logical/subscriber/copy.rs | 28 +++++++++++-------- .../client/query_engine/two_pc/manager.rs | 12 ++++++++ 4 files changed, 53 insertions(+), 11 deletions(-) diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index 355ce8641..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}, }; @@ -219,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 { @@ -240,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..6f7c1d2eb 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() { + if 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 649f5a50a..4d7f3e707 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/pgdog/src/backend/replication/logical/subscriber/copy.rs @@ -263,19 +263,25 @@ impl CopySubscriber { let txn = TwoPcTransaction::new(); let identifier = self.cluster.identifier(); - 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?; + async { + let _guard_phase_1 = manager + .transaction_state(&txn, &identifier, TwoPcPhase::Phase1) + .await?; + self.two_pc_on_shards(&txn, TwoPcPhase::Phase1).await?; - manager.done(&txn).await?; + let _guard_phase_2 = manager + .transaction_state(&txn, &identifier, TwoPcPhase::Phase2) + .await?; + self.two_pc_on_shards(&txn, TwoPcPhase::Phase2).await?; - Ok(()) + manager.done(&txn).await?; + Ok(()) + } + .await + .map_err(|source| Error::TwoPcCleanupPending { + transaction: txn, + source: Box::new(source), + }) } async fn two_pc_on_shards( 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 e2ea72dcc..ab87b722f 100644 --- a/pgdog/src/frontend/client/query_engine/two_pc/manager.rs +++ b/pgdog/src/frontend/client/query_engine/two_pc/manager.rs @@ -154,6 +154,18 @@ impl Manager { Ok(()) } + /// Block until the monitor has removed this transaction from the manager. + /// + /// No-op if the transaction was never registered or is already gone. + pub async fn wait_until_cleaned_up(&self, transaction: TwoPcTransaction) { + loop { + if !self.inner.lock().transactions.contains_key(&transaction) { + return; + } + self.notify.notify.notified().await; + } + } + /// 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 From 23e79bd6690724b6fe8c274244f1496dd6e70c99 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Thu, 25 Jun 2026 22:26:53 +0530 Subject: [PATCH 5/7] fix clippy errors --- .../replication/logical/publisher/parallel_sync.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pgdog/src/backend/replication/logical/publisher/parallel_sync.rs b/pgdog/src/backend/replication/logical/publisher/parallel_sync.rs index 6f7c1d2eb..8b8f90b21 100644 --- a/pgdog/src/backend/replication/logical/publisher/parallel_sync.rs +++ b/pgdog/src/backend/replication/logical/publisher/parallel_sync.rs @@ -104,10 +104,10 @@ impl ParallelSync { sleep(backoff).await; - if self.dest.two_pc_enabled() { - if let Some(txn) = err.two_pc_cleanup_transaction() { - Manager::get().wait_until_cleaned_up(txn).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 From 54d9cb0800c63912c6bd551dfcf3fc42e46104d7 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Mon, 29 Jun 2026 22:22:58 +0530 Subject: [PATCH 6/7] fix add wait for monitor cleanup and duration for cleaup loop --- .../client/query_engine/two_pc/manager.rs | 21 ++++++++++++++++--- pgdog/src/main.rs | 6 +++++- 2 files changed, 23 insertions(+), 4 deletions(-) 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 ab87b722f..93889e17b 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,7 @@ 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::{ @@ -154,15 +154,30 @@ impl Manager { Ok(()) } - /// Block until the monitor has removed this transaction from the manager. + /// 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; } - self.notify.notify.notified().await; + 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) => {} + } } } 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); } From 90543d3dde93a64f52f143f2d8b95e83feaaf610 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Tue, 30 Jun 2026 16:29:48 +0530 Subject: [PATCH 7/7] fix fmt errors --- pgdog/src/frontend/client/query_engine/two_pc/manager.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 93889e17b..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::{Instant, interval, sleep}}; +use tokio::{ + select, spawn, + sync::Notify, + time::{Instant, interval, sleep}, +}; use tracing::{debug, error, info, warn}; use crate::{