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..65252a67d --- /dev/null +++ b/pgdog/src/backend/replication/logical/subscriber/pipeline.rs @@ -0,0 +1,423 @@ +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(), + parse_ack_sync_points: VecDeque::new(), + ready_to_commit_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> { + 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 + }; + 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). + parse_ack_sync_points: VecDeque<(usize, oneshot::Sender<()>)>, + // Commit / out-of-transaction prepare sync points, resolved on ReadyForQuery. + 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>, +} + +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 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(); + 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 { + // We enqueue the number of acknowledgements we are waiting for and the waiter to resolve. + SyncPointKind::ParseAcks(remaining) => { + 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); + } + } + } + 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(); + } + // 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.parse_ack_sync_points.front_mut() + { + *remaining -= 1; + *remaining == 0 + } else { + false + }; + if resolved && let Some((_, done)) = self.parse_ack_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 + && 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(); + } + } + // When Ready to commit sync point arrives we resolve the waiter. + 'Z' => { + if let Some(done) = self.ready_to_commit_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.parse_ack_sync_points.pop_front() { + let _ = done.send(()); + } + while let Some(done) = self.ready_to_commit_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..47a84748f 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,34 @@ 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). 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 +259,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 +432,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 +631,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 +894,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 +909,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 +944,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 {