diff --git a/pgdog/src/backend/pool/address.rs b/pgdog/src/backend/pool/address.rs index 197a3f110..6a38d558e 100644 --- a/pgdog/src/backend/pool/address.rs +++ b/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)] diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 53fd035f0..c2c6b83e8 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/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; diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index 341481612..d6f77730e 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/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 { + 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 { diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index ca9386db0..c2c820157 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/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("source and destination clusters are identical")] + SourceDestinationIdentical, + + #[error("destination cluster has no shards")] + DestinationNoShards, } impl From 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()); } diff --git a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs b/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs index 8be44945c..a84c306b4 100644 --- a/pgdog/src/backend/replication/logical/publisher/publisher_impl.rs +++ b/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; diff --git a/pgdog/src/backend/replication/logical/subscriber/copy.rs b/pgdog/src/backend/replication/logical/subscriber/copy.rs index 3a35ecf94..3821f5156 100644 --- a/pgdog/src/backend/replication/logical/subscriber/copy.rs +++ b/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,7 +13,6 @@ 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, @@ -21,7 +20,10 @@ use crate::{ }, }; -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, connections: Vec, stmt: CopyStatement, @@ -48,12 +50,12 @@ impl CopySubscriber { pub(crate) fn new( copy_stmt: &CopyStatement, source: &Cluster, - cluster: &Cluster, + dest: &Cluster, ) -> Result { let ast = pg_raw_parse::parse(©_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::(); 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(©, &cluster, &cluster).unwrap(); + let mut subscriber = CopySubscriber::new(©, &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 diff --git a/pgdog/src/backend/replication/logical/subscriber/duplicate_check.rs b/pgdog/src/backend/replication/logical/subscriber/duplicate_check.rs new file mode 100644 index 000000000..adb319343 --- /dev/null +++ b/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 { + 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) + } +} diff --git a/pgdog/src/backend/replication/logical/subscriber/mod.rs b/pgdog/src/backend/replication/logical/subscriber/mod.rs index 6291a10bb..25db0c4c2 100644 --- a/pgdog/src/backend/replication/logical/subscriber/mod.rs +++ b/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; diff --git a/pgdog/src/backend/replication/logical/subscriber/parallel_connection.rs b/pgdog/src/backend/replication/logical/subscriber/parallel_connection.rs index 4243d61ce..d9e0f0c0a 100644 --- a/pgdog/src/backend/replication/logical/subscriber/parallel_connection.rs +++ b/pgdog/src/backend/replication/logical/subscriber/parallel_connection.rs @@ -40,6 +40,7 @@ pub(crate) struct ParallelConnection { rx: Receiver, stop: CancellationToken, address: Address, + shard: usize, } impl ParallelConnection { @@ -78,7 +79,7 @@ impl ParallelConnection { } // Move server connection into its own Tokio task. - pub(crate) fn new(server: Server) -> Result { + pub(crate) fn new(server: Server, shard: usize) -> Result { // Ideally we don't hardcode these. PgDog // can use a lot of memory if this is high. let (tx1, rx1) = channel(4096); @@ -104,6 +105,7 @@ impl ParallelConnection { tx: tx1, rx: rx2, stop, + shard, }) } @@ -117,6 +119,11 @@ impl ParallelConnection { _ => Err(Error::ParallelConnection), } } + + /// Get the shard number this connection is connected to. + pub(super) fn shard_number(&self) -> usize { + self.shard + } } // Stop the background task and kill the connection. @@ -211,7 +218,7 @@ mod test { #[tokio::test] async fn test_parallel_connection() { let server = test_server().await; - let mut parallel = ParallelConnection::new(server).unwrap(); + let mut parallel = ParallelConnection::new(server, 0).unwrap(); parallel .send_one(&Parse::named("test", "SELECT $1::bigint").into()) diff --git a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs index 2a2de8612..ed33aa6f6 100644 --- a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs +++ b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs @@ -80,11 +80,12 @@ pub(crate) struct PipelinedConnection { tx: Sender, shared: Arc>, address: Address, + shard: usize, } impl PipelinedConnection { /// This moves `server` into a background task and returns a handle to it. - pub(crate) fn new(server: Server) -> Result { + pub(crate) fn new(server: Server, shard: usize) -> Result { let (tx, rx) = channel(4096); let shared = Arc::new(Mutex::new(Shared::default())); let address = server.addr().clone(); @@ -103,9 +104,14 @@ impl PipelinedConnection { tx, shared, address, + shard, }) } + pub(super) fn shard_number(&self) -> usize { + self.shard + } + /// Server address. pub(crate) fn addr(&self) -> &Address { &self.address @@ -424,7 +430,7 @@ mod test { #[tokio::test] async fn prepare_execute_drain_commit() { let server = test_server().await; - let conn = PipelinedConnection::new(server).unwrap(); + let conn = PipelinedConnection::new(server, 0).unwrap(); // Prepare + create a temp table (out of transaction: uses Sync). conn.prepare( @@ -466,7 +472,7 @@ mod test { #[tokio::test] async fn in_transaction_prepare_uses_flush() { let server = test_server().await; - let conn = PipelinedConnection::new(server).unwrap(); + let conn = PipelinedConnection::new(server, 0).unwrap(); // In-transaction prepare sends Flush and waits for ParseComplete acks // (ParseAcks path) rather than committing with Sync. @@ -488,7 +494,7 @@ mod test { #[tokio::test] async fn prepare_invalid_sql_sync_returns_error() { let server = test_server().await; - let conn = PipelinedConnection::new(server).unwrap(); + let conn = PipelinedConnection::new(server, 0).unwrap(); // Out-of-transaction prepare of invalid SQL: Postgres replies with an // ErrorResponse, which is latched and surfaced through the Sync path. @@ -507,7 +513,7 @@ mod test { #[tokio::test] async fn prepare_invalid_sql_flush_returns_error() { let server = test_server().await; - let conn = PipelinedConnection::new(server).unwrap(); + let conn = PipelinedConnection::new(server, 0).unwrap(); // In-transaction prepare uses Flush, so Postgres sends no ReadyForQuery // on error. The parked ParseAcks waiter can only be released by the @@ -529,7 +535,7 @@ mod test { use tokio::time::timeout; let server = test_server().await; - let conn = PipelinedConnection::new(server).unwrap(); + let conn = PipelinedConnection::new(server, 0).unwrap(); // Valid prepare (succeeds), then a fire-and-forget execute that errors // only at execution time: division by zero. The ErrorResponse arrives @@ -583,7 +589,7 @@ mod test { use tokio::time::sleep; let server = test_server().await; - let conn = PipelinedConnection::new(server).unwrap(); + let conn = PipelinedConnection::new(server, 0).unwrap(); // Fire-and-forget DML that fails at execution time (division by zero). conn.prepare(&[Parse::named("__drain_div", "SELECT 1 / $1::int")], false) @@ -613,7 +619,7 @@ mod test { #[tokio::test] async fn direct_dml_zero_rows_counts_missed() { let server = test_server().await; - let conn = PipelinedConnection::new(server).unwrap(); + let conn = PipelinedConnection::new(server, 0).unwrap(); // Scratch table with one row (id = 1). conn.prepare( diff --git a/pgdog/src/backend/replication/logical/subscriber/stream.rs b/pgdog/src/backend/replication/logical/subscriber/stream.rs index dc054d59d..9c00903b9 100644 --- a/pgdog/src/backend/replication/logical/subscriber/stream.rs +++ b/pgdog/src/backend/replication/logical/subscriber/stream.rs @@ -18,14 +18,12 @@ use super::super::publisher::{NonIdentityColumnsPresence, tables_missing_unique_ use super::super::{ Error, TableValidationError, TableValidationErrorKind, ensure_validation, publisher::Table, }; -use super::PipelinedConnection; -use super::StreamContext; use super::omni_ownership::OmniOwnership; +use super::{OverlappingShardsCheck, PipelinedConnection, StreamContext}; use crate::net::messages::replication::logical::tuple_data::{Identifier, TupleData}; use crate::net::messages::replication::logical::update::Update as XLogUpdate; use crate::{ backend::{Cluster, ConnectReason, Server}, - config::Role, frontend::router::parser::Shard, net::{ Bind, CopyData, ErrorResponse, FromBytes, Parse, Protocol, Sync, ToBytes, @@ -94,8 +92,11 @@ impl Statement { #[derive(Debug)] pub(crate) struct StreamSubscriber { + /// Source cluster. + source: Cluster, + /// Destination cluster. - cluster: Cluster, + dest: Cluster, // Relation markers sent by the publisher. // Happens once per connection. @@ -136,10 +137,16 @@ pub(crate) struct StreamSubscriber { } impl StreamSubscriber { - pub(crate) fn new(cluster: &Cluster, tables: &[Table], partition: OmniOwnership) -> Self { - let cluster = cluster.logical_stream(); + pub(crate) fn new( + source: &Cluster, + dest: &Cluster, + tables: &[Table], + partition: OmniOwnership, + ) -> Self { + let dest = dest.with_replication_settings_override(); Self { - cluster, + source: source.clone(), + dest, relations: HashMap::new(), statements: HashMap::new(), table_lsns: HashMap::new(), @@ -167,31 +174,34 @@ impl StreamSubscriber { } } - // Connect to all the shards. - // - // The transaction-control prepare and the omni FULL-identity validation run - // synchronously on the raw `Server` connections (both are one-shot, - // request/response query flows). Only once they succeed are the connections - // moved into their per-shard pipelined tasks for the streaming apply path. + // Connect to all shards. pub(crate) async fn connect(&mut self) -> Result<(), Error> { - let mut conns: Vec = vec![]; + let mut conns = vec![]; + let overlap_check = OverlappingShardsCheck::new(&self.source); - for shard in self.cluster.shards() { - let primary = shard - .pools_with_roles() - .iter() - .find(|(r, _)| r == &Role::Primary) - .ok_or(Error::NoPrimary)? - .1 - .standalone(ConnectReason::Replication) - .await?; - conns.push(primary); + let dest_shards = self.dest.shards(); + + if dest_shards.is_empty() { + return Err(Error::DestinationNoShards); + } + + for (shard_number, shard) in self.dest.shards().iter().enumerate() { + if overlap_check.overlaps(shard)? { + continue; + } + + let primary = shard.primary_standalone(ConnectReason::Replication).await?; + conns.push((shard_number, primary)); + } + + if conns.is_empty() { + return Err(Error::SourceDestinationIdentical); } // Transaction control statements. // // TODO: Figure out if we need to use them? - for server in &mut conns { + for (_, server) in &mut conns { let begin = Parse::named("__pgdog_repl_begin", "BEGIN"); let commit = Parse::named("__pgdog_repl_commit", "COMMIT"); @@ -217,20 +227,24 @@ impl StreamSubscriber { let omni_full: Vec = self .tables .values() - .filter(|t| { - t.is_identity_full() && !t.is_sharded(&self.cluster.sharding_schema().tables) - }) + .filter(|t| t.is_identity_full() && !t.is_sharded(&self.dest.sharding_schema().tables)) .cloned() .collect(); if !omni_full.is_empty() { - self.validate_full_identity_omni_has_unique_index(&mut conns, &omni_full) - .await?; + self.validate_full_identity_omni_has_unique_index( + &mut conns + .iter_mut() + .map(|(_, server)| server) + .collect::>(), + &omni_full, + ) + .await?; } // Hand each connection to its background pipelining task. self.connections = conns .into_iter() - .map(PipelinedConnection::new) + .map(|(number, server)| PipelinedConnection::new(server, number)) .collect::, _>>()?; Ok(()) @@ -254,7 +268,12 @@ impl StreamSubscriber { // hands the `Bind` over without cloning it, and multi-target // writes clone once per extra target. let mut pending: Option = None; - for shard in 0..n_conns { + for (idx, shard) in self + .connections + .iter() + .enumerate() + .map(|(idx, conn)| (idx, conn.shard_number())) + { let target = match val { // With a single destination shard the router collapses Shard::All // to Direct(0), bypassing the partition ownership check. Apply @@ -267,7 +286,7 @@ impl StreamSubscriber { if !target { continue; } - if let Some(previous) = pending.replace(shard) { + if let Some(previous) = pending.replace(idx) { self.connections[previous] .execute(bind.clone(), is_direct) .await?; @@ -295,7 +314,7 @@ impl StreamSubscriber { } else { statements.insert.parse() }; - let ctx = StreamContext::new(&self.cluster, &insert.tuple_data, parse).await?; + let ctx = StreamContext::new(&self.dest, &insert.tuple_data, parse).await?; { let (shard, bind) = ctx.into_parts(); self.send(&shard, bind).await?; @@ -389,7 +408,7 @@ impl StreamSubscriber { .update .parse() .clone(); - let ctx = StreamContext::new(&self.cluster, new, &parse).await?; + let ctx = StreamContext::new(&self.dest, new, &parse).await?; { let (shard, bind) = ctx.into_parts(); self.send(&shard, bind).await?; @@ -416,7 +435,7 @@ impl StreamSubscriber { let shape_stmt = self .ensure_update_shape_for(oid, &table, &present, false) .await?; - let ctx = StreamContext::new(&self.cluster, &partial_new, shape_stmt.parse()).await?; + let ctx = StreamContext::new(&self.dest, &partial_new, shape_stmt.parse()).await?; { let (shard, bind) = ctx.into_parts(); self.send(&shard, bind).await?; @@ -472,7 +491,7 @@ impl StreamSubscriber { /// Route a tuple to its shard without constructing a `Bind`. /// Used when the bind merges multiple tuples (FULL identity UPDATE/DELETE). async fn shard_for(&self, tuple: &TupleData, parse: &Parse) -> Result { - Ok(StreamContext::new(&self.cluster, tuple, parse) + Ok(StreamContext::new(&self.dest, tuple, parse) .await? .shard() .clone()) @@ -723,7 +742,7 @@ impl StreamSubscriber { debug!("queries for table {} already prepared", dest_key); } else { - let omni = !table.is_sharded(&self.cluster.sharding_schema().tables); + let omni = !table.is_sharded(&self.dest.sharding_schema().tables); let statements = if table.is_identity_full() { // ── FULL identity path ────────────────────────────────────────────── @@ -932,7 +951,7 @@ impl StreamSubscriber { /// the complete set of missing indexes across the cluster in a single error. async fn validate_full_identity_omni_has_unique_index( &self, - servers: &mut [Server], + servers: &mut [&mut Server], tables: &[Table], ) -> Result<(), Error> { // Fan out to all shards concurrently; each gets one IN-list query. @@ -1028,7 +1047,7 @@ mod tests { fn make_subscriber() -> StreamSubscriber { let cluster = Cluster::new_test(&config()); - StreamSubscriber::new(&cluster, &[], OmniOwnership::test()) + StreamSubscriber::new(&Cluster::default(), &cluster, &[], OmniOwnership::test()) } #[test] diff --git a/pgdog/src/backend/replication/logical/subscriber/tests.rs b/pgdog/src/backend/replication/logical/subscriber/tests.rs index 1c271ee77..7ba1e9e99 100644 --- a/pgdog/src/backend/replication/logical/subscriber/tests.rs +++ b/pgdog/src/backend/replication/logical/subscriber/tests.rs @@ -247,12 +247,22 @@ fn x_update(u: XLogUpdate) -> CopyData { fn make_subscriber() -> StreamSubscriber { let cluster = Cluster::new_test(&config()); let tables = vec![make_sharded_table(), make_sharded_test_b_table()]; - StreamSubscriber::new(&cluster, &tables, OmniOwnership::test()) + StreamSubscriber::new( + &Cluster::default(), + &cluster, + &tables, + OmniOwnership::test(), + ) } fn make_subscriber_with_tables(tables: Vec
) -> StreamSubscriber { let cluster = Cluster::new_test(&config()); - StreamSubscriber::new(&cluster, &tables, OmniOwnership::test()) + StreamSubscriber::new( + &Cluster::default(), + &cluster, + &tables, + OmniOwnership::test(), + ) } fn make_subscriber_with_tables_two_databases( @@ -260,13 +270,18 @@ fn make_subscriber_with_tables_two_databases( partition: OmniOwnership, ) -> StreamSubscriber { let cluster = Cluster::new_test_two_databases(&config()); - StreamSubscriber::new(&cluster, &tables, partition) + StreamSubscriber::new(&Cluster::default(), &cluster, &tables, partition) } fn make_subscriber_single_shard() -> StreamSubscriber { let cluster = Cluster::new_test_single_shard(&config()); let tables = vec![make_sharded_table(), make_sharded_test_b_table()]; - StreamSubscriber::new(&cluster, &tables, OmniOwnership::test()) + StreamSubscriber::new( + &Cluster::default(), + &cluster, + &tables, + OmniOwnership::test(), + ) } /// Count rows matching the given `WHERE` predicate using a separate connection. @@ -576,7 +591,12 @@ async fn partition_leaves_share_destination() { leaf_b.table.parent_name = "sharded".to_string(); let cluster = Cluster::new_test_single_shard(&config()); - let mut sub = StreamSubscriber::new(&cluster, &[leaf_a, leaf_b], OmniOwnership::test()); + let mut sub = StreamSubscriber::new( + &Cluster::default(), + &cluster, + &[leaf_a, leaf_b], + OmniOwnership::test(), + ); let mut verify = test_server().await; sub.connect().await.unwrap(); @@ -1508,6 +1528,7 @@ fn omni_insert_copy_data(oid: Oid, a: &str, b: &str) -> CopyData { async fn full_identity_nothing_rejected() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_replica_identity_nothing_table()], OmniOwnership::test(), @@ -1546,6 +1567,7 @@ async fn full_identity_nothing_rejected() { async fn full_identity_omni_no_unique_index_rejected() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_omni_table()], OmniOwnership::test(), @@ -1589,6 +1611,7 @@ async fn full_identity_omni_no_unique_index_rejected() { async fn full_identity_insert_sharded() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_sharded_table()], OmniOwnership::test(), @@ -1619,6 +1642,7 @@ async fn full_identity_insert_sharded() { async fn full_identity_update_fast_path() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_sharded_table()], OmniOwnership::test(), @@ -1678,6 +1702,7 @@ async fn full_identity_update_fast_path() { async fn full_identity_update_slow_path() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_sharded_table()], OmniOwnership::test(), @@ -1742,6 +1767,7 @@ async fn full_identity_update_slow_path() { async fn full_identity_update_slow_path_realistic_old_tuple() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_sharded_table()], OmniOwnership::test(), @@ -1803,6 +1829,7 @@ async fn full_identity_update_slow_path_realistic_old_tuple() { async fn full_identity_update_all_toasted_is_noop() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_sharded_table()], OmniOwnership::test(), @@ -1849,6 +1876,7 @@ async fn full_identity_update_all_toasted_is_noop() { async fn full_identity_delete() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_sharded_table()], OmniOwnership::test(), @@ -1891,6 +1919,7 @@ async fn full_identity_delete() { async fn full_identity_insert_omni_dedup() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_omni_dedup_table()], OmniOwnership::test(), @@ -1954,6 +1983,7 @@ async fn full_identity_insert_omni_dedup() { async fn full_identity_update_duplicate_rows() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_dup_rows_table()], OmniOwnership::test(), @@ -2024,6 +2054,7 @@ async fn full_identity_update_duplicate_rows() { async fn full_identity_delete_duplicate_rows() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_dup_rows_table()], OmniOwnership::test(), @@ -2095,6 +2126,7 @@ async fn full_identity_delete_duplicate_rows() { async fn full_identity_update_matches_null_column() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_dup_rows_table()], OmniOwnership::test(), @@ -2160,6 +2192,7 @@ async fn full_identity_update_matches_null_column() { async fn full_identity_delete_matches_null_column() { let cluster = Cluster::new_test_single_shard(&config()); let mut sub = StreamSubscriber::new( + &Cluster::default(), &cluster, &[make_full_identity_dup_rows_table()], OmniOwnership::test(),