feat: allow source and destination clusters to overlap

pgdog/src/backend/pool/address.rs

@@ -197,6 +197,14 @@ impl Address {
         }
     }
 
+    /// Return true if both addresses are pointing to the same physical Postgres
+    /// database.
+    pub(crate) fn same_database(&self, other: &Self) -> bool {
+        self.host == other.host
+            && self.port == other.port
+            && self.database_name == other.database_name
+    }
+
     /// Test convention: `new_test()` represents a primary. Tests that need
     /// a replica do `Address { configured_role: Role::Replica, ..new_test() }`.
     #[cfg(test)]

pgdog/src/backend/pool/cluster.rs

@@ -419,7 +419,7 @@ impl Cluster {
     }
 
     /// Change config to work with logical replication streaming.
-    pub(crate) fn logical_stream(&self) -> Self {
+    pub(crate) fn with_replication_settings_override(&self) -> Self {
         let mut cluster = self.clone();
         // Disable rewrites, we are only sending valid statements.
         cluster.rewrite.enabled = false;

pgdog/src/backend/pool/shard/mod.rs

@@ -10,17 +10,14 @@ use tokio::sync::SetOnce;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, info};
 
-use crate::backend::PubSubListener;
-use crate::backend::Schema;
-use crate::backend::databases::User;
-use crate::backend::pool::lb::ban::Ban;
-use crate::backend::pub_sub::listener::Listener;
-use crate::backend::schema::SchemaCache;
-use crate::config::{LoadBalancingStrategy, ReadWriteSplit, Role};
-use crate::net::Parameters;
-use crate::net::messages::FrontendPid;
-
 use super::{Error, Guard, LoadBalancer, Pool, PoolConfig, Request};
+use crate::backend::pool::Address;
+use crate::backend::{
+    ConnectReason, PubSubListener, Schema, Server, databases::User, pool::lb::ban::Ban,
+    pub_sub::listener::Listener, schema::SchemaCache,
+};
+use crate::config::{LoadBalancingStrategy, ReadWriteSplit, Role};
+use crate::net::{Parameters, messages::FrontendPid};
 
 pub(crate) mod failover_signal;
 pub(crate) mod monitor;
@@ -83,6 +80,27 @@ impl Shard {
         self.lb.get_primary(request).await
     }
 
+    /// Get a standalone (throw-away) connection to the primary database
+    /// of this shard.
+    pub(crate) async fn primary_standalone(&self, reason: ConnectReason) -> Result<Server, Error> {
+        self.lb
+            .primary_target()
+            .ok_or(Error::NoPrimary)?
+            .pool
+            .standalone(reason)
+            .await
+    }
+
+    /// Get the address of the primary database of this shard.
+    pub(crate) fn primary_address(&self) -> Result<&Address, Error> {
+        Ok(self
+            .lb
+            .primary_target()
+            .ok_or(Error::NoPrimary)?
+            .pool
+            .addr())
+    }
+
     /// Get connection to one of the replica databases, using the configured
     /// load balancing algorithm.
     pub(crate) async fn replica(&self, request: &Request) -> Result<Guard, Error> {

pgdog/src/backend/replication/logical/error.rs

@@ -76,6 +76,7 @@ pub(crate) enum Error {
 
     #[error("router: {0}")]
     Router(#[from] crate::frontend::router::Error),
+
     #[error("sharding key lookup failed: {0}")]
     Lookup(String),
 
@@ -106,9 +107,6 @@ pub(crate) enum Error {
     #[error("parse int")]
     ParseInt(#[from] ParseIntError),
 
-    #[error("shard has no primary")]
-    NoPrimary,
-
     #[error("parser: {0}")]
     Parser(#[from] crate::frontend::router::parser::Error),
 
@@ -195,6 +193,12 @@ pub(crate) enum Error {
         #[source]
         source: Box<Error>,
     },
+
+    #[error("source and destination clusters are identical")]
+    SourceDestinationIdentical,
+
+    #[error("destination cluster has no shards")]
+    DestinationNoShards,
 }
 
 impl From<ErrorResponse> for Error {
@@ -231,8 +235,7 @@ impl Error {
             Self::Net(inner) => inner.is_retryable(),
             Self::Pool(inner) => inner.is_retryable(),
             Self::Backend(inner) => inner.is_retryable(),
-            // No connection yet, or primary is down.
-            Self::NotConnected | Self::NoPrimary => true,
+            Self::NotConnected => true,
             // Replication stalled; temporary slot is gone, next attempt starts fresh.
             Self::ReplicationTimeout => true,
             // Postgres sent a transient error (e.g. admin_shutdown, cannot_connect_now).
@@ -266,7 +269,6 @@ mod tests {
         assert!(Error::Pool(PE::NoPrimary).is_retryable());
         assert!(Error::Pool(PE::CheckoutTimeout).is_retryable());
         assert!(Error::NotConnected.is_retryable());
-        assert!(Error::NoPrimary.is_retryable());
         assert!(Error::ReplicationTimeout.is_retryable());
     }
 

pgdog/src/backend/replication/logical/publisher/publisher_impl.rs

@@ -214,7 +214,7 @@ impl Publisher {
             // Each subscriber owns a partition of destination shards for omni-table DML
             // (dest_shard % n_sources == source_shard), preventing cross-subscriber deadlocks.
             let mut stream =
-                StreamSubscriber::new(dest, tables, OmniOwnership::new(number, n_sources));
+                StreamSubscriber::new(source, dest, tables, OmniOwnership::new(number, n_sources));
 
             // Take ownership of the slot for replication.
             let mut slot = self
@@ -813,7 +813,8 @@ mod test {
         let cfg = config();
         let cluster = Cluster::new_test(&cfg);
         cluster.launch();
-        let mut stream = StreamSubscriber::new(&cluster, &[], OmniOwnership::test());
+        let mut stream =
+            StreamSubscriber::new(&Cluster::default(), &cluster, &[], OmniOwnership::test());
         stream.connect().await.unwrap();
 
         let result = stream.handle(begin_copy_data(1)).await;
@@ -833,7 +834,8 @@ mod test {
         let cfg = config();
         let cluster = Cluster::new_test(&cfg);
         cluster.launch();
-        let mut stream = StreamSubscriber::new(&cluster, &[], OmniOwnership::test());
+        let mut stream =
+            StreamSubscriber::new(&Cluster::default(), &cluster, &[], OmniOwnership::test());
         stream.connect().await.unwrap();
 
         let result = stream.handle(commit_copy_data(1)).await;

pgdog/src/backend/replication/logical/subscriber/copy.rs

@@ -3,7 +3,7 @@
 
 use futures::future::join_all;
 use pg_raw_parse::Node;
-use tracing::debug;
+use tracing::{debug, warn};
 
 use crate::frontend::client::query_engine::TwoPcPhase;
 use crate::frontend::client::query_engine::two_pc::{
@@ -13,15 +13,17 @@ use crate::frontend::client::query_engine::two_pc::{
 use crate::frontend::router::parser::Error as ParseError;
 use crate::{
     backend::{Cluster, ConnectReason, replication::subscriber::ParallelConnection},
-    config::Role,
     frontend::router::parser::{CopyParser, Shard},
     net::{
         CopyData, CopyDone, ErrorResponse, FromBytes, Message, Protocol, ProtocolMessage, Query,
         ToBytes,
     },
 };
 
-use super::super::{CopyStatement, Error};
+use super::{
+    super::{CopyStatement, Error},
+    OverlappingShardsCheck,
+};
 
 // Not really needed, but we're currently
 // sharding 3 CopyData messages at a time.
@@ -31,8 +33,8 @@ static BUFFER_SIZE: usize = 3;
 #[derive(Debug)]
 pub(crate) struct CopySubscriber {
     copy: CopyParser,
-    /// Destination cluster.
-    cluster: Cluster,
+    dest: Cluster,
+    source: Cluster,
     buffer: Vec<CopyData>,
     connections: Vec<ParallelConnection>,
     stmt: CopyStatement,
@@ -48,12 +50,12 @@ impl CopySubscriber {
     pub(crate) fn new(
         copy_stmt: &CopyStatement,
         source: &Cluster,
-        cluster: &Cluster,
+        dest: &Cluster,
     ) -> Result<Self, Error> {
         let ast = pg_raw_parse::parse(&copy_stmt.copy_in()).map_err(ParseError::from)?;
         let stmt = ast.stmts().next().ok_or(ParseError::EmptyQuery)?;
         let mut copy = if let Node::CopyStmt(stmt) = stmt {
-            CopyParser::new(stmt, cluster).map_err(|_| Error::MissingData)?
+            CopyParser::new(stmt, dest).map_err(|_| Error::MissingData)?
         } else {
             return Err(Error::MissingData);
         };
@@ -63,7 +65,8 @@ impl CopySubscriber {
 
         Ok(Self {
             copy,
-            cluster: cluster.clone(),
+            dest: dest.clone(),
+            source: source.clone(),
             buffer: vec![],
             connections: vec![],
             stmt: copy_stmt.clone(),
@@ -73,20 +76,32 @@ impl CopySubscriber {
 
     /// Connect to all shards. One connection per primary.
     pub(crate) async fn connect(&mut self) -> Result<(), Error> {
-        let mut servers = vec![];
-        for shard in self.cluster.shards() {
-            let primary = shard
-                .pools_with_roles()
-                .iter()
-                .find(|(role, _)| role == &Role::Primary)
-                .ok_or(Error::NoPrimary)?
-                .1
-                .standalone(ConnectReason::Replication)
-                .await?;
-            servers.push(ParallelConnection::new(primary)?);
+        let mut connections = vec![];
+        let overlap_check = OverlappingShardsCheck::new(&self.source);
+        let destination_shards = self.dest.shards();
+
+        if destination_shards.is_empty() {
+            return Err(Error::DestinationNoShards);
         }
 
-        self.connections = servers;
+        for (shard_number, shard) in destination_shards.iter().enumerate() {
+            if overlap_check.overlaps(shard)? {
+                warn!(
+                    "skipping data sync to {} because it is part of the source cluster",
+                    shard.primary_address()?,
+                );
+                continue;
+            }
+
+            let primary = shard.primary_standalone(ConnectReason::Replication).await?;
+            connections.push(ParallelConnection::new(primary, shard_number)?);
+        }
+
+        if connections.is_empty() {
+            return Err(Error::SourceDestinationIdentical);
+        }
+
+        self.connections = connections;
 
         Ok(())
     }
@@ -227,13 +242,14 @@ 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.
-        if self.cluster.two_pc_enabled() {
+        if self.dest.two_pc_enabled() {
             self.commit_two_pc().await?;
         } else {
-            for (shard, server) in self.connections.iter_mut().enumerate() {
+            for server in &mut self.connections {
                 if let Err(error) =
                     Self::send_and_confirm(server, Query::new("COMMIT").into()).await
                 {
+                    let shard = server.shard_number();
                     tracing::error!(
                         "COMMIT failed on destination shard {shard} during copy_done: {error}; \
                          shards committed before it stay committed, the rest roll back on \
@@ -250,7 +266,7 @@ impl CopySubscriber {
     async fn commit_two_pc(&mut self) -> Result<(), Error> {
         let manager = Manager::get();
         let txn = TwoPcTransaction::new();
-        let identifier = self.cluster.identifier();
+        let identifier = self.dest.identifier();
 
         async {
             let _guard_phase_1 = manager
@@ -280,14 +296,14 @@ impl CopySubscriber {
     ) -> Result<(), Error> {
         let mut futures = Vec::new();
 
-        for (shard, server) in self.connections.iter_mut().enumerate() {
+        for server in &mut self.connections {
             // 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, shard, phase),
+                phase => phase_control(txn, server.shard_number(), phase),
             };
             futures.push(Self::send_and_confirm(server, Query::new(query).into()));
         }
@@ -317,16 +333,16 @@ impl CopySubscriber {
         let bytes = result.iter().map(|row| row.len()).sum::<usize>();
 
         for row in &result {
-            for (shard, server) in self.connections.iter_mut().enumerate() {
+            for server in &mut self.connections {
                 match row.shard() {
                     Shard::All => server.send_one(&row.message().into()).await?,
                     Shard::Direct(destination) => {
-                        if *destination == shard {
+                        if *destination == server.shard_number() {
                             server.send_one(&row.message().into()).await?;
                         }
                     }
                     Shard::Multi(multi) => {
-                        if multi.contains(&shard) {
+                        if multi.contains(&server.shard_number()) {
                             server.send_one(&row.message().into()).await?;
                         }
                     }
@@ -402,7 +418,7 @@ mod test {
             .await
             .unwrap();
 
-        let mut subscriber = CopySubscriber::new(&copy, &cluster, &cluster).unwrap();
+        let mut subscriber = CopySubscriber::new(&copy, &Cluster::default(), &cluster).unwrap();
         subscriber.start_copy().await.unwrap();
 
         let header = CopyData::new(&Header::default().to_bytes());
@@ -445,7 +461,7 @@ mod test {
         crate::logger();
 
         let server = test_server().await;
-        let mut conn = ParallelConnection::new(server).unwrap();
+        let mut conn = ParallelConnection::new(server, 0).unwrap();
 
         // RAISE WARNING emits a NoticeResponse ('N') before the statement's
         // CommandComplete. Without async-message skipping this was misread as
@@ -468,7 +484,7 @@ mod test {
         crate::logger();
 
         let server = test_server().await;
-        let mut conn = ParallelConnection::new(server).unwrap();
+        let mut conn = ParallelConnection::new(server, 0).unwrap();
 
         // A NoticeResponse precedes the ErrorResponse. The notice is skipped, the
         // error is surfaced, and drain_to_ready consumes the trailing ReadyForQuery

pgdog/src/backend/replication/logical/subscriber/duplicate_check.rs

@@ -0,0 +1,29 @@
+//! Check that the source and destination clusters
+//! don't have overlapping shards.
+
+use super::super::Error;
+use crate::backend::{Cluster, Shard};
+
+pub(super) struct OverlappingShardsCheck<'a> {
+    source: &'a Cluster,
+}
+
+impl<'a> OverlappingShardsCheck<'a> {
+    /// Create check.
+    pub(super) fn new(source: &'a Cluster) -> Self {
+        Self { source }
+    }
+
+    /// Check if the destination shard overlaps with any shards in the source cluster.
+    pub(super) fn overlaps(&self, shard: &Shard) -> Result<bool, Error> {
+        let address = shard.primary_address()?;
+
+        for source_shard in self.source.shards() {
+            if source_shard.primary_address()?.same_database(address) {
+                return Ok(true);
+            }
+        }
+
+        Ok(false)
+    }
+}

pgdog/src/backend/replication/logical/subscriber/mod.rs

@@ -1,5 +1,6 @@
 pub(crate) mod context;
 pub(crate) mod copy;
+pub(crate) mod duplicate_check;
 pub(crate) mod omni_ownership;
 pub(crate) mod parallel_connection;
 pub(crate) mod pipeline;
@@ -10,5 +11,6 @@ mod tests;
 
 pub(crate) use context::StreamContext;
 pub(crate) use copy::CopySubscriber;
+use duplicate_check::OverlappingShardsCheck;
 pub(crate) use parallel_connection::ParallelConnection;
 pub(crate) use pipeline::PipelinedConnection;