refactor: drop parallel_connections wrapper

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

@@ -81,9 +81,6 @@ pub enum Error {
     #[error("shard {0} has no replication slot")]
     NoReplicationSlot(usize),
 
-    #[error("parallel connection error")]
-    ParallelConnection,
-
     #[error("no replicas available for table sync")]
     NoReplicas,
 

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

@@ -44,7 +44,7 @@ impl ParallelSync {
                 .permit
                 .acquire()
                 .await
-                .map_err(|_| Error::ParallelConnection)?;
+                .map_err(|_| Error::DataSyncAborted)?;
 
             if self.tx.is_closed() {
                 return Err(Error::DataSyncAborted);
@@ -64,9 +64,7 @@ impl ParallelSync {
                 }
             };
 
-            self.tx
-                .send(result)
-                .map_err(|_| Error::ParallelConnection)?;
+            self.tx.send(result).map_err(|_| Error::DataSyncAborted)?;
 
             Ok::<(), Error>(())
         })

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

@@ -1,12 +1,13 @@
 //! Shard COPY stream from one source
 //! between N shards.
 
+use futures::future::try_join_all;
 use pg_query::{parse_raw, NodeEnum};
 use pgdog_config::QueryParserEngine;
 use tracing::debug;
 
 use crate::{
-    backend::{replication::subscriber::ParallelConnection, Cluster, ConnectReason},
+    backend::{Cluster, ConnectReason, Server},
     config::Role,
     frontend::router::parser::{CopyParser, Shard},
     net::{CopyData, CopyDone, ErrorResponse, FromBytes, Protocol, Query, ToBytes},
@@ -24,7 +25,7 @@ pub struct CopySubscriber {
     copy: CopyParser,
     cluster: Cluster,
     buffer: Vec<CopyData>,
-    connections: Vec<ParallelConnection>,
+    connections: Vec<Server>,
     stmt: CopyStatement,
     bytes_sharded: usize,
 }
@@ -85,7 +86,7 @@ impl CopySubscriber {
                 .1
                 .standalone(ConnectReason::Replication)
                 .await?;
-            servers.push(ParallelConnection::new(primary)?);
+            servers.push(primary);
         }
 
         self.connections = servers;
@@ -95,38 +96,38 @@ impl CopySubscriber {
 
     /// Disconnect from all shards.
     pub async fn disconnect(&mut self) -> Result<(), Error> {
-        for conn in std::mem::take(&mut self.connections) {
-            conn.reattach().await?;
-        }
+        self.connections.clear();
 
         Ok(())
     }
 
     /// Start COPY on all shards.
     pub async fn start_copy(&mut self) -> Result<(), Error> {
-        let stmt = Query::new(self.stmt.copy_in());
-
         if self.connections.is_empty() {
             self.connect().await?;
         }
 
-        for server in &mut self.connections {
-            debug!("{} [{}]", stmt.query(), server.addr());
+        let stmt = Query::new(self.stmt.copy_in());
 
-            server.send_one(&stmt.clone().into()).await?;
-            server.flush().await?;
+        // Start COPY IN on all shards concurrently.
+        try_join_all(self.connections.iter_mut().map(|server| {
+            let msg: crate::net::ProtocolMessage = stmt.clone().into();
+            debug!("{} [{}]", stmt.query(), server.addr());
 
-            let msg = server.read().await?;
-            match msg.code() {
-                'G' => (),
-                'E' => {
-                    return Err(Error::PgError(Box::new(ErrorResponse::from_bytes(
-                        msg.to_bytes()?,
-                    )?)))
+            async move {
+                server.send_one(&msg).await?;
+                server.flush().await?;
+                let reply = server.read().await?;
+                match reply.code() {
+                    'G' => Ok(()),
+                    'E' => Err(Error::PgError(Box::new(ErrorResponse::from_bytes(
+                        reply.to_bytes()?,
+                    )?))),
+                    c => Err(Error::OutOfSync(c)),
                 }
-                c => return Err(Error::OutOfSync(c)),
             }
-        }
+        }))
+        .await?;
 
         Ok(())
     }
@@ -135,20 +136,20 @@ impl CopySubscriber {
     pub async fn copy_done(&mut self) -> Result<(), Error> {
         self.flush().await?;
 
-        for server in &mut self.connections {
+        // Finalise COPY on all shards concurrently.
+        try_join_all(self.connections.iter_mut().map(|server| async move {
             server.send_one(&CopyDone.into()).await?;
             server.flush().await?;
 
-            let command_complete = server.read().await?;
-            match command_complete.code() {
+            let cc = server.read().await?;
+            match cc.code() {
                 'E' => {
-                    let error = ErrorResponse::from_bytes(command_complete.to_bytes()?)?;
+                    let error = ErrorResponse::from_bytes(cc.to_bytes()?)?;
                     if error.code == "08P01" && error.message == "insufficient data left in message"
                     {
                         return Err(Error::BinaryFormatMismatch(Box::new(error)));
-                    } else {
-                        return Err(Error::PgError(Box::new(error)));
                     }
+                    return Err(Error::PgError(Box::new(error)));
                 }
                 'C' => (),
                 c => return Err(Error::OutOfSync(c)),
@@ -158,7 +159,9 @@ impl CopySubscriber {
             if rfq.code() != 'Z' {
                 return Err(Error::OutOfSync(rfq.code()));
             }
-        }
+            Ok(())
+        }))
+        .await?;
 
         Ok(())
     }
@@ -174,12 +177,19 @@ impl CopySubscriber {
     }
 
     async fn flush(&mut self) -> Result<(usize, usize), Error> {
+        if self.buffer.is_empty() {
+            return Ok((0, 0));
+        }
+
         let result = self.copy.shard(&self.buffer)?;
         self.buffer.clear();
 
         let rows = result.len();
         let bytes = result.iter().map(|row| row.len()).sum::<usize>();
+        self.bytes_sharded += bytes;
 
+        // Route each row to the right shard(s). send_one is a buffered write
+        // so this loop does no I/O — no concurrency needed here.
         for row in &result {
             for (shard, server) in self.connections.iter_mut().enumerate() {
                 match row.shard() {
@@ -198,7 +208,8 @@ impl CopySubscriber {
             }
         }
 
-        self.bytes_sharded += result.iter().map(|c| c.len()).sum::<usize>();
+        // Flush all shards concurrently — this is the actual socket write.
+        try_join_all(self.connections.iter_mut().map(|s| s.flush())).await?;
 
         Ok((rows, bytes))
     }

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

@@ -1,12 +1,10 @@
 pub mod context;
 pub mod copy;
-pub mod parallel_connection;
 pub mod stream;
 
 #[cfg(test)]
 mod tests;
 
 pub use context::StreamContext;
 pub use copy::CopySubscriber;
-pub use parallel_connection::ParallelConnection;
 pub use stream::StreamSubscriber;