From ec8d2b807cbc240b582edde70e144cf25496d7f8 Mon Sep 17 00:00:00 2001 From: vahid Date: Tue, 13 Oct 2026 00:34:38 +0300 Subject: [PATCH 1/2] feat: send standby status updates when the applied LSN changes update_applied_lsn now wakes the worker, which sends the new position once status_interval has passed since its last status update, also on an idle stream. Before, an idle stream reported it only at the next idle wakeup, server message or keepalive reply request. --- CHANGELOG.md | 1 + Cargo.toml | 1 + src/client/tokio_client.rs | 5 +- src/client/worker.rs | 388 +++++++++++++++++++++++++++++++++- src/config.rs | 3 + tests/integration_postgres.rs | 84 ++++++++ 6 files changed, 469 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9aac209..8a6694f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/). - **BREAKING**: these two changes are intended for the next breaking release, 0.5.0. - **BREAKING**: `PgWireError` is now `#[non_exhaustive]`. Exhaustive `match` on it needs a `_` arm; future variants can then be added without another breaking release. - **BREAKING**: refusing an oversized message is terminal. The refused payload is left on the wire rather than skipped, because it may carry WAL changes that must not be silently dropped — skipping it would let a direct `MessageReader` caller retry and continue with missing data. Every subsequent `read()` returns the same `MessageTooLarge`, so the reader and its connection must be discarded. The replication worker already ends the stream on the first error. A bogus length field therefore cannot make the reader consume an arbitrary number of bytes to resynchronize. +- **Change-driven standby feedback**: `update_applied_lsn` now wakes the worker, which sends the new applied LSN once `status_interval` has passed since its last standby status update, also on an idle stream. Previously an idle stream reported it only at the next `idle_wakeup_interval` wakeup, server message or keepalive reply request, so a consumer acknowledging after traffic stopped could hold WAL for that long. Updates coalesce into one status carrying the newest LSN; no wakeup is added while nothing new is pending, and `update_applied_lsn` stays a compare-and-swap plus one atomic load except for the first raise after the worker last looked. Keepalive replies, idle wakeup feedback and feedback under backpressure are unchanged. --- diff --git a/Cargo.toml b/Cargo.toml index 066b4a6..503f600 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -73,6 +73,7 @@ tracing-subscriber = { version = "0.3", optional = true, features = [ [dev-dependencies] hex = "0.4" +tokio = { version = "1", features = ["test-util"] } proptest = "1" tempfile = "3" tokio-postgres = "0.7" diff --git a/src/client/tokio_client.rs b/src/client/tokio_client.rs index ae34ca8..55eca14 100644 --- a/src/client/tokio_client.rs +++ b/src/client/tokio_client.rs @@ -157,8 +157,9 @@ impl ReplicationClient { /// Update the applied/durable LSN reported to the server. /// /// Semantics: call this only once you have durably persisted all events up to `lsn`. - /// This update is monotonic and cheap; wire feedback is still governed by the worker’s - /// `status_interval` and keepalive reply requests. + /// This update is monotonic and cheap. A raised LSN wakes the worker, which sends it + /// once `status_interval` has passed since its last standby status update, also on an + /// idle stream; updates in between coalesce into one status carrying the newest LSN. #[inline] pub fn update_applied_lsn(&self, lsn: Lsn) { self.progress.update_applied(lsn); diff --git a/src/client/worker.rs b/src/client/worker.rs index 901da45..5b35c98 100644 --- a/src/client/worker.rs +++ b/src/client/worker.rs @@ -1,5 +1,5 @@ use bytes::Bytes; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; use tokio::io::{AsyncRead, AsyncWrite, BufReader}; use tokio::sync::{mpsc, watch}; @@ -20,16 +20,24 @@ use crate::protocol::replication::{ /// Shared replication progress updated by the consumer and read by the worker. /// -/// Stored as an AtomicU64 so progress updates are cheap and monotonic -/// without async backpressure. +/// The applied LSN is an AtomicU64 so progress updates are cheap and monotonic +/// without async backpressure. A raise also wakes an idle worker, which sends +/// it promptly instead of at its next wakeup. Only the first raise after the +/// worker last looked bumps the watch generation; later ones coalesce into it +/// and cost one more atomic load. pub struct SharedProgress { applied: AtomicU64, + /// Set by the first raise the worker has not looked at yet. + signalled: AtomicBool, + changed: watch::Sender, } impl SharedProgress { pub fn new(start: Lsn) -> Self { Self { applied: AtomicU64::new(start.as_u64()), + signalled: AtomicBool::new(false), + changed: watch::Sender::new(0), } } @@ -46,15 +54,36 @@ impl SharedProgress { let mut cur = self.applied.load(Ordering::Relaxed); while new > cur { + // SeqCst pairs with `take_signal`: either the worker reads this + // raise, or this raise sees the signal cleared and wakes it. match self .applied - .compare_exchange_weak(cur, new, Ordering::Release, Ordering::Relaxed) + .compare_exchange_weak(cur, new, Ordering::SeqCst, Ordering::Relaxed) { - Ok(_) => break, + Ok(_) => { + if !self.signalled.load(Ordering::Relaxed) + && !self.signalled.swap(true, Ordering::SeqCst) + { + self.changed.send_modify(|g| *g = g.wrapping_add(1)); + } + break; + } Err(observed) => cur = observed, } } } + + /// Wakes when the applied LSN is raised. + pub fn subscribe(&self) -> watch::Receiver { + self.changed.subscribe() + } + + /// The worker is about to read the applied LSN: the next raise wakes it + /// again. Returns the applied LSN read after clearing the signal. + fn take_signal(&self) -> Lsn { + self.signalled.store(false, Ordering::SeqCst); + Lsn::from_u64(self.applied.load(Ordering::SeqCst)) + } } /// Events emitted by the replication worker. @@ -133,6 +162,8 @@ pub struct WorkerState { stop_rx: watch::Receiver, out: mpsc::Sender>, metrics: Arc, + /// The applied LSN in the last status update written to the server, if any. + last_sent: Option, } impl WorkerState { @@ -149,6 +180,7 @@ impl WorkerState { stop_rx, out, metrics, + last_sent: None, } } @@ -234,6 +266,8 @@ impl WorkerState { ) -> Result<()> { let mut last_status_sent = Instant::now() - self.cfg.status_interval; let mut last_applied = self.progress.load_applied(); + // Wakes the wait phase when the consumer raises the applied LSN. + let mut applied_rx = self.progress.subscribe(); // Cancellation-safe message reader, partial reads survive dropped futures. let mut reader = MessageReader::new().with_max_message_size(self.cfg.max_message_size); // How many messages to process in the tight loop before checking @@ -295,6 +329,16 @@ impl WorkerState { // mid-message. `MessageReader::read` is cancellation-safe — partial // header/payload state lives on `reader` and is preserved across the // drop, so the next iteration resumes the read without losing bytes. + // + // An applied LSN newer than the last one sent is pending: it is sent + // once `status_interval` has passed since the last status update. + // Otherwise a raise wakes the loop. Marking the watch seen and + // clearing the signal before reading the LSN means a raise after + // this point still wakes it. + applied_rx.mark_unchanged(); + let applied = self.progress.take_signal(); + let pending = self.last_sent.is_none_or(|sent| applied > sent); + let feedback_due = last_status_sent + self.cfg.status_interval; let msg = tokio::select! { biased; @@ -306,6 +350,17 @@ impl WorkerState { continue; } + _ = tokio::time::sleep_until(feedback_due), if pending => { + let applied = self.progress.load_applied(); + last_applied = applied; + self.send_feedback(stream, applied, false).await?; + last_status_sent = Instant::now(); + continue; + } + + // The sender lives in `self.progress`, so this never closes. + _ = applied_rx.changed(), if !pending => continue, + msg_result = tokio::time::timeout( self.cfg.idle_wakeup_interval, reader.read(stream), @@ -540,6 +595,7 @@ impl WorkerState { let _ = write_copy_done(stream).await; return Ok(true); } + continue; } permit = self.out.reserve() => { @@ -557,12 +613,12 @@ impl WorkerState { return Ok(false); } - _ = tokio::time::sleep_until(deadline) => { - let applied = self.progress.load_applied(); - self.send_feedback(stream, applied, false).await?; - *last_status_sent = Instant::now(); - } + _ = tokio::time::sleep_until(deadline) => {} } + // The status interval elapsed while parked. + let applied = self.progress.load_applied(); + self.send_feedback(stream, applied, false).await?; + *last_status_sent = Instant::now(); } } @@ -679,8 +735,11 @@ impl WorkerState { } /// Send standby status update to server. + /// + /// The position is recorded as sent only once the write succeeded, so a + /// failed send leaves it pending. async fn send_feedback( - &self, + &mut self, stream: &mut S, applied: Lsn, reply_requested: bool, @@ -688,6 +747,7 @@ impl WorkerState { let client_time = current_pg_timestamp(); let payload = encode_standby_status_update(applied, client_time, reply_requested); write_copy_data(stream, &payload).await?; + self.last_sent = Some(applied); self.metrics.record_feedback_sent(applied); Ok(()) } @@ -1230,4 +1290,310 @@ mod tests { let _ = stop_tx.send(true); let _ = worker_task.await; } + + // ── Change-driven standby status updates ────────────────────────────── + + /// A KeepAlive ('k') CopyData frame. + fn keepalive_frame(wal_end: u64, reply_requested: bool) -> Vec { + let mut payload = vec![b'k']; + payload.extend_from_slice(&(wal_end as i64).to_be_bytes()); + payload.extend_from_slice(&0i64.to_be_bytes()); + payload.push(u8::from(reply_requested)); + let mut frame = vec![b'd']; + frame.extend_from_slice(&((payload.len() + 4) as i32).to_be_bytes()); + frame.extend_from_slice(&payload); + frame + } + + /// A status update read by the scripted server. + #[derive(Debug, PartialEq)] + struct Status { + applied: u64, + reply: bool, + } + + /// The scripted server end of a worker's stream. + struct Server { + end: tokio::io::DuplexStream, + reader: MessageReader, + } + + impl Server { + /// The next standby status update, skipping anything else. + async fn status(&mut self) -> Status { + loop { + let msg = self.reader.read(&mut self.end).await.expect("stream open"); + if msg.tag == b'd' && msg.payload.first() == Some(&b'r') { + let p = &msg.payload; + return Status { + applied: u64::from_be_bytes(p[1..9].try_into().unwrap()), + reply: p[33] == 1, + }; + } + } + } + + /// No status update arrives within `within`. + async fn quiet(&mut self, within: std::time::Duration) -> bool { + tokio::time::timeout(within, self.status()).await.is_err() + } + + async fn send(&mut self, frame: &[u8]) { + use tokio::io::AsyncWriteExt; + self.end.write_all(frame).await.unwrap(); + self.end.flush().await.unwrap(); + } + } + + struct Running { + server: Server, + task: tokio::task::JoinHandle>, + stop: watch::Sender, + _events: mpsc::Receiver>, + } + + fn feedback_cfg() -> ReplicationConfig { + ReplicationConfig { + status_interval: std::time::Duration::from_secs(1), + idle_wakeup_interval: std::time::Duration::from_secs(10), + ..Default::default() + } + } + + /// Run `stream_loop` against a scripted server over a duplex of `capacity` + /// bytes. + fn run_stream(progress: &Arc, capacity: usize) -> Running { + run_stream_with(feedback_cfg(), progress, capacity) + } + + fn run_stream_with( + cfg: ReplicationConfig, + progress: &Arc, + capacity: usize, + ) -> Running { + let (stop, stop_rx) = watch::channel(false); + let (tx, events) = mpsc::channel(16); + let mut worker = WorkerState::new( + cfg, + Arc::clone(progress), + stop_rx, + tx, + Arc::new(ReplicationMetrics::default()), + ); + let (worker_end, end) = tokio::io::duplex(capacity); + let task = tokio::spawn(async move { + let mut buf = BufReader::with_capacity(1024, worker_end); + worker.stream_loop(&mut buf).await + }); + Running { + server: Server { + end, + reader: MessageReader::new(), + }, + task, + stop, + _events: events, + } + } + + const SECOND: std::time::Duration = std::time::Duration::from_secs(1); + + #[tokio::test(start_paused = true)] + async fn an_update_is_sent_within_the_status_interval() { + let progress = Arc::new(SharedProgress::new(Lsn(0))); + let mut r = run_stream(&progress, 64 * 1024); + assert_eq!( + r.server.status().await, + Status { + applied: 0, + reply: false + } + ); + // Idle with nothing new: no status before the idle wakeup. + assert!(r.server.quiet(3 * SECOND).await); + + // The interval since the last status has passed: sent at once. + let t = Instant::now(); + progress.update_applied(Lsn(100)); + assert_eq!( + r.server.status().await, + Status { + applied: 100, + reply: false + } + ); + assert!(t.elapsed() < SECOND); + + // The next one waits for the interval, no longer. + progress.update_applied(Lsn(200)); + let t = Instant::now(); + assert_eq!( + r.server.status().await, + Status { + applied: 200, + reply: false + } + ); + assert!(t.elapsed() <= SECOND, "took {:?}", t.elapsed()); + r.task.abort(); + } + + #[tokio::test(start_paused = true)] + async fn coalesced_updates_send_the_newest_once() { + let progress = Arc::new(SharedProgress::new(Lsn(0))); + let mut r = run_stream(&progress, 64 * 1024); + assert_eq!(r.server.status().await.applied, 0); + let t = Instant::now(); + for lsn in 1..=100 { + progress.update_applied(Lsn(lsn)); + } + assert_eq!( + r.server.status().await, + Status { + applied: 100, + reply: false + } + ); + assert!(t.elapsed() <= SECOND); + // Nothing else until the idle wakeup. + assert!(r.server.quiet(8 * SECOND).await); + r.task.abort(); + } + + #[tokio::test(start_paused = true)] + async fn an_update_during_a_send_stays_pending() { + let progress = Arc::new(SharedProgress::new(Lsn(0))); + // Too small for one status message: the first send blocks mid-write. + let mut r = run_stream(&progress, 8); + for _ in 0..10 { + tokio::task::yield_now().await; + } + progress.update_applied(Lsn(300)); + // The blocked send carries the position loaded before the update... + assert_eq!(r.server.status().await.applied, 0); + // ...and the update follows within the interval. + let t = Instant::now(); + assert_eq!(r.server.status().await.applied, 300); + assert!(t.elapsed() <= SECOND); + r.task.abort(); + } + + #[tokio::test(start_paused = true)] + async fn a_new_session_starts_with_the_newest_position() { + let progress = Arc::new(SharedProgress::new(Lsn(0))); + progress.update_applied(Lsn(500)); + let mut r = run_stream(&progress, 64 * 1024); + assert_eq!(r.server.status().await.applied, 500); + r.task.abort(); + } + + #[tokio::test(start_paused = true)] + async fn a_failed_send_leaves_the_position_for_the_next_session() { + let progress = Arc::new(SharedProgress::new(Lsn(0))); + let mut r = run_stream(&progress, 64 * 1024); + assert_eq!(r.server.status().await.applied, 0); + progress.update_applied(Lsn(700)); + // The server goes away before the pending status is written. + drop(r.server); + let ended = tokio::time::timeout(2 * SECOND, r.task).await.unwrap(); + assert!(ended.unwrap().is_err(), "the failed send ends the session"); + let mut next = run_stream(&progress, 64 * 1024); + assert_eq!(next.server.status().await.applied, 700); + next.task.abort(); + } + + #[tokio::test(start_paused = true)] + async fn a_keepalive_requesting_a_reply_is_answered_at_once() { + let progress = Arc::new(SharedProgress::new(Lsn(40))); + let mut r = run_stream(&progress, 64 * 1024); + assert_eq!(r.server.status().await.applied, 40); + // No local change, inside the status interval. + let t = Instant::now(); + r.server.send(&keepalive_frame(90, true)).await; + assert_eq!( + r.server.status().await, + Status { + applied: 40, + reply: true + } + ); + assert!(t.elapsed() < SECOND); + r.task.abort(); + } + + #[tokio::test(start_paused = true)] + async fn a_completely_idle_connection_keeps_sending_status() { + let progress = Arc::new(SharedProgress::new(Lsn(10))); + let mut r = run_stream(&progress, 64 * 1024); + let t = Instant::now(); + let mut statuses = 0; + while t.elapsed() < std::time::Duration::from_secs(65) { + assert_eq!(r.server.status().await.applied, 10); + statuses += 1; + } + // One at start, then one per idle wakeup (10 s): not none, which + // would let wal_sender_timeout disconnect, and not one per interval. + assert!((6..=8).contains(&statuses), "{statuses} statuses in 65 s"); + assert!(!r.task.is_finished()); + r.task.abort(); + } + + #[tokio::test(start_paused = true)] + async fn stopping_with_the_timer_armed_ends_at_once() { + let progress = Arc::new(SharedProgress::new(Lsn(0))); + let mut r = run_stream(&progress, 64 * 1024); + assert_eq!(r.server.status().await.applied, 0); + progress.update_applied(Lsn(900)); // pending until the interval passes + let t = Instant::now(); + r.stop.send(true).unwrap(); + let ended = tokio::time::timeout(SECOND, r.task) + .await + .expect("the worker ends without waiting for the timer"); + assert!(ended.unwrap().is_ok()); + assert!(t.elapsed() < SECOND); + } + + /// Raises from another thread race the worker clearing its signal and + /// reading the position: whatever the interleaving, the newest position of + /// every burst is sent within the status interval once the burst ends. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn racing_raises_never_strand_the_newest_position() { + use std::time::Duration; + let interval = Duration::from_millis(50); + let cfg = ReplicationConfig { + status_interval: interval, + idle_wakeup_interval: Duration::from_secs(60), + ..Default::default() + }; + let progress = Arc::new(SharedProgress::new(Lsn(0))); + let mut r = run_stream_with(cfg, &progress, 64 * 1024); + let mut newest = 0u64; + for burst in 0..50u64 { + let p = Arc::clone(&progress); + let from = newest; + let len = 1 + (burst * 7919) % 5000; + newest = std::thread::spawn(move || { + for lsn in from + 1..=from + len { + p.update_applied(Lsn(lsn)); + } + from + len + }) + .join() + .unwrap(); + let ended = Instant::now(); + loop { + let status = tokio::time::timeout(4 * interval, r.server.status()) + .await + .unwrap_or_else(|_| { + panic!("burst {burst}: {newest} not sent within {:?}", 4 * interval) + }); + if status.applied == newest { + break; + } + assert!(status.applied < newest); + } + assert!(ended.elapsed() <= 4 * interval); + } + r.task.abort(); + } } diff --git a/src/config.rs b/src/config.rs index 14e800f..cf0eecd 100644 --- a/src/config.rs +++ b/src/config.rs @@ -382,6 +382,9 @@ pub struct ReplicationConfig { /// with no incoming messages, the client will send a standby status update /// (feedback) and continue waiting. /// + /// A raised applied LSN does not wait for this wakeup: it is sent within + /// `status_interval` of `update_applied_lsn`. + /// /// This effectively bounds how long the worker can stay blocked in a read /// while idle. /// diff --git a/tests/integration_postgres.rs b/tests/integration_postgres.rs index abf02ab..0ee3711 100644 --- a/tests/integration_postgres.rs +++ b/tests/integration_postgres.rs @@ -612,6 +612,90 @@ async fn postgres_replication_batch_insert() -> Result<()> { Ok(()) } +/// The slot's confirmed_flush_lsn, read from the server. +async fn confirmed_flush(client: &tokio_postgres::Client, slot: &str) -> Result { + let row = client + .query_one( + "SELECT confirmed_flush_lsn::text FROM pg_replication_slots WHERE slot_name = $1", + &[&slot], + ) + .await?; + Lsn::parse(&row.get::<_, String>(0)).context("parse confirmed_flush_lsn") +} + +/// On an idle stream, an applied LSN reaches the server within +/// `status_interval`, not at the next idle wakeup or server keepalive. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn postgres_idle_feedback_is_change_driven() -> Result<()> { + init_tracing(); + + let host_port: u16 = std::env::var("PG_ITEST_PORT") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or_else(get_available_port); + let container = postgres_image(host_port) + .start() + .await + .expect("start postgres"); + follow_container_logs(&container).await; + + let client = wait_for_pg_ready(host_port, Duration::from_secs(30)).await?; + setup_publication_and_slot(&client, "slot_fb", "pub_fb").await?; + let base_lsn = current_wal_lsn(&client).await?; + + // status_interval 1 s, idle wakeup 15 s. + let mut repl = ReplicationClient::connect(replication_config( + host_port, "slot_fb", "pub_fb", base_lsn, None, + )) + .await?; + recv_keepalive(&mut repl, Duration::from_secs(10)).await?; + + client + .execute("INSERT INTO t(id, v) VALUES (1, 'a')", &[]) + .await?; + let deadline = Instant::now() + Duration::from_secs(10); + let commit_end = loop { + anyhow::ensure!(Instant::now() < deadline, "no commit received"); + match tokio::time::timeout(Duration::from_secs(1), repl.recv()).await { + Ok(Ok(Some(ReplicationEvent::Commit { end_lsn, .. }))) => break end_lsn, + Ok(Ok(Some(_))) | Err(_) => {} + Ok(Ok(None)) => anyhow::bail!("stream ended"), + Ok(Err(e)) => return Err(e.into()), + } + }; + + // Not acknowledged yet: the slot stays behind the commit. + tokio::time::sleep(Duration::from_secs(3)).await; + let before = confirmed_flush(&client, "slot_fb").await?; + anyhow::ensure!( + before < commit_end, + "slot advanced to {before} before the commit {commit_end} was acknowledged" + ); + + // Acknowledge on the idle stream; the server sees it within the interval. + let acked = Instant::now(); + repl.update_applied_lsn(commit_end); + loop { + let flushed = confirmed_flush(&client, "slot_fb").await?; + if flushed >= commit_end { + break; + } + anyhow::ensure!( + acked.elapsed() < Duration::from_secs(2), + "confirmed_flush_lsn {flushed} did not reach {commit_end} within 2 s" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + info!( + "confirmed_flush_lsn reached the acknowledged commit in {:?}", + acked.elapsed() + ); + + repl.stop(); + let _ = repl.join().await; + Ok(()) +} + /// Test error handling: connecting with a nonexistent slot should fail. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn postgres_replication_invalid_slot_error() -> Result<()> { From 4d1eba1ed67338752202f34d1e48368ac31a9d26 Mon Sep 17 00:00:00 2001 From: vahid Date: Tue, 13 Oct 2026 09:42:30 +0300 Subject: [PATCH 2/2] fix: never lose a wake when a raise races the worker clearing its signal The fast-path load of the signalled flag in update_applied_lsn was Relaxed, so it could still see a flag the worker had already cleared after reading the LSN before the raise: the raise then skipped the wake and the worker kept waiting until its idle wakeup. Make it SeqCst, move the lock-free core into ProgressCell, and add a loom model (run in CI) proving a raise concurrent with send or clear is sent, pending or wakes the worker. --- .github/workflows/ci.yml | 6 ++ Cargo.toml | 5 ++ src/client/mod.rs | 1 + src/client/progress.rs | 118 +++++++++++++++++++++++++++++++++++++++ src/client/worker.rs | 46 ++++----------- 5 files changed, 142 insertions(+), 34 deletions(-) create mode 100644 src/client/progress.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8271183..f3a0701 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -45,6 +45,12 @@ jobs: env: RUSTDOCFLAGS: "-D warnings" + - name: Loom model (progress signalling) + run: cargo test --release --lib loom_model + env: + RUSTFLAGS: --cfg pgwire_loom + CARGO_TARGET_DIR: target/loom + msrv: name: MSRV check runs-on: ubuntu-latest diff --git a/Cargo.toml b/Cargo.toml index 503f600..d305143 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,6 +71,9 @@ tracing-subscriber = { version = "0.3", optional = true, features = [ "env-filter", ] } +[target.'cfg(pgwire_loom)'.dev-dependencies] +loom = "0.7" + [dev-dependencies] hex = "0.4" tokio = { version = "1", features = ["test-util"] } @@ -98,3 +101,5 @@ name = "with_mtls" path = "examples/with_mtls.rs" required-features = ["tls-rustls"] +[lints.rust] +unexpected_cfgs = { level = "warn", check-cfg = ['cfg(pgwire_loom)'] } diff --git a/src/client/mod.rs b/src/client/mod.rs index 3d3a464..c7305ea 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -71,6 +71,7 @@ //! ``` mod metrics; +mod progress; mod tokio_client; mod worker; diff --git a/src/client/progress.rs b/src/client/progress.rs new file mode 100644 index 0000000..0f2b711 --- /dev/null +++ b/src/client/progress.rs @@ -0,0 +1,118 @@ +//! The lock-free core of [`super::worker::SharedProgress`]: the applied LSN and +//! the flag that decides which raise wakes the worker. +//! +//! The consumer raises the applied LSN; the worker clears the flag and then +//! reads the LSN before deciding whether anything is pending. Both sides use +//! SeqCst: either the worker's read sees a raise, or that raise sees the flag +//! cleared and asks for a wake. A raise is never both unseen and unsignalled. +//! The `loom` model in this module checks that for every interleaving. + +#[cfg(all(test, pgwire_loom))] +use loom::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +#[cfg(not(all(test, pgwire_loom)))] +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + +pub(crate) struct ProgressCell { + applied: AtomicU64, + /// Set by the first raise the worker has not looked at yet. + signalled: AtomicBool, +} + +impl ProgressCell { + pub(crate) fn new(start: u64) -> Self { + Self { + applied: AtomicU64::new(start), + signalled: AtomicBool::new(false), + } + } + + pub(crate) fn load(&self) -> u64 { + self.applied.load(Ordering::Acquire) + } + + /// Monotonic raise to `new` (a lower value is a no-op). Returns true when + /// the caller must wake the worker: this is the first raise since the + /// worker last looked. + pub(crate) fn raise(&self, new: u64) -> bool { + let mut cur = self.applied.load(Ordering::Relaxed); + while new > cur { + match self + .applied + .compare_exchange_weak(cur, new, Ordering::SeqCst, Ordering::Relaxed) + { + Ok(_) => { + // The fast-path load must be SeqCst too: a Relaxed load + // may still see a flag the worker already cleared, after + // the worker read the LSN before this raise. + return !self.signalled.load(Ordering::SeqCst) + && !self.signalled.swap(true, Ordering::SeqCst); + } + Err(observed) => cur = observed, + } + } + false + } + + /// The worker is about to decide whether a position is pending: clear the + /// signal, then read the applied LSN. A raise after the clear wakes it. + pub(crate) fn take_signal(&self) -> u64 { + self.signalled.store(false, Ordering::SeqCst); + self.applied.load(Ordering::SeqCst) + } +} + +#[cfg(all(test, pgwire_loom))] +mod loom_model { + use super::ProgressCell; + use loom::sync::atomic::{AtomicU64, Ordering}; + use loom::sync::Arc; + use loom::thread; + + /// The worker's wait-phase cycle against two concurrent raises. The watch + /// generation stands in for the wake channel (bumped with SeqCst by a + /// raise that asks for a wake; read before `take_signal`, as the worker + /// marks the watch seen). In every interleaving, once both sides are done, + /// the newest raised position has been sent, is still pending, or has a + /// wake waiting: it is never lost. + #[test] + fn a_raise_concurrent_with_send_or_clear_is_never_lost() { + loom::model(|| { + let cell = Arc::new(ProgressCell::new(0)); + let generation = Arc::new(AtomicU64::new(0)); + + let consumer = { + let (cell, generation) = (Arc::clone(&cell), Arc::clone(&generation)); + thread::spawn(move || { + for lsn in [1, 2] { + if cell.raise(lsn) { + generation.fetch_add(1, Ordering::SeqCst); + } + } + }) + }; + + // Two worker iterations: each marks the watch seen, clears the + // signal, reads the position and sends it if it is pending. + let mut last_sent = 0; + let mut seen = 0; + let mut pending = false; + for _ in 0..2 { + seen = generation.load(Ordering::SeqCst); + let applied = cell.take_signal(); + pending = applied > last_sent; + if pending { + last_sent = applied; // the send + pending = false; + } + } + consumer.join().unwrap(); + + let newest = cell.load(); + let woken = generation.load(Ordering::SeqCst) != seen; + assert!( + last_sent == newest || pending || woken, + "lost raise: newest {newest}, sent {last_sent}, no wake pending" + ); + }); + } +} diff --git a/src/client/worker.rs b/src/client/worker.rs index 5b35c98..fec86f6 100644 --- a/src/client/worker.rs +++ b/src/client/worker.rs @@ -1,11 +1,11 @@ use bytes::Bytes; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; use tokio::io::{AsyncRead, AsyncWrite, BufReader}; use tokio::sync::{mpsc, watch}; use tokio::time::Instant; use super::metrics::ReplicationMetrics; +use super::progress::ProgressCell; use crate::config::ReplicationConfig; use crate::error::{PgWireError, Result}; use crate::lsn::Lsn; @@ -20,56 +20,35 @@ use crate::protocol::replication::{ /// Shared replication progress updated by the consumer and read by the worker. /// -/// The applied LSN is an AtomicU64 so progress updates are cheap and monotonic -/// without async backpressure. A raise also wakes an idle worker, which sends -/// it promptly instead of at its next wakeup. Only the first raise after the -/// worker last looked bumps the watch generation; later ones coalesce into it -/// and cost one more atomic load. +/// The applied LSN lives in a lock-free [`ProgressCell`] so progress updates +/// are cheap and monotonic without async backpressure. A raise also wakes an +/// idle worker, which sends it promptly instead of at its next wakeup. Only the +/// first raise after the worker last looked bumps the watch generation; later +/// ones coalesce into it and cost one more atomic load. pub struct SharedProgress { - applied: AtomicU64, - /// Set by the first raise the worker has not looked at yet. - signalled: AtomicBool, + cell: ProgressCell, changed: watch::Sender, } impl SharedProgress { pub fn new(start: Lsn) -> Self { Self { - applied: AtomicU64::new(start.as_u64()), - signalled: AtomicBool::new(false), + cell: ProgressCell::new(start.as_u64()), changed: watch::Sender::new(0), } } #[inline] pub fn load_applied(&self) -> Lsn { - Lsn::from_u64(self.applied.load(Ordering::Acquire)) + Lsn::from_u64(self.cell.load()) } /// Monotonic update: if `lsn` is lower than the currently stored applied LSN, /// this is a no-op. #[inline] pub fn update_applied(&self, lsn: Lsn) { - let new = lsn.as_u64(); - let mut cur = self.applied.load(Ordering::Relaxed); - - while new > cur { - // SeqCst pairs with `take_signal`: either the worker reads this - // raise, or this raise sees the signal cleared and wakes it. - match self - .applied - .compare_exchange_weak(cur, new, Ordering::SeqCst, Ordering::Relaxed) - { - Ok(_) => { - if !self.signalled.load(Ordering::Relaxed) - && !self.signalled.swap(true, Ordering::SeqCst) - { - self.changed.send_modify(|g| *g = g.wrapping_add(1)); - } - break; - } - Err(observed) => cur = observed, - } + if self.cell.raise(lsn.as_u64()) { + self.changed.send_modify(|g| *g = g.wrapping_add(1)); } } @@ -81,8 +60,7 @@ impl SharedProgress { /// The worker is about to read the applied LSN: the next raise wakes it /// again. Returns the applied LSN read after clearing the signal. fn take_signal(&self) -> Lsn { - self.signalled.store(false, Ordering::SeqCst); - Lsn::from_u64(self.applied.load(Ordering::SeqCst)) + Lsn::from_u64(self.cell.take_signal()) } }