From e1972b6e480d0bcb13bcd90c2c2b513965885d18 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Thu, 9 Jul 2026 00:39:04 +0530 Subject: [PATCH 1/5] faster replication v1 --- .../src/backend/replication/logical/error.rs | 7 + .../replication/logical/subscriber/mod.rs | 2 + .../logical/subscriber/pipeline.rs | 415 ++++++++++++++++++ .../replication/logical/subscriber/stream.rs | 202 ++++----- 4 files changed, 514 insertions(+), 112 deletions(-) create mode 100644 pgdog/src/backend/replication/logical/subscriber/pipeline.rs diff --git a/pgdog/src/backend/replication/logical/error.rs b/pgdog/src/backend/replication/logical/error.rs index 3f52c5d19..211cd8032 100644 --- a/pgdog/src/backend/replication/logical/error.rs +++ b/pgdog/src/backend/replication/logical/error.rs @@ -144,6 +144,9 @@ pub enum Error { #[error("parallel connection error")] ParallelConnection, + #[error("pipelined connection task closed")] + PipelineClosed, + #[error("no replicas available for table sync")] NoReplicas, @@ -274,6 +277,10 @@ impl Error { // either the ParallelConnection wrapper should be removed or // the proper error should be propagated Self::ParallelConnection => true, + // A pipelined shard task closed (write failed, socket died, or the + // handle was dropped). The transaction is torn down; retry from a + // fresh connection. + Self::PipelineClosed => true, _ => false, } } diff --git a/pgdog/src/backend/replication/logical/subscriber/mod.rs b/pgdog/src/backend/replication/logical/subscriber/mod.rs index 27d1853b9..a2f871694 100644 --- a/pgdog/src/backend/replication/logical/subscriber/mod.rs +++ b/pgdog/src/backend/replication/logical/subscriber/mod.rs @@ -2,6 +2,7 @@ pub mod context; pub mod copy; pub mod omni_ownership; pub mod parallel_connection; +pub mod pipeline; pub mod stream; #[cfg(test)] @@ -10,4 +11,5 @@ mod tests; pub use context::StreamContext; pub use copy::CopySubscriber; pub use parallel_connection::ParallelConnection; +pub use pipeline::PipelinedConnection; pub use stream::StreamSubscriber; diff --git a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs new file mode 100644 index 000000000..67fb54389 --- /dev/null +++ b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs @@ -0,0 +1,415 @@ +use std::collections::VecDeque; +use std::sync::Arc; + +use parking_lot::Mutex; +use tokio::select; +use tokio::spawn; +use tokio::sync::{ + mpsc::{Receiver, Sender, channel}, + oneshot, +}; +use tracing::trace; + +use super::stream::MissedRows; +use crate::backend::Server; +use crate::backend::pool::Address; +use crate::net::{ + Bind, CommandComplete, ErrorResponse, Execute, Flush, FromBytes, Message, Parse, Protocol, + ProtocolMessage, Sync, ToBytes, +}; + +use super::super::Error; + +// State shared between the handle and its background listener task. +#[derive(Debug, Default)] +struct Shared { + // First error observed on this connection. Sticky until taken. + error: Option, + // Rows a direct-to-shard DML expected to touch but didn't (0 rows affected). + missed: MissedRows, +} + +// How a sync point completes. +enum SyncPointKind { + // Wait for `n` ParseComplete responses (in-transaction prepare, `Flush`). + ParseAcks(usize), + // Wait for a single ReadyForQuery (`Sync`: commit, or out-of-transaction prepare). + ReadyForQuery, +} + +// Work sent from the handle to the listener task. +enum Command { + // Fire-and-forget DML: Bind/Execute/Flush. `is_direct` drives missed-row counting. + Execute { + messages: Vec, + is_direct: bool, + }, + // A synchronization point: write `messages`, then resolve `done` per `kind`. + SyncPoint { + messages: Vec, + kind: SyncPointKind, + done: oneshot::Sender<()>, + }, + // Wait until every outstanding DML ack has been read (no messages sent). + DrainAcks { + done: oneshot::Sender<()>, + }, +} + +/// Pipelined destination connection: a handle over a background task that +/// owns the `Server` and reconciles responses. +#[derive(Debug)] +pub struct PipelinedConnection { + tx: Sender, + shared: Arc>, + address: Address, +} + +impl PipelinedConnection { + /// This moves `server` into a background task and returns a handle to it. + pub fn new(server: Server) -> Result { + let (tx, rx) = channel(4096); + let shared = Arc::new(Mutex::new(Shared::default())); + let address = server.addr().clone(); + + let listener = Listener { + rx, + server, + shared: shared.clone(), + direct: VecDeque::new(), + prepare_sync_points: VecDeque::new(), + ready_sync_points: VecDeque::new(), + drain_waiter: None, + }; + + spawn(listener.run()); + + Ok(Self { + tx, + shared, + address, + }) + } + + /// Server address. + pub fn addr(&self) -> &Address { + &self.address + } + + /// Enqueue a DML statement (`Bind/Execute/Flush`) without waiting for its + /// response. `is_direct` marks a single-shard write whose 0-row result + /// counts as a missed row. + pub async fn execute(&self, bind: &Bind, is_direct: bool) -> Result<(), Error> { + let messages = vec![bind.clone().into(), Execute::new().into(), Flush.into()]; + self.tx + .send(Command::Execute { + messages, + is_direct, + }) + .await + .map_err(|_| Error::PipelineClosed) + } + + /// Prepare `parses` and wait for the acknowledgments. Inside a transaction + /// uses `Flush` (must not commit the open implicit transaction); otherwise + /// uses `Sync`. + pub async fn prepare(&self, parses: &[Parse], in_transaction: bool) -> Result<(), Error> { + if parses.is_empty() { + return Ok(()); + } + let mut messages: Vec = parses.iter().map(|p| p.clone().into()).collect(); + let kind = if in_transaction { + messages.push(Flush.into()); + SyncPointKind::ParseAcks(parses.len()) + } else { + messages.push(Sync.into()); + SyncPointKind::ReadyForQuery + }; + self.sync_point(messages, kind).await + } + + /// Send `Sync` and wait for `ReadyForQuery` (commits the open implicit + /// transaction on this shard). + pub async fn sync_and_drain(&self) -> Result<(), Error> { + self.sync_point(vec![Sync.into()], SyncPointKind::ReadyForQuery) + .await + } + + /// Wait until every outstanding DML acknowledgment has been read. Used at + /// commit to confirm the shard is error-free before any shard is `Sync`ed. + pub async fn drain_acks(&self) -> Result<(), Error> { + let (done, rx) = oneshot::channel(); + self.tx + .send(Command::DrainAcks { done }) + .await + .map_err(|_| Error::PipelineClosed)?; + self.await_done(rx).await + } + + /// Non-blocking peek + take of the latched error. `Some` means the shard + /// has errored and the transaction must be rolled back. + pub fn take_error(&self) -> Option { + self.shared.lock().error.take() + } + + /// Drain the missed-row counters accumulated by the listener. + pub(crate) fn take_missed_rows(&self) -> MissedRows { + std::mem::take(&mut self.shared.lock().missed) + } + + async fn sync_point( + &self, + messages: Vec, + kind: SyncPointKind, + ) -> Result<(), Error> { + let (done, rx) = oneshot::channel(); + self.tx + .send(Command::SyncPoint { + messages, + kind, + done, + }) + .await + .map_err(|_| Error::PipelineClosed)?; + self.await_done(rx).await + } + + // Wait for a sync point to resolve, then report any latched error. If the + // task died before signaling, surface the latched error or a generic one. + async fn await_done(&self, rx: oneshot::Receiver<()>) -> Result<(), Error> { + match rx.await { + Ok(()) => match self.take_error() { + Some(err) => Err(err), + None => Ok(()), + }, + Err(_) => Err(self.take_error().unwrap_or(Error::PipelineClosed)), + } + } +} + +// Background task: owns the `Server`, writes queued messages, and reconciles +// responses (counts acks, records missed rows, latches the first error). +struct Listener { + rx: Receiver, + server: Server, + shared: Arc>, + // One entry per outstanding DML op, in send order: `is_direct`. Popped on + // each CommandComplete. + direct: VecDeque, + // In-transaction prepare sync points: (remaining ParseComplete acks, waiter). + prepare_sync_points: VecDeque<(usize, oneshot::Sender<()>)>, + // Commit / out-of-transaction prepare sync points, resolved on ReadyForQuery. + ready_sync_points: VecDeque>, + // Waiter that resolves once every DML ack has been read. At most one is + // outstanding: `commit()` issues drain_acks sequentially, one per connection. + drain_waiter: Option>, +} + +impl Listener { + async fn run(mut self) { + loop { + select! { + command = self.rx.recv() => { + match command { + Some(command) => { + if self.handle_command(command).await.is_err() { + break; + } + } + None => break, + } + } + + message = self.server.read() => { + match message { + Ok(message) => self.handle_response(message), + Err(err) => { + self.latch_error(err.into()); + self.wake_all(); + break; + } + } + } + } + } + } + + async fn handle_command(&mut self, command: Command) -> Result<(), Error> { + let errored = self.shared.lock().error.is_some(); + + match command { + Command::Execute { + messages, + is_direct, + } => { + if errored { + return Ok(()); + } + if let Err(err) = self.write(&messages).await { + self.latch_error(err); + self.wake_all(); + return Err(Error::PipelineClosed); + } + self.direct.push_back(is_direct); + } + Command::SyncPoint { + messages, + kind, + done, + } => { + if errored { + let _ = done.send(()); + return Ok(()); + } + if let Err(err) = self.write(&messages).await { + self.latch_error(err); + let _ = done.send(()); + self.wake_all(); + return Err(Error::PipelineClosed); + } + match kind { + // `prepare()` never enqueues an empty batch, so `remaining >= 1`. + SyncPointKind::ParseAcks(remaining) => { + self.prepare_sync_points.push_back((remaining, done)); + } + SyncPointKind::ReadyForQuery => self.ready_sync_points.push_back(done), + } + } + Command::DrainAcks { done } => { + if errored || self.direct.is_empty() { + let _ = done.send(()); + } else { + // At most one drain is outstanding (issued one-at-a-time by + // commit()). Overwriting would drop a prior sender, which + // surfaces as an error on its receiver rather than a hang. + self.drain_waiter = Some(done); + } + } + } + + Ok(()) + } + + fn handle_response(&mut self, message: Message) { + let code = message.code(); + trace!("[{}] --> {}", self.address(), code); + + match code { + 'E' => { + let err = ErrorResponse::from_bytes(message.to_bytes()) + .map(|resp| Error::PgError(Box::new(resp))) + .unwrap_or(Error::PipelineClosed); + self.latch_error(err); + // After an error Postgres skips until Sync; abandon tracking and + // resolve every waiter so the handle can roll back. + self.direct.clear(); + self.wake_all(); + } + // ParseComplete: advance the front in-transaction prepare sync point. + '1' => { + let resolved = if let Some((remaining, _)) = self.prepare_sync_points.front_mut() { + *remaining -= 1; + *remaining == 0 + } else { + false + }; + if resolved { + if let Some((_, done)) = self.prepare_sync_points.pop_front() { + let _ = done.send(()); + } + } + } + // BindComplete: nothing to account for. + '2' => {} + // CommandComplete: match to the DML op and count missed rows. + 'C' => { + let is_direct = self.direct.pop_front().unwrap_or(false); + if is_direct { + if let Ok(complete) = CommandComplete::try_from(message) { + if matches!(complete.rows(), Ok(Some(0))) { + self.shared.lock().missed.record(complete.tag()); + } + } + } + if self.direct.is_empty() { + self.wake_drain(); + } + } + // ReadyForQuery: resolve the front commit sync point. + 'Z' => { + if let Some(done) = self.ready_sync_points.pop_front() { + let _ = done.send(()); + } + } + // NoticeResponse / ParameterStatus / NotificationResponse / etc. + _ => {} + } + } + + // Write messages to the socket and flush so the bytes reach Postgres. + async fn write(&mut self, messages: &[ProtocolMessage]) -> Result<(), Error> { + for message in messages { + self.server.send_one(message).await?; + } + self.server.flush().await?; + Ok(()) + } + + fn latch_error(&self, err: Error) { + let mut shared = self.shared.lock(); + if shared.error.is_none() { + shared.error = Some(err); + } + } + + fn wake_all(&mut self) { + while let Some((_, done)) = self.prepare_sync_points.pop_front() { + let _ = done.send(()); + } + while let Some(done) = self.ready_sync_points.pop_front() { + let _ = done.send(()); + } + self.wake_drain(); + } + + fn wake_drain(&mut self) { + if let Some(done) = self.drain_waiter.take() { + let _ = done.send(()); + } + } + + fn address(&self) -> &Address { + self.server.addr() + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::{backend::server::test::test_server, net::Parse}; + + #[tokio::test] + async fn prepare_and_commit_drain() { + let server = test_server().await; + let conn = PipelinedConnection::new(server).unwrap(); + + // Prepare out of a transaction (uses Sync internally). + conn.prepare(&[Parse::named("__pipe_test", "SELECT $1::bigint")], false) + .await + .unwrap(); + + // A bare Sync must drain cleanly with no error. + conn.sync_and_drain().await.unwrap(); + assert!(conn.take_error().is_none()); + } + + #[tokio::test] + async fn drain_acks_with_no_outstanding_is_immediate() { + let server = test_server().await; + let conn = PipelinedConnection::new(server).unwrap(); + + // No DML enqueued: drain_acks resolves immediately without error. + conn.drain_acks().await.unwrap(); + assert!(conn.take_error().is_none()); + } +} diff --git a/pgdog/src/backend/replication/logical/subscriber/stream.rs b/pgdog/src/backend/replication/logical/subscriber/stream.rs index a006db7e7..b698eeeb7 100644 --- a/pgdog/src/backend/replication/logical/subscriber/stream.rs +++ b/pgdog/src/backend/replication/logical/subscriber/stream.rs @@ -18,6 +18,7 @@ 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 crate::net::messages::replication::logical::tuple_data::{Identifier, TupleData}; @@ -27,8 +28,7 @@ use crate::{ config::Role, frontend::router::parser::Shard, net::{ - Bind, CommandComplete, CopyData, ErrorResponse, Execute, Flush, FromBytes, Parse, Protocol, - Sync, ToBytes, + Bind, CopyData, ErrorResponse, FromBytes, Parse, Protocol, Sync, ToBytes, replication::{ Commit as XLogCommit, Delete as XLogDelete, Insert as XLogInsert, Relation, StatusUpdate, UpdateIdentity, xlog_data::XLogPayload, @@ -116,8 +116,9 @@ pub struct StreamSubscriber { // watermark on commit so equal-LSN rows in the same transaction are not skipped. changed_tables: HashSet, - // Connections to shards. - connections: Vec, + // Pipelined connections to shards. DML is pushed without waiting per event; + // a background task per connection reads responses and reconciles them. + connections: Vec, // Last commit LSN acked to Postgres. Reported in status updates; never // advances mid-transaction so KeepAlive replies can't skip an open transaction. @@ -130,9 +131,6 @@ pub struct StreamSubscriber { // Bytes sharded bytes_sharded: usize, - // Missed rows. - missed_rows: MissedRows, - // Determines which destination shards this subscriber owns for omni tables. partition: OmniOwnership, } @@ -164,15 +162,19 @@ impl StreamSubscriber { bytes_sharded: 0, lsn_changed: true, in_transaction: false, - missed_rows: MissedRows::default(), keys: HashMap::default(), partition, } } // 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. pub async fn connect(&mut self) -> Result<(), Error> { - let mut conns = vec![]; + let mut conns: Vec = vec![]; for shard in self.cluster.shards() { let primary = shard @@ -186,12 +188,10 @@ impl StreamSubscriber { conns.push(primary); } - self.connections = conns; - // Transaction control statements. // // TODO: Figure out if we need to use them? - for server in &mut self.connections { + for server in &mut conns { let begin = Parse::named("__pgdog_repl_begin", "BEGIN"); let commit = Parse::named("__pgdog_repl_commit", "COMMIT"); @@ -223,23 +223,42 @@ impl StreamSubscriber { .cloned() .collect(); if !omni_full.is_empty() { - self.validate_full_identity_omni_has_unique_index(&omni_full) + self.validate_full_identity_omni_has_unique_index(&mut conns, &omni_full) .await?; } + // Hand each connection to its background pipelining task. + self.connections = conns + .into_iter() + .map(PipelinedConnection::new) + .collect::, _>>()?; + Ok(()) } // Dispatch a pre-built bind to the matching shard(s). + // + // Fire-and-forget: the `Bind/Execute/Flush` is enqueued on each owning + // shard's pipelined connection without waiting for the response. Responses + // are reconciled by the per-shard background task; missed-row counting and + // error detection happen there. A cheap, non-blocking fail-fast peek surfaces + // an already-latched error so we roll back instead of piling work onto a + // transaction Postgres has already aborted. The authoritative check is the + // commit drain (`commit()`). async fn send(&mut self, val: &Shard, bind: &Bind) -> Result<(), Error> { + // Fail-fast: the replicated transaction spans every shard, so an error + // latched on any connection means whole transaction is aborted. + for conn in &self.connections { + if let Some(err) = conn.take_error() { + return Err(err); + } + } + // Locals avoid borrowing self inside the iter_mut closure. let partition = self.partition; let n_conns = self.connections.len(); - let mut conns: Vec<_> = self - .connections - .iter_mut() - .enumerate() - .filter(|(shard, _)| match val { + let targets: Vec = (0..n_conns) + .filter(|shard| match val { // With a single destination shard the router collapses Shard::All // to Direct(0), bypassing the partition ownership check. Apply // partition.owns() for all variants when there is only one connection @@ -248,48 +267,11 @@ impl StreamSubscriber { Shard::Multi(multi) if n_conns > 1 => multi.contains(shard), _ => partition.owns(*shard), }) - .map(|(_, server)| server) .collect(); - for conn in &mut conns { - conn.send(&vec![bind.clone().into(), Execute::new().into(), Flush.into()].into()) - .await?; - } - - for conn in &mut conns { - conn.flush().await?; - } - - for conn in &mut conns { - // Keep server connections always synchronized. - for _ in 0..2 { - let msg = conn.read().await?; - match msg.code() { - 'C' => { - let cmd = CommandComplete::try_from(msg)?; - let rows = cmd - .rows()? - .ok_or(Error::CommandCompleteNoRows(cmd.clone()))?; - // A direct-to-shard update indicates a row has changed on source. - // This row must exist on the destination, or we missed some data during sync. - if rows == 0 && val.is_direct() { - match cmd.tag() { - "UPDATE" => self.missed_rows.update += 1, - "DELETE" => self.missed_rows.delete += 1, - "INSERT" => self.missed_rows.insert += 1, - _ => (), - } - } - } - '2' => (), - 'E' => { - return Err(Error::PgError(Box::new(ErrorResponse::from_bytes( - msg.to_bytes(), - )?))); - } - c => return Err(Error::SendOutOfSync(c)), - } - } + let is_direct = val.is_direct(); + for &idx in &targets { + self.connections[idx].execute(bind, is_direct).await?; } Ok(()) @@ -458,37 +440,11 @@ impl StreamSubscriber { /// not in a transaction). async fn prepare_statements(&mut self, parses: &[Parse]) -> Result<(), Error> { let in_txn = self.in_transaction; - let mut msgs: Vec<_> = parses.iter().map(|p| p.clone().into()).collect(); - msgs.push(if in_txn { Flush.into() } else { Sync.into() }); - let payload = msgs.into(); - - for server in &mut self.connections { + for server in &self.connections { for p in parses { debug!("preparing \"{}\" [{}]", p.query(), server.addr()); } - server.send(&payload).await?; - } - - let num_acks = if in_txn { - parses.len() - } else { - parses.len() + 1 - }; - for server in &mut self.connections { - for _ in 0..num_acks { - let msg = server.read().await?; - trace!("[{}] --> {:?}", server.addr(), msg); - match msg.code() { - 'E' => { - return Err(Error::PgError(Box::new(ErrorResponse::from_bytes( - msg.to_bytes(), - )?))); - } - 'Z' => break, - '1' => continue, - c => return Err(Error::RelationOutOfSync(c)), - } - } + server.prepare(parses, in_txn).await?; } Ok(()) } @@ -683,26 +639,27 @@ impl StreamSubscriber { // Handle Commit message. // - // Send Sync to all shards, ensuring they close the transaction. + // Two-phase to preserve "all shards commit or none do": + // 1. Drain every shard — wait until all outstanding DML has been acked + // and confirm none errored. A `?` here (an errored shard) returns before + // any shard is `Sync`ed, so `handle()` drops every connection and Postgres + // rolls back the open implicit transactions cluster-wide. + // 2. Only once every shard is drained and clean, send `Sync` to commit. + // + // NOTE: phase 2 is not atomic across shards. If an earlier shard's `Sync` + // commits and a later shard's `Sync` fails (e.g. a deferred constraint that + // only fires at COMMIT), the earlier shard is already durable and cannot be + // rolled back. Phase 1 narrows this window to commit-time-only failures. + // TODO: use 2PC (see the `two_pc` path) for true cross-shard atomicity. async fn commit(&mut self, commit: XLogCommit) -> Result<(), Error> { - for server in &mut self.connections { - server.send_one(&Sync.into()).await?; - server.flush().await?; + // Phase 1: drain outstanding DML without committing. + for server in &self.connections { + server.drain_acks().await?; } - for server in &mut self.connections { - // Drain responses from server. - let msg = server.read().await?; - trace!("[{}] --> {:?}", server.addr(), msg); - - match msg.code() { - 'E' => { - return Err(Error::PgError(Box::new(ErrorResponse::from_bytes( - msg.to_bytes(), - )?))); - } - 'Z' => (), - c => return Err(Error::CommitOutOfSync(c)), - } + + // Phase 2: every shard is clean — commit each one. + for server in &self.connections { + server.sync_and_drain().await?; } let transaction_lsn = self.lsn; @@ -945,9 +902,13 @@ impl StreamSubscriber { self.in_transaction } - /// Get and reset missing rows. + /// Aggregate and reset the missed-row counters across all shard connections. pub(crate) fn missed_rows(&mut self) -> MissedRows { - std::mem::take(&mut self.missed_rows) + let mut total = MissedRows::default(); + for conn in &self.connections { + total.merge(conn.take_missed_rows()); + } + total } /// Verify every destination shard has a qualifying unique index for all `tables`. @@ -956,15 +917,15 @@ impl StreamSubscriber { /// Queries all shards in parallel (one bulk query per shard) then surfaces /// the complete set of missing indexes across the cluster in a single error. async fn validate_full_identity_omni_has_unique_index( - &mut self, + &self, + servers: &mut [Server], tables: &[Table], ) -> Result<(), Error> { // Fan out to all shards concurrently; each gets one IN-list query. - let per_shard: Vec> = - try_join_all(self.connections.iter_mut().map(|dest_server| { - tables_missing_unique_index(tables.iter().map(|t| &t.table), dest_server) - })) - .await?; + let per_shard: Vec> = try_join_all(servers.iter_mut().map(|dest_server| { + tables_missing_unique_index(tables.iter().map(|t| &t.table), dest_server) + })) + .await?; // Flatten; ensure_validation! deduplicates and sorts before reporting. let errors: Vec = per_shard @@ -991,6 +952,23 @@ impl MissedRows { pub(crate) fn non_zero(&self) -> bool { self.insert > 0 || self.delete > 0 || self.update > 0 } + + /// Count a direct-to-shard DML that touched 0 rows, keyed by command tag. + pub(crate) fn record(&mut self, tag: &str) { + match tag { + "UPDATE" => self.update += 1, + "DELETE" => self.delete += 1, + "INSERT" => self.insert += 1, + _ => (), + } + } + + /// Fold another shard's counts into this one. + pub(crate) fn merge(&mut self, other: MissedRows) { + self.insert += other.insert; + self.update += other.update; + self.delete += other.delete; + } } impl Display for MissedRows { From 5c7f1c4004bd62a360b13078259c665e2ce357be Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Mon, 13 Jul 2026 16:27:36 +0530 Subject: [PATCH 2/5] cosmetic changes --- .../logical/subscriber/pipeline.rs | 43 ++++++++++++------- .../replication/logical/subscriber/stream.rs | 8 ---- 2 files changed, 27 insertions(+), 24 deletions(-) diff --git a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs index 67fb54389..09c8883d6 100644 --- a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs +++ b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs @@ -77,8 +77,8 @@ impl PipelinedConnection { server, shared: shared.clone(), direct: VecDeque::new(), - prepare_sync_points: VecDeque::new(), - ready_sync_points: VecDeque::new(), + parse_ack_sync_points: VecDeque::new(), + ready_to_commit_sync_points: VecDeque::new(), drain_waiter: None, }; @@ -114,14 +114,15 @@ impl PipelinedConnection { /// uses `Flush` (must not commit the open implicit transaction); otherwise /// uses `Sync`. pub async fn prepare(&self, parses: &[Parse], in_transaction: bool) -> Result<(), Error> { - if parses.is_empty() { - return Ok(()); - } let mut messages: Vec = parses.iter().map(|p| p.clone().into()).collect(); let kind = if in_transaction { + // If in transaction we send flush and wait for the acknowledgements + // since we donot want to commit the open transaction. messages.push(Flush.into()); SyncPointKind::ParseAcks(parses.len()) } else { + // Since we are not in a transaction we send Sync and wait for postgres + // to becomes ready for the next command. messages.push(Sync.into()); SyncPointKind::ReadyForQuery }; @@ -197,9 +198,9 @@ struct Listener { // each CommandComplete. direct: VecDeque, // In-transaction prepare sync points: (remaining ParseComplete acks, waiter). - prepare_sync_points: VecDeque<(usize, oneshot::Sender<()>)>, + parse_ack_sync_points: VecDeque<(usize, oneshot::Sender<()>)>, // Commit / out-of-transaction prepare sync points, resolved on ReadyForQuery. - ready_sync_points: VecDeque>, + ready_to_commit_sync_points: VecDeque>, // Waiter that resolves once every DML ack has been read. At most one is // outstanding: `commit()` issues drain_acks sequentially, one per connection. drain_waiter: Option>, @@ -242,9 +243,14 @@ impl Listener { messages, is_direct, } => { + // If any connection has already reported and error in the shared state, + // we do not need to enqueue the command as postgres will ignore the command. if errored { return Ok(()); } + // If our command fails to be executed, we latch the error to the shared state. + // we also wake up all the sync point and resolve all the sync point waiters to + // resolve with the error. if let Err(err) = self.write(&messages).await { self.latch_error(err); self.wake_all(); @@ -268,11 +274,14 @@ impl Listener { return Err(Error::PipelineClosed); } match kind { - // `prepare()` never enqueues an empty batch, so `remaining >= 1`. + // We enqueue the number of acknowledgements we are waiting for and the waiter to resolve. SyncPointKind::ParseAcks(remaining) => { - self.prepare_sync_points.push_back((remaining, done)); + self.parse_ack_sync_points.push_back((remaining, done)); + } + // We enqueue the waiter to resolve when postgres is ready to commit. + SyncPointKind::ReadyForQuery => { + self.ready_to_commit_sync_points.push_back(done); } - SyncPointKind::ReadyForQuery => self.ready_sync_points.push_back(done), } } Command::DrainAcks { done } => { @@ -305,9 +314,11 @@ impl Listener { self.direct.clear(); self.wake_all(); } - // ParseComplete: advance the front in-transaction prepare sync point. + // as sync points arrive we decrement the remaining count and resolve th waiter + // when the remaining count becomes 0. '1' => { - let resolved = if let Some((remaining, _)) = self.prepare_sync_points.front_mut() { + let resolved = if let Some((remaining, _)) = self.parse_ack_sync_points.front_mut() + { *remaining -= 1; *remaining == 0 } else { @@ -335,9 +346,9 @@ impl Listener { self.wake_drain(); } } - // ReadyForQuery: resolve the front commit sync point. + // When Ready to commit sync point arrives we resolve the waiter. 'Z' => { - if let Some(done) = self.ready_sync_points.pop_front() { + if let Some(done) = self.ready_to_commit_sync_points.pop_front() { let _ = done.send(()); } } @@ -363,10 +374,10 @@ impl Listener { } fn wake_all(&mut self) { - while let Some((_, done)) = self.prepare_sync_points.pop_front() { + while let Some((_, done)) = self.parse_ack_sync_points.pop_front() { let _ = done.send(()); } - while let Some(done) = self.ready_sync_points.pop_front() { + while let Some(done) = self.ready_to_commit_sync_points.pop_front() { let _ = done.send(()); } self.wake_drain(); diff --git a/pgdog/src/backend/replication/logical/subscriber/stream.rs b/pgdog/src/backend/replication/logical/subscriber/stream.rs index b698eeeb7..47a84748f 100644 --- a/pgdog/src/backend/replication/logical/subscriber/stream.rs +++ b/pgdog/src/backend/replication/logical/subscriber/stream.rs @@ -237,14 +237,6 @@ impl StreamSubscriber { } // Dispatch a pre-built bind to the matching shard(s). - // - // Fire-and-forget: the `Bind/Execute/Flush` is enqueued on each owning - // shard's pipelined connection without waiting for the response. Responses - // are reconciled by the per-shard background task; missed-row counting and - // error detection happen there. A cheap, non-blocking fail-fast peek surfaces - // an already-latched error so we roll back instead of piling work onto a - // transaction Postgres has already aborted. The authoritative check is the - // commit drain (`commit()`). async fn send(&mut self, val: &Shard, bind: &Bind) -> Result<(), Error> { // Fail-fast: the replicated transaction spans every shard, so an error // latched on any connection means whole transaction is aborted. From d93587a6c92611de8bfff681c96d639a639334fd Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Mon, 13 Jul 2026 16:35:14 +0530 Subject: [PATCH 3/5] minor correction --- pgdog/src/backend/replication/logical/subscriber/pipeline.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs index 09c8883d6..4c049baa6 100644 --- a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs +++ b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs @@ -325,7 +325,7 @@ impl Listener { false }; if resolved { - if let Some((_, done)) = self.prepare_sync_points.pop_front() { + if let Some((_, done)) = self.parse_ack_sync_points.pop_front() { let _ = done.send(()); } } From 9d5f770f205e04702c47713e9a6bc6effe1c78b9 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Mon, 13 Jul 2026 16:44:51 +0530 Subject: [PATCH 4/5] making clippy happy --- .../logical/subscriber/pipeline.rs | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs index 4c049baa6..23059a03b 100644 --- a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs +++ b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs @@ -324,10 +324,10 @@ impl Listener { } else { false }; - if resolved { - if let Some((_, done)) = self.parse_ack_sync_points.pop_front() { - let _ = done.send(()); - } + if resolved + && let Some((_, done)) = self.parse_ack_sync_points.pop_front() + { + let _ = done.send(()); } } // BindComplete: nothing to account for. @@ -335,12 +335,11 @@ impl Listener { // CommandComplete: match to the DML op and count missed rows. 'C' => { let is_direct = self.direct.pop_front().unwrap_or(false); - if is_direct { - if let Ok(complete) = CommandComplete::try_from(message) { - if matches!(complete.rows(), Ok(Some(0))) { - self.shared.lock().missed.record(complete.tag()); - } - } + if is_direct + && let Ok(complete) = CommandComplete::try_from(message) + && matches!(complete.rows(), Ok(Some(0))) + { + self.shared.lock().missed.record(complete.tag()); } if self.direct.is_empty() { self.wake_drain(); From 06112832eaaf30ceb75ec542666457f12152f063 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Mon, 13 Jul 2026 16:49:33 +0530 Subject: [PATCH 5/5] making fmt happy --- pgdog/src/backend/replication/logical/subscriber/pipeline.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs index 23059a03b..65252a67d 100644 --- a/pgdog/src/backend/replication/logical/subscriber/pipeline.rs +++ b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs @@ -324,9 +324,7 @@ impl Listener { } else { false }; - if resolved - && let Some((_, done)) = self.parse_ack_sync_points.pop_front() - { + if resolved && let Some((_, done)) = self.parse_ack_sync_points.pop_front() { let _ = done.send(()); } }