elefant_client/postgres_client/replication/
stream.rs

1use super::message_reader::parse_replication_message;
2use super::messages::*;
3use crate::pool::ConnectionFactory;
4use crate::postgres_client::PostgresClient;
5use crate::protocol::frame_reader::ByteSliceWriter;
6use crate::protocol::{BackendMessage, CopyData, FrontendMessage};
7use crate::{reborrow_until_polonius, ElefantClientError};
8
9pub struct ReplicationStream<'a, F: ConnectionFactory> {
10    client: &'a mut PostgresClient<F>,
11    status_buf: Vec<u8>,
12    last_received_lsn: Lsn,
13}
14
15/// Result of reading a single raw replication message, used internally
16/// to separate parsing (which borrows the frame buffer) from actions
17/// like sending keepalive replies (which need mutable access).
18enum RawReadResult<'a> {
19    Message(ReplicationMessage<'a>),
20    KeepaliveReply(Lsn),
21    EndOfStream,
22}
23
24impl<'a, F: ConnectionFactory> ReplicationStream<'a, F> {
25    pub(super) fn new(client: &'a mut PostgresClient<F>) -> Self {
26        Self {
27            client,
28            status_buf: Vec::with_capacity(34),
29            last_received_lsn: Lsn(0),
30        }
31    }
32
33    /// Reads the next replication message from the stream.
34    ///
35    /// Returns `Ok(None)` when the server ends the replication stream
36    /// (sends `CopyDone`). Automatically replies to keepalive messages
37    /// that have `reply_requested` set, so callers never see those.
38    pub async fn next_message(
39        &mut self,
40    ) -> Result<Option<ReplicationMessage<'_>>, ElefantClientError> {
41        loop {
42            let result = {
43                let client: &mut PostgresClient<F> = reborrow_until_polonius!(&mut *self.client);
44                let msg = client.read_next_backend_message().await?;
45                match msg {
46                    BackendMessage::CopyData(cd) => {
47                        let repl_msg = parse_replication_message(cd.data)?;
48                        match &repl_msg {
49                            ReplicationMessage::XLogData(xlog) => {
50                                self.last_received_lsn = xlog.end_lsn;
51                                RawReadResult::Message(repl_msg)
52                            }
53                            ReplicationMessage::PrimaryKeepalive(ka) => {
54                                if ka.end_lsn > self.last_received_lsn {
55                                    self.last_received_lsn = ka.end_lsn;
56                                }
57                                if ka.reply_requested {
58                                    RawReadResult::KeepaliveReply(self.last_received_lsn)
59                                } else {
60                                    RawReadResult::Message(repl_msg)
61                                }
62                            }
63                        }
64                    }
65                    BackendMessage::CopyDone => RawReadResult::EndOfStream,
66                    _ => {
67                        return Err(ElefantClientError::UnexpectedBackendMessage(format!(
68                            "Expected CopyData or CopyDone during replication, got {msg:?}"
69                        )));
70                    }
71                }
72            };
73
74            match result {
75                RawReadResult::Message(msg) => return Ok(Some(msg)),
76                RawReadResult::EndOfStream => return Ok(None),
77                RawReadResult::KeepaliveReply(lsn) => {
78                    self.send_status_update(lsn, lsn, Lsn(0)).await?;
79                }
80            }
81        }
82    }
83
84    pub async fn send_status_update(
85        &mut self,
86        write_lsn: Lsn,
87        flush_lsn: Lsn,
88        apply_lsn: Lsn,
89    ) -> Result<(), ElefantClientError> {
90        // Timestamp: microseconds since PostgreSQL epoch (2000-01-01)
91        let now = std::time::SystemTime::now()
92            .duration_since(std::time::UNIX_EPOCH)
93            .unwrap_or(std::time::Duration::ZERO);
94        let pg_epoch_offset_us = 946_684_800i64 * 1_000_000;
95        let pg_timestamp = (now.as_micros() as i64) - pg_epoch_offset_us;
96
97        self.status_buf.clear();
98        let mut writer = ByteSliceWriter::new(&mut self.status_buf);
99        writer.write_u8(b'r');
100        writer.write_u64(write_lsn.0);
101        writer.write_u64(flush_lsn.0);
102        writer.write_u64(apply_lsn.0);
103        writer.write_i64(pg_timestamp);
104        writer.write_u8(0); // no reply requested
105
106        self.client
107            .connection
108            .write_frontend_message(&FrontendMessage::CopyData(CopyData {
109                data: &self.status_buf,
110            }))
111            .await?;
112        self.client.connection.flush().await?;
113        Ok(())
114    }
115}
elefant_client/postgres_client/replication/stream.rs

elefant_client/postgres_client/replication/
stream.rs
