Levkk lazy conns 4

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

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

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<Guard>,
-    state: Box<MultiShard>,
+pub(crate) struct LinkedServer {
+    server: Guard,
+    // Shard number.
+    shard: usize,
+    // Parameters were sync'ed.
+    linked: bool,
 }
 
-impl From<Vec<Guard>> for MultiBinding {
-    fn from(value: Vec<Guard>) -> 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<LinkedServer>,
+    state: Box<MultiShard>,
+}
+
 impl MultiBinding {
     /// Create new multi-shard binding.
     pub(crate) fn new(servers: Vec<Guard>, shard_indices: Vec<usize>, 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: false,
+                })
+                .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<usize, Error> {
+        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<CopyRow>) -> 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<Guard>;
+    type Target = Vec<LinkedServer>;
 
     fn deref(&self) -> &Self::Target {
         &self.servers

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<usize>,
     /// 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<usize>, 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();

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