From 1276f761219ac6cd5c545a093cc565ba03f8f6c9 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Thu, 24 Sep 2026 15:50:13 -0700 Subject: [PATCH 1/2] ok --- pgdog/src/backend/pool/connection/binding.rs | 26 ++--- pgdog/src/backend/pool/connection/mod.rs | 20 +++- .../pool/connection/multi_shard/binding.rs | 96 +++++++++++++++---- .../pool/connection/multi_shard/mod.rs | 23 +---- .../pool/connection/multi_shard/test.rs | 22 ++--- 5 files changed, 116 insertions(+), 71 deletions(-) diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index f80a9ed77..a140e32e2 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -16,7 +16,7 @@ use futures::future::join_all; use super::*; use crate::util::safe_sleep; -use multi_shard::MultiBinding; +use multi_shard::{LinkedServer, MultiBinding}; /// The server(s) the client is connected to. #[derive(Debug, Default)] @@ -241,12 +241,12 @@ impl Binding { } pub(crate) async fn two_pc_on_guards( - servers: &mut [Guard], + servers: &mut [LinkedServer], transaction: TwoPcTransaction, phase: TwoPcPhase, ignore_missing: bool, ) -> Result<(), Error> { - let mut futures = Vec::new(); + let mut futures = vec![]; for (shard, server) in servers.iter_mut().enumerate() { let query = phase_control(transaction, shard, phase); futures.push(server.execute(query)); @@ -282,7 +282,8 @@ impl Binding { ) -> Result<(), Error> { match self { Binding::MultiShard(servers) => { - Self::two_pc_on_guards(servers, transaction, phase, ignore_missing).await + Self::two_pc_on_guards(servers.deref_mut(), transaction, phase, ignore_missing) + .await } _ => Err(Error::TwoPcMultiShardOnly), @@ -300,21 +301,10 @@ impl Binding { Binding::Direct(server, ..) => { server.link_client(id, params, transaction_start_stmt).await } - Binding::MultiShard(servers) => { - let futures = servers - .iter_mut() - .map(|server| server.link_client(id, params, transaction_start_stmt)); - let results = join_all(futures).await; - let mut max = 0; - for result in results { - let synced = result?; - if max < synced { - max = synced; - } - } - Ok(max) - } + Binding::MultiShard(servers) => Ok(servers + .link_client(id, params, transaction_start_stmt) + .await?), _ => Ok(0), } diff --git a/pgdog/src/backend/pool/connection/mod.rs b/pgdog/src/backend/pool/connection/mod.rs index f5fa9cf88..ac86ecd90 100644 --- a/pgdog/src/backend/pool/connection/mod.rs +++ b/pgdog/src/backend/pool/connection/mod.rs @@ -88,6 +88,22 @@ impl Connection { Ok(()) } + #[allow(unused)] + pub(crate) async fn ensure_connected( + &mut self, + request: &Request, + route: &Route, + ) -> Result<(), Error> { + match self.binding { + Binding::Direct(_, _) => Ok(()), + Binding::NotConnected => Err(Error::NotConnected), + Binding::MultiShard(ref mut servers) => { + Ok(servers.ensure_connected(request, route).await?) + } + Binding::Admin(_) => Ok(()), + } + } + /// Send client request to mirrors. pub(crate) fn mirror(&mut self, buffer: &crate::frontend::ClientRequest) { for mirror in self.cluster.mirrors() { @@ -356,7 +372,9 @@ impl Connection { pub(crate) async fn cancel_query(&self) -> Result<(), Error> { let servers: Vec<&Guard> = match self.binding { Binding::Direct(ref server, ..) => vec![server], - Binding::MultiShard(ref servers) => servers.iter().collect(), + Binding::MultiShard(ref servers) => { + servers.iter().map(|server| server.deref()).collect() + } _ => return Ok(()), }; diff --git a/pgdog/src/backend/pool/connection/multi_shard/binding.rs b/pgdog/src/backend/pool/connection/multi_shard/binding.rs index 8818283ba..5e4c05406 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/binding.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/binding.rs @@ -3,39 +3,95 @@ use std::ops::{Deref, DerefMut}; use futures::future::join_all; use crate::backend::Error; +use crate::backend::pool::Request; use crate::frontend::ClientRequest; use crate::frontend::router::parser::Shard; use crate::frontend::router::{CopyRow, Route}; -use crate::net::{Message, ProtocolMessage}; +use crate::net::{FrontendPid, Message, Parameters, ProtocolMessage}; use super::super::Guard; use super::MultiShard; -/// Handle talking to multiple servers for cross-shard queries. #[derive(Debug)] -pub(crate) struct MultiBinding { - servers: Vec, - state: Box, +pub(crate) struct LinkedServer { + server: Guard, + // Shard number. + shard: usize, + // Parameters were sync'ed. + linked: bool, } -impl From> for MultiBinding { - fn from(value: Vec) -> Self { - Self { - servers: value, - state: Box::default(), - } +impl Deref for LinkedServer { + type Target = Guard; + + fn deref(&self) -> &Self::Target { + &self.server + } +} + +impl DerefMut for LinkedServer { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.server } } +/// Handle talking to multiple servers for cross-shard queries. +#[derive(Debug)] +pub(crate) struct MultiBinding { + servers: Vec, + state: Box, +} + impl MultiBinding { /// Create new multi-shard binding. pub(crate) fn new(servers: Vec, shard_indices: Vec, route: &Route) -> Self { Self { - state: Box::new(MultiShard::new(shard_indices, route)), - servers, + state: Box::new(MultiShard::new(servers.len(), route)), + servers: servers + .into_iter() + .zip(shard_indices.into_iter()) + .map(|(server, shard)| LinkedServer { + server, + shard, + linked: true, + }) + .collect(), } } + #[allow(unused)] + pub(crate) async fn ensure_connected( + &mut self, + request: &Request, + route: &Route, + ) -> Result<(), super::Error> { + Ok(()) + } + + pub(crate) async fn link_client( + &mut self, + client_id: FrontendPid, + params: &Parameters, + transaction_start_stmt: Option<&str>, + ) -> Result { + let futures = self + .servers + .iter_mut() + .filter(|server| !server.linked) + .map(|server| server.link_client(client_id, params, transaction_start_stmt)); + let results = join_all(futures).await; + + let mut max = 0; + for result in results { + let synced = result?; + if max < synced { + max = synced; + } + } + + Ok(max) + } + /// Read-only handle to internal state. pub(crate) fn state(&self) -> &MultiShard { &self.state @@ -83,11 +139,11 @@ impl MultiBinding { let mut shards_sent = self.servers.len(); let mut futures = Vec::new(); - for (position, server) in self.servers.iter_mut().enumerate() { + for server in self.servers.iter_mut() { // Map positional index to actual shard number. // When only a subset of shards is connected (Shard::Multi binding), // positional indices don't match actual shard numbers. - let shard = self.state.shard_number(position); + let shard = server.shard; let send = match client_request.route().shard() { Shard::Direct(s) => { shards_sent = 1; @@ -134,8 +190,8 @@ impl MultiBinding { } let mut futures = Vec::new(); - for (position, server) in self.servers.iter_mut().enumerate() { - let shard = self.state.shard_number(position); + for server in self.servers.iter_mut() { + let shard = server.shard; let send = match route.shard() { Shard::Direct(s) => *s == shard, Shard::Multi(shards) => shards.contains(&shard), @@ -157,8 +213,8 @@ impl MultiBinding { /// Send COPY rows to all shards. pub(crate) async fn send_copy(&mut self, rows: Vec) -> Result<(), Error> { for row in rows { - for (position, server) in self.servers.iter_mut().enumerate() { - let shard = self.state.shard_number(position); + for server in self.servers.iter_mut() { + let shard = server.shard; match row.shard() { Shard::Direct(row_shard) => { if shard == *row_shard { @@ -190,7 +246,7 @@ impl MultiBinding { } impl Deref for MultiBinding { - type Target = Vec; + type Target = Vec; fn deref(&self) -> &Self::Target { &self.servers diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index b828905db..ebf47a364 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -21,7 +21,7 @@ mod error; mod test; mod validator; -pub(crate) use binding::MultiBinding; +pub(crate) use binding::{LinkedServer, MultiBinding}; pub(crate) use error::Error; use validator::Validator; @@ -53,10 +53,6 @@ pub(crate) struct MultiShard { shards: usize, /// Route the query is taking. route: Route, - /// Maps positional index in the servers vec to actual shard number. - /// When all shards are connected, this is `[0, 1, 2, ...]`. - /// When only a subset is connected (e.g. shards 0 and 2), this is `[0, 2]`. - shard_indices: Vec, /// In-flight request state. request_state: RequestState, /// Sorting/aggregate buffer. @@ -71,29 +67,14 @@ pub(crate) struct MultiShard { impl MultiShard { /// New multi-shard state given the actual shard indices connected. - pub(super) fn new(shard_indices: Vec, route: &Route) -> Self { - let shards = shard_indices.len(); + pub(super) fn new(shards: usize, route: &Route) -> Self { Self { shards, - shard_indices, route: route.clone(), ..Default::default() } } - /// Map a positional index to the actual shard number. - /// - /// These can diverge since we can connect to less shards than there - /// are in the config, while the query parser will produce shard numbers - /// relative to all configured shards. - /// - pub(super) fn shard_number(&self, position: usize) -> usize { - self.shard_indices - .get(position) - .copied() - .unwrap_or(position) - } - /// Update multi-shard state. pub(super) fn update(&mut self, shards: usize, route: &Route) { self.reset(); diff --git a/pgdog/src/backend/pool/connection/multi_shard/test.rs b/pgdog/src/backend/pool/connection/multi_shard/test.rs index 8cd70f231..9f56d93b6 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/test.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/test.rs @@ -8,7 +8,7 @@ use super::*; #[test] fn test_inconsistent_row_descriptions() { let route = Route::default(); - let mut multi_shard = MultiShard::new(vec![0, 1], &route); + let mut multi_shard = MultiShard::new(2, &route); // Create two different row descriptions let rd1 = RowDescription::new(&[Field::text("name"), Field::bigint("id")]); @@ -32,7 +32,7 @@ fn test_inconsistent_row_descriptions() { #[test] fn test_inconsistent_data_rows() { let route = Route::default(); - let mut multi_shard = MultiShard::new(vec![0, 1], &route); + let mut multi_shard = MultiShard::new(2, &route); // Set up row description first let rd = RowDescription::new(&[Field::text("name"), Field::bigint("id")]); @@ -63,7 +63,7 @@ fn test_inconsistent_data_rows() { #[test] fn test_rd_before_dr() { let mut multi_shard = MultiShard::new( - vec![0, 1, 2], + 3, &Route::read(ShardWithPriority::new_default_unset(Shard::All)), ); let rd = RowDescription::new(&[Field::bigint("id")]); @@ -131,7 +131,7 @@ fn test_distinct_state_resets_between_requests() { Default::default(), Some(DistinctBy::Row), ); - let mut multi_shard = MultiShard::new(vec![0, 1], &route); + let mut multi_shard = MultiShard::new(2, &route); let row_description = RowDescription::new(&[Field::bigint("id")]); let mut data_row = DataRow::new(); data_row.add(1_i64); @@ -173,7 +173,7 @@ fn test_distinct_state_resets_between_requests() { #[test] fn test_ready_for_query_error_preservation() { let route = Route::default(); - let mut multi_shard = MultiShard::new(vec![0, 1], &route); + let mut multi_shard = MultiShard::new(2, &route); // Create ReadyForQuery messages - one with transaction error, one normal let rfq_error = ReadyForQuery::error(); @@ -201,7 +201,7 @@ fn test_ready_for_query_error_preservation() { fn test_omni_command_complete_not_summed() { // For omni-sharded tables, we should NOT sum row counts across shards. let route = Route::write(ShardWithPriority::new_table_omni(Shard::All)).with_omnisharded(true); - let mut multi_shard = MultiShard::new(vec![0, 1, 2], &route); + let mut multi_shard = MultiShard::new(3, &route); let backend1 = BackendPid::for_test(1); let backend2 = BackendPid::for_test(2); @@ -240,7 +240,7 @@ fn test_omni_command_complete_not_summed() { fn test_omni_command_complete_uses_first_shard_row_count() { // For omni, we use the first shard's row count for consistency with DataRow behavior. let route = Route::write(ShardWithPriority::new_table_omni(Shard::All)).with_omnisharded(true); - let mut multi_shard = MultiShard::new(vec![0, 1], &route); + let mut multi_shard = MultiShard::new(2, &route); let backend1 = BackendPid::for_test(1); let backend2 = BackendPid::for_test(2); @@ -273,7 +273,7 @@ fn test_omni_command_complete_uses_first_shard_row_count() { fn test_omni_data_rows_only_from_first_server() { // For omni-sharded tables with RETURNING, only forward DataRows from the first server. let route = Route::write(ShardWithPriority::new_table_omni(Shard::All)).with_omnisharded(true); - let mut multi_shard = MultiShard::new(vec![0, 1], &route); + let mut multi_shard = MultiShard::new(2, &route); let backend1 = BackendPid::for_test(1); let backend2 = BackendPid::for_test(2); @@ -319,7 +319,7 @@ fn test_omni_data_rows_only_from_first_server() { fn test_pipelined_describe_forwards_every_group() { for shards in [1, 2] { let mut multi_shard = MultiShard::new( - (0..shards).collect(), + shards, &Route::read(ShardWithPriority::new_default_unset(Shard::All)), ); @@ -351,7 +351,7 @@ fn test_pipelined_describe_forwards_every_group() { #[test] fn test_bind_result_formats_apply_per_statement() { let mut multi_shard = MultiShard::new( - vec![0, 1], + 2, &Route::read(ShardWithPriority::new_default_unset(Shard::All)), ); @@ -389,7 +389,7 @@ fn test_bind_result_formats_apply_per_statement() { #[test] fn test_ready_for_query_drops_pending_binds() { let mut multi_shard = MultiShard::new( - vec![0, 1], + 2, &Route::read(ShardWithPriority::new_default_unset(Shard::All)), ); From eae5439ed8f63d3eb11be6e455e02465b9018c17 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Thu, 24 Sep 2026 18:17:58 -0700 Subject: [PATCH 2/2] fix link --- pgdog/src/backend/pool/connection/multi_shard/binding.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pgdog/src/backend/pool/connection/multi_shard/binding.rs b/pgdog/src/backend/pool/connection/multi_shard/binding.rs index 5e4c05406..6e4226318 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/binding.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/binding.rs @@ -53,7 +53,7 @@ impl MultiBinding { .map(|(server, shard)| LinkedServer { server, shard, - linked: true, + linked: false, }) .collect(), } @@ -73,7 +73,7 @@ impl MultiBinding { client_id: FrontendPid, params: &Parameters, transaction_start_stmt: Option<&str>, - ) -> Result { + ) -> Result { let futures = self .servers .iter_mut()