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/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..d305143 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,8 +71,12 @@ 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"] } proptest = "1" tempfile = "3" tokio-postgres = "0.7" @@ -97,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/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..fec86f6 100644 --- a/src/client/worker.rs +++ b/src/client/worker.rs @@ -1,11 +1,11 @@ use bytes::Bytes; -use std::sync::atomic::{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,41 +20,48 @@ 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 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, + cell: ProgressCell, + changed: watch::Sender, } impl SharedProgress { pub fn new(start: Lsn) -> Self { Self { - applied: AtomicU64::new(start.as_u64()), + 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 { - match self - .applied - .compare_exchange_weak(cur, new, Ordering::Release, Ordering::Relaxed) - { - Ok(_) => break, - Err(observed) => cur = observed, - } + if self.cell.raise(lsn.as_u64()) { + self.changed.send_modify(|g| *g = g.wrapping_add(1)); } } + + /// 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 { + Lsn::from_u64(self.cell.take_signal()) + } } /// Events emitted by the replication worker. @@ -133,6 +140,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 +158,7 @@ impl WorkerState { stop_rx, out, metrics, + last_sent: None, } } @@ -234,6 +244,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 +307,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 +328,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 +573,7 @@ impl WorkerState { let _ = write_copy_done(stream).await; return Ok(true); } + continue; } permit = self.out.reserve() => { @@ -557,12 +591,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 +713,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 +725,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 +1268,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<()> {