pgsync_postgres/replication/
stream.rs

1//! `EventStream` / `StreamFactory` implementations over a replication connection.
2
3use super::conn::ReplicationConnection;
4use super::pgoutput::{LogicalMessage, ReplicationMessage, decode, encode_standby_status};
5use super::slot::{identify_system, read_slot, start_replication_command};
6use crate::url::ConnectionInfo;
7use pgsync_core::checkpoint::Checkpoint;
8use pgsync_core::config::SafetyConfig;
9use pgsync_core::event::{Op, SourceEvent, now_pg_micros};
10use pgsync_core::traits::{EventStream, Feedback, StreamFactory};
11use pgsync_core::{Error, Lsn, Result};
12use std::sync::Arc;
13use std::time::Duration;
14use tracing::{debug, info};
15
16pub struct PgEventStream {
17    conn: ReplicationConnection,
18    /// Bound on one feedback write; a peer that stops reading is treated as gone.
19    write_timeout: Duration,
20}
21
22impl PgEventStream {
23    pub fn new(conn: ReplicationConnection, write_timeout: Duration) -> Self {
24        PgEventStream {
25            conn,
26            write_timeout,
27        }
28    }
29
30    pub async fn close(self) {
31        self.conn.close().await;
32    }
33}
34
35impl EventStream for PgEventStream {
36    async fn next_event(&mut self) -> Result<SourceEvent> {
37        loop {
38            let data = self.conn.next_copy_data().await?;
39            match decode(data)? {
40                ReplicationMessage::PrimaryKeepalive { wal_end, .. } => {
41                    return Ok(SourceEvent::Keepalive { wal_end });
42                }
43                ReplicationMessage::XLogData { body, .. } => match body {
44                    LogicalMessage::Begin {
45                        final_lsn,
46                        commit_ts,
47                        xid,
48                    } => {
49                        return Ok(SourceEvent::Begin {
50                            final_lsn,
51                            commit_ts,
52                            xid,
53                        });
54                    }
55                    LogicalMessage::Commit {
56                        flags,
57                        commit_lsn,
58                        end_lsn,
59                        commit_ts,
60                    } => {
61                        if flags != 0 {
62                            return Err(Error::unsupported(format!(
63                                "commit message with flags {flags:#x}"
64                            )));
65                        }
66                        return Ok(SourceEvent::Commit {
67                            commit_lsn,
68                            end_lsn,
69                            commit_ts,
70                        });
71                    }
72                    LogicalMessage::Relation(r) => return Ok(SourceEvent::Relation(Arc::new(r))),
73                    LogicalMessage::Insert { relation, new } => {
74                        return Ok(SourceEvent::Op(Op::Insert { relation, new }));
75                    }
76                    LogicalMessage::Update {
77                        relation,
78                        key,
79                        old,
80                        new,
81                    } => {
82                        return Ok(SourceEvent::Op(Op::Update {
83                            relation,
84                            key,
85                            old,
86                            new,
87                        }));
88                    }
89                    LogicalMessage::Delete { relation, key, old } => {
90                        return Ok(SourceEvent::Op(Op::Delete { relation, key, old }));
91                    }
92                    LogicalMessage::Truncate {
93                        relations,
94                        cascade,
95                        restart_identity,
96                    } => {
97                        return Ok(SourceEvent::Op(Op::Truncate {
98                            relations,
99                            cascade,
100                            restart_identity,
101                        }));
102                    }
103                    LogicalMessage::Origin { .. } | LogicalMessage::Type { .. } => continue,
104                    LogicalMessage::Message {
105                        prefix,
106                        transactional,
107                        ..
108                    } => {
109                        debug!(
110                            prefix,
111                            transactional, "ignoring logical decoding message (not a data change)"
112                        );
113                        continue;
114                    }
115                    LogicalMessage::Unsupported { tag } => {
116                        return Err(Error::unsupported(format!(
117                            "pgoutput message {:?} is not supported by this version of PGSync",
118                            tag as char
119                        )));
120                    }
121                },
122            }
123        }
124    }
125
126    async fn send_feedback(&mut self, fb: Feedback) -> Result<()> {
127        let payload = encode_standby_status(
128            fb.written,
129            fb.flushed,
130            fb.applied,
131            now_pg_micros(),
132            fb.reply_requested,
133        );
134        match tokio::time::timeout(self.write_timeout, self.conn.send_copy_data(&payload)).await {
135            Ok(r) => r,
136            Err(_) => Err(Error::transient(format!(
137                "feedback write did not complete within {:?}; source not reading",
138                self.write_timeout
139            ))),
140        }
141    }
142}
143
144/// Opens replication streams for `run`, verifying source identity and slot continuity.
145pub struct PgStreamFactory {
146    pub info: ConnectionInfo,
147    pub safety: SafetyConfig,
148    pub timeout: Duration,
149    pub slot: String,
150    pub publication: String,
151    pub max_message_bytes: usize,
152}
153
154impl StreamFactory for PgStreamFactory {
155    type Stream = PgEventStream;
156
157    async fn connect(&mut self, cp: &Checkpoint) -> Result<(PgEventStream, Lsn)> {
158        let mut conn =
159            ReplicationConnection::connect(&self.info, &self.safety, self.timeout).await?;
160        conn.set_max_message_bytes(self.max_message_bytes);
161        let sys = identify_system(&mut conn).await?;
162        if sys.system_id != cp.source_system_id {
163            return Err(Error::resync_required(format!(
164                "source system identifier changed (checkpoint {}, server {}); continuity cannot be proven",
165                cp.source_system_id, sys.system_id
166            )));
167        }
168        if sys.dbname.as_deref() != Some(cp.source_dbname.as_str()) {
169            return Err(Error::preflight(format!(
170                "connected to database {:?} but the job was bootstrapped against {:?}",
171                sys.dbname, cp.source_dbname
172            )));
173        }
174        let Some(slot) = read_slot(&mut conn, &self.slot).await? else {
175            return Err(Error::resync_required(format!(
176                "replication slot {:?} does not exist; the job must be re-bootstrapped",
177                self.slot
178            )));
179        };
180        verify_slot_continuity(&slot, cp)?;
181        let cmd = start_replication_command(&self.slot, cp.committed_lsn, &self.publication);
182        conn.start_replication(&cmd).await?;
183        info!(slot = %self.slot, from = %cp.committed_lsn, confirmed_flush = ?slot.confirmed_flush_lsn, "START_REPLICATION accepted");
184        Ok((PgEventStream::new(conn, self.timeout), sys.xlogpos))
185    }
186}
187
188/// Fail with `ResyncRequired` whenever the slot cannot deliver every transaction after the
189/// checkpoint, or with `Preflight` when it is the wrong kind of slot.
190pub fn verify_slot_continuity(slot: &super::slot::SlotInfo, cp: &Checkpoint) -> Result<()> {
191    if slot.slot_type.as_deref() != Some("logical") || slot.plugin.as_deref() != Some("pgoutput") {
192        return Err(Error::preflight(format!(
193            "slot {:?} is not a logical pgoutput slot (type {:?}, plugin {:?})",
194            slot.slot_name, slot.slot_type, slot.plugin
195        )));
196    }
197    if slot.database.as_deref() != Some(cp.source_dbname.as_str()) {
198        return Err(Error::preflight(format!(
199            "slot {:?} belongs to database {:?}, expected {:?}",
200            slot.slot_name, slot.database, cp.source_dbname
201        )));
202    }
203    if slot.wal_status.as_deref() == Some("lost") || slot.invalidation_reason.is_some() {
204        return Err(Error::resync_required(format!(
205            "slot {:?} has lost WAL (wal_status={:?}, invalidation_reason={:?}); resynchronisation required",
206            slot.slot_name, slot.wal_status, slot.invalidation_reason
207        )));
208    }
209    if let Some(confirmed) = slot.confirmed_flush_lsn
210        && confirmed > cp.acked_lsn
211    {
212        return Err(Error::resync_required(format!(
213            "slot {:?} confirmed_flush_lsn {} is beyond the checkpoint's acknowledged position {}; \
214             another consumer advanced the slot and transactions may be missing",
215            slot.slot_name, confirmed, cp.acked_lsn
216        )));
217    }
218    Ok(())
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224    use crate::replication::slot::SlotInfo;
225    use pgsync_core::ErrorKind;
226
227    fn slot() -> SlotInfo {
228        SlotInfo {
229            slot_name: "s".into(),
230            plugin: Some("pgoutput".into()),
231            slot_type: Some("logical".into()),
232            database: Some("app".into()),
233            confirmed_flush_lsn: Some(Lsn(100)),
234            wal_status: Some("reserved".into()),
235            ..Default::default()
236        }
237    }
238    fn cp() -> Checkpoint {
239        let mut cp = Checkpoint::new_job("j", "sys", "app", "s", "p");
240        cp.record_commit(Lsn(90), Lsn(100), 1, 1);
241        cp
242    }
243
244    #[test]
245    fn continuity_checks() {
246        verify_slot_continuity(&slot(), &cp()).unwrap();
247        let mut s = slot();
248        s.confirmed_flush_lsn = Some(Lsn(101));
249        assert_eq!(
250            verify_slot_continuity(&s, &cp()).unwrap_err().kind(),
251            ErrorKind::ResyncRequired
252        );
253        let mut s = slot();
254        s.wal_status = Some("lost".into());
255        assert_eq!(
256            verify_slot_continuity(&s, &cp()).unwrap_err().kind(),
257            ErrorKind::ResyncRequired
258        );
259        let mut s = slot();
260        s.invalidation_reason = Some("wal_removed".into());
261        assert_eq!(
262            verify_slot_continuity(&s, &cp()).unwrap_err().kind(),
263            ErrorKind::ResyncRequired
264        );
265        let mut s = slot();
266        s.plugin = Some("wal2json".into());
267        assert_eq!(
268            verify_slot_continuity(&s, &cp()).unwrap_err().kind(),
269            ErrorKind::Preflight
270        );
271        let mut s = slot();
272        s.database = Some("other".into());
273        assert_eq!(
274            verify_slot_continuity(&s, &cp()).unwrap_err().kind(),
275            ErrorKind::Preflight
276        );
277    }
278}
pgsync_postgres/replication/stream.rs

pgsync_postgres/replication/
stream.rs
