feat: 2pc logic for resharding

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));
         }
 

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<ErrorResponse>),
 
+    #[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<Error>,
+    },
 }
 
 impl From<ErrorResponse> for Error {
@@ -237,9 +248,18 @@ impl From<TableValidationError> for Error {
 }
 
 impl Error {
+    /// Two-phase commit transaction that still needs manager cleanup, if any.
+    pub fn two_pc_cleanup_transaction(&self) -> Option<TwoPcTransaction> {
+        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(),

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

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);

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<User>,

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;

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'"
+        );
+    }
+}

pgdog/src/main.rs

@@ -165,7 +165,11 @@ async fn pgdog(command: Option<Commands>) -> Result<(), Box<dyn std::error::Erro
         Some(ref command) => {
             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);
                 }