pgwire_replication/client/
tokio_client.rs

1use crate::config::ReplicationConfig;
2use crate::error::{PgWireError, Result};
3use crate::lsn::Lsn;
4
5use tokio::net::TcpStream;
6#[cfg(unix)]
7use tokio::net::UnixStream;
8
9use tokio::sync::{mpsc, watch};
10use tokio::task::JoinHandle;
11
12use std::sync::Arc;
13
14#[cfg(not(feature = "tls-rustls"))]
15use crate::config::SslMode;
16
17use super::metrics::ReplicationMetrics;
18use super::worker::{
19    GateCommand, QueryRow, ReplicationEvent, ReplicationEventReceiver, SharedProgress, WorkerState,
20};
21
22/// The IDENTIFY_SYSTEM answer of a replication session.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct IdentifySystem {
25    /// The cluster's system identifier (decimal text).
26    pub system_id: String,
27    /// The server's current timeline.
28    pub timeline: u32,
29    /// The current WAL flush location.
30    pub xlogpos: Lsn,
31    /// The session's database (`None` for a physical session).
32    pub dbname: Option<String>,
33}
34
35/// The TIMELINE_HISTORY answer: the history file of a timeline.
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct TimelineHistory {
38    /// The history file's name, such as `00000002.history`.
39    pub filename: String,
40    /// The history file's contents.
41    pub content: bytes::Bytes,
42}
43
44/// PostgreSQL logical replication client.
45///
46/// This client spawns a background worker task that maintains the replication
47/// connection and streams events to the consumer via a bounded channel.
48///
49/// # Example
50///
51/// ```no_run
52/// use pgwire_replication::client::{ReplicationClient, ReplicationEvent};
53/// use pgwire_replication::config::ReplicationConfig;
54///
55/// #[tokio::main]
56/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
57///     let config = ReplicationConfig::new(
58///         "localhost",
59///         "postgres",
60///         "password",
61///         "mydb",
62///         "my_slot",
63///         "my_pub",
64///     );
65///
66///     let mut client = ReplicationClient::connect(config).await?;
67///
68///     while let Some(ev) = client.recv().await? {
69///         match ev {
70///             ReplicationEvent::XLogData { data, wal_end, .. } => {
71///                 process_change(&data);
72///                 client.update_applied_lsn(wal_end);
73///             }
74///             ReplicationEvent::KeepAlive { .. } => {}
75///             ReplicationEvent::StoppedAt { reached } => {
76///                 println!("Reached stop LSN: {reached}");
77///                 break;
78///             }
79///             _ => {}
80///         }
81///     }
82///
83///     Ok(())
84/// }
85///
86/// fn process_change(_data: &bytes::Bytes) {
87///     // user-defined
88/// }
89/// ```
90pub struct ReplicationClient {
91    rx: ReplicationEventReceiver,
92    progress: Arc<SharedProgress>,
93    stop_tx: watch::Sender<bool>,
94    metrics: Arc<ReplicationMetrics>,
95    join: Option<JoinHandle<std::result::Result<(), PgWireError>>>,
96    /// Becomes `true` once the server accepted START_REPLICATION.
97    started: watch::Receiver<bool>,
98    /// Commands to a gated session before START_REPLICATION. `None`: not
99    /// gated, or already started.
100    gate: Option<mpsc::Sender<GateCommand>>,
101}
102
103impl ReplicationClient {
104    /// Connect to PostgreSQL and start streaming replication events.
105    ///
106    /// This establishes a TCP connection (optionally upgrading to TLS),
107    /// authenticates, and starts the replication stream. Events are buffered
108    /// in a channel of size `config.buffer_events`.
109    ///
110    /// # Errors
111    ///
112    /// Returns an error if:
113    /// - TCP connection fails
114    /// - TLS handshake fails (when enabled)
115    /// - Authentication fails
116    /// - Replication slot doesn't exist
117    /// - Publication doesn't exist
118    /// - Unix socket does not exist (when host starts with `/`)
119    /// - TLS requested with Unix socket connection
120    pub async fn connect(cfg: ReplicationConfig) -> Result<Self> {
121        Self::spawn(cfg, None)
122    }
123
124    /// Connect and authenticate like [`connect`](Self::connect), but do not
125    /// start replication: the session waits for
126    /// [`simple_query`](Self::simple_query) (and the IDENTIFY_SYSTEM /
127    /// TIMELINE_HISTORY helpers) until [`start`](Self::start) sends
128    /// START_REPLICATION. Everything runs on the one authenticated session, so
129    /// what is read before the start describes the server the stream then
130    /// opens on. Dropping the client before `start` ends the session.
131    pub async fn connect_gated(cfg: ReplicationConfig) -> Result<Self> {
132        let (gate_tx, gate_rx) = mpsc::channel(1);
133        let mut client = Self::spawn(cfg, Some(gate_rx))?;
134        client.gate = Some(gate_tx);
135        Ok(client)
136    }
137
138    fn spawn(cfg: ReplicationConfig, gate: Option<mpsc::Receiver<GateCommand>>) -> Result<Self> {
139        let (tx, rx) = mpsc::channel(cfg.buffer_events);
140
141        // Progress is shared via atomics: cheap, monotonic, no async backpressure.
142        let progress = Arc::new(SharedProgress::with_timing(
143            cfg.ack_lsn.unwrap_or(cfg.start_lsn),
144            cfg.feedback_observer.is_some(),
145        ));
146
147        let (stop_tx, stop_rx) = watch::channel(false);
148        let (started_tx, started_rx) = watch::channel(false);
149
150        let metrics = Arc::new(ReplicationMetrics::default());
151
152        let progress_for_worker = Arc::clone(&progress);
153        let metrics_for_worker = Arc::clone(&metrics);
154        let cfg_for_worker = cfg.clone();
155
156        let join = tokio::spawn(async move {
157            let mut worker = WorkerState::new(
158                cfg_for_worker,
159                progress_for_worker,
160                stop_rx,
161                tx,
162                metrics_for_worker,
163            );
164            worker.notify_started(started_tx);
165            if let Some(gate) = gate {
166                worker.gate_before_start(gate);
167            }
168            let res = run_worker(&mut worker, &cfg).await;
169            if let Err(ref e) = res {
170                tracing::error!("replication worker terminated with error: {e}");
171            }
172            res
173        });
174
175        Ok(Self {
176            rx,
177            progress,
178            stop_tx,
179            metrics,
180            join: Some(join),
181            started: started_rx,
182            gate: None,
183        })
184    }
185
186    /// Run `sql` with the simple query protocol on a gated session before
187    /// [`start`](Self::start); returns its rows (text format). A server error
188    /// leaves the session usable.
189    pub async fn simple_query(&mut self, sql: &str) -> Result<Vec<QueryRow>> {
190        let gate = self.gate.as_ref().ok_or_else(|| {
191            PgWireError::Internal("simple_query needs a gated session that has not started".into())
192        })?;
193        let (reply_tx, reply_rx) = tokio::sync::oneshot::channel();
194        let sent = gate
195            .send(GateCommand::Query {
196                sql: sql.to_string(),
197                reply: reply_tx,
198            })
199            .await;
200        if sent.is_ok() {
201            if let Ok(result) = reply_rx.await {
202                return result;
203            }
204        }
205        // The worker ended (connect or authentication failed, or the session
206        // broke): report why.
207        self.gate = None;
208        match self.handle_worker_shutdown().await {
209            Err(e) => Err(e),
210            Ok(_) => Err(PgWireError::Internal(
211                "replication session ended before the query ran".into(),
212            )),
213        }
214    }
215
216    /// IDENTIFY_SYSTEM on a gated session.
217    pub async fn identify_system(&mut self) -> Result<IdentifySystem> {
218        let rows = self.simple_query("IDENTIFY_SYSTEM").await?;
219        let row = single_row(rows, 4, "IDENTIFY_SYSTEM")?;
220        let text = |i: usize| -> Result<String> {
221            row[i]
222                .as_ref()
223                .map(|b| String::from_utf8_lossy(b).into_owned())
224                .ok_or_else(|| {
225                    PgWireError::Protocol(format!("IDENTIFY_SYSTEM: column {i} is NULL"))
226                })
227        };
228        let timeline = text(1)?.parse::<u32>().map_err(|e| {
229            PgWireError::Protocol(format!("IDENTIFY_SYSTEM: invalid timeline: {e}"))
230        })?;
231        let xlogpos = Lsn::parse(&text(2)?)
232            .map_err(|e| PgWireError::Protocol(format!("IDENTIFY_SYSTEM: invalid xlogpos: {e}")))?;
233        Ok(IdentifySystem {
234            system_id: text(0)?,
235            timeline,
236            xlogpos,
237            dbname: text(3).ok(),
238        })
239    }
240
241    /// TIMELINE_HISTORY for `timeline` on a gated session (the server has no
242    /// history file for timeline 1).
243    pub async fn timeline_history(&mut self, timeline: u32) -> Result<TimelineHistory> {
244        let rows = self
245            .simple_query(&format!("TIMELINE_HISTORY {timeline}"))
246            .await?;
247        let mut row = single_row(rows, 2, "TIMELINE_HISTORY")?;
248        let null =
249            |i: usize| PgWireError::Protocol(format!("TIMELINE_HISTORY: column {i} is NULL"));
250        let content = row[1].take().ok_or_else(|| null(1))?;
251        let filename = row[0].take().ok_or_else(|| null(0))?;
252        Ok(TimelineHistory {
253            filename: String::from_utf8_lossy(&filename).into_owned(),
254            content,
255        })
256    }
257
258    /// Send START_REPLICATION on a gated session; then
259    /// [`wait_started`](Self::wait_started) reports whether the server
260    /// accepted it.
261    pub async fn start(&mut self) -> Result<()> {
262        let gate = self.gate.take().ok_or_else(|| {
263            PgWireError::Internal("start needs a gated session that has not started".into())
264        })?;
265        // A closed gate means the worker already ended; wait_started reports
266        // its error.
267        let _ = gate.send(GateCommand::Start).await;
268        Ok(())
269    }
270
271    /// Wait until the server has accepted START_REPLICATION: from then on the
272    /// replication slot is held by this session, and no other session can
273    /// acquire or advance it. [`connect`](Self::connect) returns before that,
274    /// as connecting, authentication and START_REPLICATION run in the
275    /// background worker. Returns the worker's error if it ended first.
276    pub async fn wait_started(&mut self) -> Result<()> {
277        if self.started.wait_for(|started| *started).await.is_ok() {
278            return Ok(());
279        }
280        match self.handle_worker_shutdown().await {
281            Err(e) => Err(e),
282            Ok(_) => Err(PgWireError::Internal(
283                "replication worker ended before replication started".into(),
284            )),
285        }
286    }
287
288    /// Receive the next replication event.
289    ///
290    /// - `Ok(Some(event))` => received an event
291    /// - `Ok(None)`        => replication ended normally (stop requested or stop_at_lsn reached)
292    /// - `Err(e)`          => replication ended abnormally
293    pub async fn recv(&mut self) -> Result<Option<ReplicationEvent>> {
294        match self.rx.recv().await {
295            Some(Ok(ev)) => Ok(Some(ev)),
296            Some(Err(e)) => Err(e),
297            None => self.handle_worker_shutdown().await,
298        }
299    }
300
301    async fn handle_worker_shutdown(&mut self) -> Result<Option<ReplicationEvent>> {
302        let join = self
303            .join
304            .take()
305            .ok_or_else(|| PgWireError::Internal("replication worker already joined".into()))?;
306
307        match join.await {
308            Ok(Ok(())) => Ok(None),
309            Ok(Err(e)) => Err(e),
310            Err(join_err) => Err(PgWireError::Task(format!(
311                "replication worker panicked: {join_err}"
312            ))),
313        }
314    }
315
316    /// Update the applied/durable LSN reported to the server.
317    ///
318    /// Semantics: call this only once you have durably persisted all events up to `lsn`.
319    /// This update is monotonic and cheap. A raised LSN wakes the worker, which sends it
320    /// once `status_interval` has passed since its last standby status update, also on an
321    /// idle stream; updates in between coalesce into one status carrying the newest LSN.
322    #[inline]
323    pub fn update_applied_lsn(&self, lsn: Lsn) {
324        self.progress.update_applied(lsn);
325    }
326
327    /// When the oldest applied-position raise not yet reported to the server
328    /// happened, with a feedback observer configured; `None` otherwise or
329    /// when everything raised was reported. Readable after the session ended,
330    /// so a caller replacing a failed session can carry it over with
331    /// [`carry_feedback_pending`](Self::carry_feedback_pending).
332    pub fn feedback_pending_since(&self) -> Option<std::time::Instant> {
333        self.progress.pending_since()
334    }
335
336    /// Carry over a pending start from a session this one replaces (from its
337    /// [`feedback_pending_since`](Self::feedback_pending_since)): the next
338    /// update that advances the reported position reports its wait from
339    /// there, or from an older start of this session. No effect without a
340    /// feedback observer.
341    pub fn carry_feedback_pending(&self, since: std::time::Instant) {
342        self.progress.carry_pending(since);
343    }
344
345    /// Returns a handle to the live replication metrics.
346    ///
347    /// The returned `Arc` shares the same counters the background worker
348    /// updates, so reads reflect current progress. Cheap to clone and call
349    /// repeatedly; nothing here blocks the worker.
350    #[inline]
351    pub fn metrics(&self) -> Arc<ReplicationMetrics> {
352        Arc::clone(&self.metrics)
353    }
354
355    /// Request the worker to stop gracefully.
356    ///
357    /// After calling this, [`recv()`](Self::recv) will return remaining buffered
358    /// events, then `Ok(None)` once the worker exits cleanly.
359    ///
360    /// This sends a CopyDone message to the server to cleanly terminate
361    /// the replication stream.
362    #[inline]
363    pub fn stop(&self) {
364        let _ = self.stop_tx.send(true);
365    }
366
367    pub fn is_running(&self) -> bool {
368        self.join
369            .as_ref()
370            .map(|j| !j.is_finished())
371            .unwrap_or(false)
372    }
373
374    /// Wait for the worker task to complete and return its result.
375    ///
376    /// This consumes the client. Use this for diagnostics or to ensure
377    /// clean shutdown after calling [`stop()`](Self::stop).
378    pub async fn join(mut self) -> Result<()> {
379        let join = self
380            .join
381            .take()
382            .ok_or_else(|| PgWireError::Task("worker already joined".into()))?;
383
384        match join.await {
385            Ok(inner) => inner,
386            Err(e) => Err(PgWireError::Task(format!("join error: {e}"))),
387        }
388    }
389
390    /// Abort the worker task immediately.
391    ///
392    /// This is a hard cancel and does not send CopyDone.
393    /// Prefer `stop()`/`shutdown()` for graceful termination.
394    pub fn abort(&mut self) {
395        if let Some(join) = self.join.take() {
396            join.abort();
397        }
398    }
399
400    /// Request a graceful stop and wait for the worker to exit.
401    pub async fn shutdown(&mut self) -> Result<()> {
402        self.stop();
403
404        // Drain events until the worker closes the channel.
405        while let Some(msg) = self.rx.recv().await {
406            match msg {
407                Ok(_ev) => {} //discard; caller can drain themselves if they need events
408                Err(e) => return Err(e),
409            }
410        }
411
412        self.join_mut().await
413    }
414
415    /// Wait for the worker task to complete and return its result.
416    async fn join_mut(&mut self) -> Result<()> {
417        let join = self
418            .join
419            .take()
420            .ok_or_else(|| PgWireError::Task("worker already joined".into()))?;
421
422        match join.await {
423            Ok(inner) => inner,
424            Err(e) => Err(PgWireError::Task(format!("join error: {e}"))),
425        }
426    }
427}
428
429impl Drop for ReplicationClient {
430    fn drop(&mut self) {
431        let _ = self.stop_tx.send(true);
432
433        // We cannot .await here. Prefer to detach a join in the background
434        // so the worker can exit cleanly without being aborted.
435        if let Some(join) = self.join.take() {
436            match tokio::runtime::Handle::try_current() {
437                Ok(handle) => {
438                    handle.spawn(async move {
439                        let _ = join.await;
440                    });
441                }
442                Err(_) => {
443                    // No Tokio runtime available (dropping outside async context).
444                    // Fall back to abort to avoid a potentially unbounded leaked task.
445                    tracing::debug!(
446                        "dropping ReplicationClient outside a Tokio runtime; aborting worker task"
447                    );
448                    join.abort();
449                }
450            }
451        }
452    }
453}
454
455/// The one row of a replication command's answer, with `columns` columns.
456fn single_row(rows: Vec<QueryRow>, columns: usize, command: &str) -> Result<QueryRow> {
457    let mut rows = rows.into_iter();
458    match (rows.next(), rows.next()) {
459        (Some(row), None) if row.len() >= columns => Ok(row),
460        _ => Err(PgWireError::Protocol(format!(
461            "{command}: expected one row of {columns} columns"
462        ))),
463    }
464}
465
466async fn run_worker(worker: &mut WorkerState, cfg: &ReplicationConfig) -> Result<()> {
467    #[cfg(unix)]
468    if cfg.is_unix_socket() {
469        if cfg.tls.mode.requires_tls() {
470            return Err(PgWireError::Tls(
471                "TLS is not supported over Unix domain sockets".into(),
472            ));
473        }
474
475        let path = cfg.unix_socket_path();
476        let mut stream = UnixStream::connect(&path).await.map_err(|e| {
477            PgWireError::Io(std::sync::Arc::new(std::io::Error::new(
478                e.kind(),
479                format!("failed to connect to Unix socket {}: {e}", path.display()),
480            )))
481        })?;
482
483        return worker.run_on_stream(&mut stream).await;
484    }
485
486    let tcp = TcpStream::connect((cfg.host.as_str(), cfg.port)).await?;
487    tcp.set_nodelay(true)?;
488
489    #[cfg(feature = "tls-rustls")]
490    {
491        use crate::tls::rustls::{maybe_upgrade_to_tls, MaybeTlsStream};
492        let upgraded = maybe_upgrade_to_tls(tcp, &cfg.tls, &cfg.host).await?;
493        match upgraded {
494            MaybeTlsStream::Plain(mut s) => worker.run_on_stream(&mut s).await,
495            MaybeTlsStream::Tls(mut s) => worker.run_on_stream(s.as_mut()).await,
496        }
497    }
498
499    #[cfg(not(feature = "tls-rustls"))]
500    {
501        if !matches!(cfg.tls.mode, SslMode::Disable) {
502            return Err(PgWireError::Tls("tls-rustls feature not enabled".into()));
503        }
504        let mut s = tcp;
505        worker.run_on_stream(&mut s).await
506    }
507}
pgwire_replication/client/tokio_client.rs

pgwire_replication/client/
tokio_client.rs
