datum_cdc/
source.rs

1use std::{
2    sync::{
3        Arc, Mutex,
4        atomic::{AtomicU64, Ordering},
5    },
6    time::{Duration, Instant},
7};
8
9use datum::{Channel, Source, SourceWithContext};
10use pgwire_replication::{ReplicationClient, ReplicationEvent};
11use tokio::sync::{mpsc, oneshot};
12
13use crate::{
14    CdcCheckpointHandle, CdcCheckpointStore, CdcError, CdcLag, CdcOffset, CdcResult, CdcStart,
15    ChangeEvent, FeedbackCommand, FeedbackCoordinator, FeedbackMode, PgLsn, PostgresCdcConfig,
16    ReconnectSettings, SlotLifecycle, SourceMetadata, TransactionMeta, admin,
17    config::SourceSettings,
18    pgoutput::{DecodedChange, PgOutputDecoder},
19    runtime,
20};
21
22/// Entry point for CDC source builders.
23pub struct CdcSource;
24
25impl CdcSource {
26    #[must_use]
27    pub fn postgres() -> PostgresCdcBuilder {
28        PostgresCdcBuilder::default()
29    }
30}
31
32/// Builder for a PostgreSQL pgoutput source.
33#[derive(Clone)]
34pub struct PostgresCdcBuilder {
35    postgres: Option<PostgresCdcConfig>,
36    slot: Option<String>,
37    publication: Option<String>,
38    start: CdcStart,
39    feedback: FeedbackMode,
40    slot_lifecycle: SlotLifecycle,
41    buffer_capacity: usize,
42    replication_buffer_events: usize,
43    status_interval: Duration,
44    idle_wakeup_interval: Duration,
45    validate_publication: bool,
46    checkpoint_store: Option<Arc<dyn CdcCheckpointStore>>,
47    reconnect: Option<ReconnectSettings>,
48}
49
50impl Default for PostgresCdcBuilder {
51    fn default() -> Self {
52        Self {
53            postgres: None,
54            slot: None,
55            publication: None,
56            start: CdcStart::default(),
57            feedback: FeedbackMode::default(),
58            slot_lifecycle: SlotLifecycle::default(),
59            buffer_capacity: 8192,
60            replication_buffer_events: 8192,
61            status_interval: Duration::from_secs(1),
62            idle_wakeup_interval: Duration::from_secs(10),
63            validate_publication: true,
64            checkpoint_store: None,
65            reconnect: Some(ReconnectSettings::default()),
66        }
67    }
68}
69
70impl PostgresCdcBuilder {
71    #[must_use]
72    pub fn connect(mut self, config: PostgresCdcConfig) -> Self {
73        self.postgres = Some(config);
74        self
75    }
76
77    pub fn connect_url(self, url: impl AsRef<str>) -> CdcResult<Self> {
78        Ok(self.connect(PostgresCdcConfig::from_url(url)?))
79    }
80
81    #[must_use]
82    pub fn slot(mut self, slot: impl Into<String>) -> Self {
83        self.slot = Some(slot.into());
84        self
85    }
86
87    #[must_use]
88    pub fn publication(mut self, publication: impl Into<String>) -> Self {
89        self.publication = Some(publication.into());
90        self
91    }
92
93    #[must_use]
94    pub fn start_from(mut self, start: CdcStart) -> Self {
95        self.start = start;
96        self
97    }
98
99    #[must_use]
100    pub fn feedback(mut self, feedback: FeedbackMode) -> Self {
101        self.feedback = feedback;
102        self
103    }
104
105    #[must_use]
106    pub fn slot_lifecycle(mut self, lifecycle: SlotLifecycle) -> Self {
107        self.slot_lifecycle = lifecycle;
108        self
109    }
110
111    #[must_use]
112    pub fn buffer_capacity(mut self, capacity: usize) -> Self {
113        self.buffer_capacity = capacity;
114        self
115    }
116
117    #[must_use]
118    pub fn replication_buffer_events(mut self, capacity: usize) -> Self {
119        self.replication_buffer_events = capacity;
120        self
121    }
122
123    #[must_use]
124    pub fn status_interval(mut self, interval: Duration) -> Self {
125        self.status_interval = interval;
126        self
127    }
128
129    #[must_use]
130    pub fn idle_wakeup_interval(mut self, interval: Duration) -> Self {
131        self.idle_wakeup_interval = interval;
132        self
133    }
134
135    #[must_use]
136    pub fn validate_publication(mut self, enabled: bool) -> Self {
137        self.validate_publication = enabled;
138        self
139    }
140
141    #[must_use]
142    pub fn checkpoint_store<S>(mut self, store: S) -> Self
143    where
144        S: CdcCheckpointStore,
145    {
146        self.checkpoint_store = Some(Arc::new(store));
147        self
148    }
149
150    #[must_use]
151    pub fn checkpoint_store_arc(mut self, store: Arc<dyn CdcCheckpointStore>) -> Self {
152        self.checkpoint_store = Some(store);
153        self
154    }
155
156    #[must_use]
157    pub fn reconnect(mut self, settings: ReconnectSettings) -> Self {
158        self.reconnect = Some(settings);
159        self
160    }
161
162    #[must_use]
163    pub fn disable_reconnect(mut self) -> Self {
164        self.reconnect = None;
165        self
166    }
167
168    pub fn build(self) -> CdcResult<Source<ChangeEvent, CdcHandle>> {
169        let settings = Arc::new(self.finalize()?);
170        Ok(source_from_settings(settings))
171    }
172
173    pub fn build_with_context(
174        self,
175    ) -> CdcResult<SourceWithContext<ChangeEvent, CdcOffset, CdcHandle>> {
176        Ok(self
177            .build()?
178            .as_source_with_context(|event| event.lsn.clone()))
179    }
180
181    pub fn restart_factory(
182        self,
183    ) -> CdcResult<impl Fn() -> Source<ChangeEvent, CdcHandle> + Clone + Send + Sync + 'static>
184    {
185        let settings = Arc::new(self.finalize()?);
186        Ok(move || source_from_settings(Arc::clone(&settings)))
187    }
188
189    fn finalize(self) -> CdcResult<SourceSettings> {
190        if self.buffer_capacity == 0 {
191            return Err(CdcError::Config(
192                "CDC source buffer capacity must be greater than zero".into(),
193            ));
194        }
195        if self.replication_buffer_events == 0 {
196            return Err(CdcError::Config(
197                "replication event buffer capacity must be greater than zero".into(),
198            ));
199        }
200        let postgres = self
201            .postgres
202            .ok_or_else(|| CdcError::Config("PostgreSQL connection is required".into()))?;
203        if !matches!(postgres.tls, crate::CdcTlsConfig::Disable) {
204            return Err(CdcError::Config(
205                "TLS replication is not enabled in this workspace build because pgwire-replication 0.3.2 forces rustls/ring and conflicts with Datum's aws-lc-rs rustls provider; use sslmode=disable/prefer for this MVP"
206                    .into(),
207            ));
208        }
209        let slot = self
210            .slot
211            .ok_or_else(|| CdcError::Config("replication slot is required".into()))?;
212        let publication = self
213            .publication
214            .ok_or_else(|| CdcError::Config("publication is required".into()))?;
215        Ok(SourceSettings {
216            source: SourceMetadata {
217                database: postgres.database.clone(),
218                slot: slot.clone(),
219                publication: publication.clone(),
220            },
221            postgres,
222            slot,
223            publication,
224            start: self.start,
225            feedback: self.feedback,
226            slot_lifecycle: self.slot_lifecycle,
227            buffer_capacity: self.buffer_capacity,
228            replication_buffer_events: self.replication_buffer_events,
229            status_interval: self.status_interval,
230            idle_wakeup_interval: self.idle_wakeup_interval,
231            validate_publication: self.validate_publication,
232            checkpoint_store: self.checkpoint_store,
233            reconnect: self.reconnect,
234        })
235    }
236}
237
238fn source_from_settings(settings: Arc<SourceSettings>) -> Source<ChangeEvent, CdcHandle> {
239    Source::channel(settings.buffer_capacity).map_materialized_value(move |channel| {
240        let (commands, command_rx) = mpsc::unbounded_channel();
241        let health = Arc::new(HealthState::default());
242        let checkpoint = CdcCheckpointHandle::new(
243            settings.slot.clone(),
244            settings.checkpoint_store.clone(),
245            commands.clone(),
246        );
247        let task_settings = Arc::clone(&settings);
248        let task_health = Arc::clone(&health);
249        runtime::spawn(async move {
250            run_carrier(task_settings, channel, command_rx, task_health).await;
251        });
252        CdcHandle {
253            settings: Arc::clone(&settings),
254            checkpoint,
255            commands,
256            health,
257        }
258    })
259}
260
261/// Materialized management handle for a CDC source.
262pub struct CdcHandle {
263    settings: Arc<SourceSettings>,
264    checkpoint: CdcCheckpointHandle,
265    commands: mpsc::UnboundedSender<FeedbackCommand>,
266    health: Arc<HealthState>,
267}
268
269impl CdcHandle {
270    #[must_use]
271    pub fn checkpoint_handle(&self) -> CdcCheckpointHandle {
272        self.checkpoint.clone()
273    }
274
275    #[must_use]
276    pub fn health(&self) -> CdcHealth {
277        self.health.snapshot()
278    }
279
280    pub fn stop(&self) -> CdcResult<()> {
281        self.commands
282            .send(FeedbackCommand::Stop)
283            .map_err(|_| CdcError::ChannelClosed)
284    }
285
286    pub async fn drop_slot(&self) -> CdcResult<()> {
287        self.drop_slot_inner(false).await
288    }
289
290    pub async fn force_drop_slot(&self) -> CdcResult<()> {
291        self.drop_slot_inner(true).await
292    }
293
294    pub async fn lag(&self) -> CdcResult<CdcLag> {
295        admin::sample_lag(&self.settings.postgres, &self.settings.slot).await
296    }
297
298    async fn drop_slot_inner(&self, force: bool) -> CdcResult<()> {
299        let (reply, wait) = oneshot::channel();
300        if self
301            .commands
302            .send(FeedbackCommand::DropSlot { force, reply })
303            .is_ok()
304        {
305            return wait
306                .await
307                .map_err(|err| CdcError::Runtime(format!("drop-slot reply lost: {err}")))?;
308        }
309        admin::drop_slot(
310            &self.settings.postgres,
311            &self.settings.slot,
312            self.settings.slot_lifecycle,
313            force,
314        )
315        .await
316    }
317}
318
319#[derive(Debug, Clone, Default, PartialEq, Eq)]
320pub struct CdcHealth {
321    pub running: bool,
322    pub last_error: Option<String>,
323    pub last_received_lsn: PgLsn,
324    pub latest_server_wal_end: PgLsn,
325    pub last_feedback_lsn: PgLsn,
326    pub emitted_events: u64,
327    pub reconnects: u64,
328}
329
330#[derive(Default)]
331struct HealthState {
332    snapshot: Mutex<CdcHealth>,
333    emitted_events: AtomicU64,
334    reconnects: AtomicU64,
335}
336
337impl HealthState {
338    fn snapshot(&self) -> CdcHealth {
339        let mut snapshot = self
340            .snapshot
341            .lock()
342            .expect("CDC health state poisoned")
343            .clone();
344        snapshot.emitted_events = self.emitted_events.load(Ordering::Acquire);
345        snapshot.reconnects = self.reconnects.load(Ordering::Acquire);
346        snapshot
347    }
348
349    fn set_running(&self, running: bool) {
350        self.update(|health| health.running = running);
351    }
352
353    fn set_error(&self, error: impl ToString) {
354        let error = error.to_string();
355        self.update(|health| health.last_error = Some(error));
356    }
357
358    fn clear_error(&self) {
359        self.update(|health| health.last_error = None);
360    }
361
362    fn record_received(&self, lsn: PgLsn) {
363        self.update(|health| health.last_received_lsn = lsn);
364    }
365
366    fn record_server_wal_end(&self, lsn: PgLsn) {
367        self.update(|health| health.latest_server_wal_end = lsn);
368    }
369
370    fn record_feedback(&self, lsn: PgLsn) {
371        self.update(|health| health.last_feedback_lsn = lsn);
372    }
373
374    fn record_emitted(&self) {
375        self.emitted_events.fetch_add(1, Ordering::AcqRel);
376    }
377
378    fn record_reconnect(&self) {
379        self.reconnects.fetch_add(1, Ordering::AcqRel);
380    }
381
382    fn update(&self, f: impl FnOnce(&mut CdcHealth)) {
383        let mut health = self.snapshot.lock().expect("CDC health state poisoned");
384        f(&mut health);
385    }
386}
387
388#[derive(Debug, Clone, Copy, PartialEq, Eq)]
389enum CarrierStop {
390    SourceClosed,
391    Stopped,
392    DroppedSlot,
393    ReplicationEnded,
394}
395
396async fn run_carrier(
397    settings: Arc<SourceSettings>,
398    channel: Channel<ChangeEvent>,
399    mut commands: mpsc::UnboundedReceiver<FeedbackCommand>,
400    health: Arc<HealthState>,
401) {
402    health.set_running(true);
403    let mut next_backoff = settings
404        .reconnect
405        .as_ref()
406        .map(|reconnect| reconnect.min_backoff);
407    let deadline = settings.start_deadline();
408    loop {
409        match run_once(&settings, &channel, &mut commands, &health).await {
410            Ok(CarrierStop::SourceClosed | CarrierStop::Stopped | CarrierStop::DroppedSlot) => {
411                break;
412            }
413            Ok(CarrierStop::ReplicationEnded) => {
414                if settings.reconnect.is_none() {
415                    break;
416                }
417            }
418            Err(error) => {
419                health.set_error(&error);
420                if !should_reconnect(&error, settings.reconnect.as_ref(), deadline) {
421                    break;
422                }
423            }
424        }
425
426        let Some(reconnect) = settings.reconnect.as_ref() else {
427            break;
428        };
429        let delay = next_backoff.unwrap_or(reconnect.min_backoff);
430        tokio::time::sleep(delay).await;
431        health.record_reconnect();
432        next_backoff = Some(delay.saturating_mul(2).min(reconnect.max_backoff));
433    }
434    channel.close();
435    health.set_running(false);
436}
437
438fn should_reconnect(
439    error: &CdcError,
440    reconnect: Option<&ReconnectSettings>,
441    deadline: Option<Instant>,
442) -> bool {
443    if reconnect.is_none() || deadline.is_some_and(|deadline| Instant::now() >= deadline) {
444        return false;
445    }
446    matches!(
447        error,
448        CdcError::Replication(_) | CdcError::Postgres(_) | CdcError::Io(_)
449    )
450}
451
452async fn run_once(
453    settings: &SourceSettings,
454    channel: &Channel<ChangeEvent>,
455    commands: &mut mpsc::UnboundedReceiver<FeedbackCommand>,
456    health: &HealthState,
457) -> CdcResult<CarrierStop> {
458    let start_lsn = admin::prepare_slot_and_start_lsn(settings).await?;
459    let mut feedback = FeedbackCoordinator::new(settings.slot.clone(), start_lsn);
460    let replication_config = settings.postgres.replication_config(
461        &settings.slot,
462        &settings.publication,
463        start_lsn,
464        settings.status_interval,
465        settings.idle_wakeup_interval,
466        settings.replication_buffer_events,
467    );
468    let mut client = ReplicationClient::connect(replication_config).await?;
469    health.clear_error();
470    let mut decoder = PgOutputDecoder::default();
471    let mut tx = None;
472
473    loop {
474        tokio::select! {
475            Some(command) = commands.recv() => {
476                if let Some(stop) = handle_command(command, settings, &mut client, &mut feedback, health).await? {
477                    return Ok(stop);
478                }
479            }
480            event = client.recv() => {
481                let Some(event) = event? else {
482                    return Ok(CarrierStop::ReplicationEnded);
483                };
484                let mut context = EventContext {
485                    settings,
486                    channel,
487                    client: &mut client,
488                    decoder: &mut decoder,
489                    feedback: &mut feedback,
490                    tx: &mut tx,
491                    health,
492                };
493                if let Some(stop) = handle_event(event, &mut context).await? {
494                    return Ok(stop);
495                }
496            }
497        }
498    }
499}
500
501async fn handle_command(
502    command: FeedbackCommand,
503    settings: &SourceSettings,
504    client: &mut ReplicationClient,
505    feedback: &mut FeedbackCoordinator,
506    health: &HealthState,
507) -> CdcResult<Option<CarrierStop>> {
508    match command {
509        FeedbackCommand::Checkpoint(offset) => {
510            match settings.feedback {
511                FeedbackMode::AfterCheckpoint => {
512                    if let Some(lsn) = feedback.checkpoint(&offset)? {
513                        client.update_applied_lsn(lsn.into());
514                        health.record_feedback(lsn);
515                    }
516                }
517            }
518            Ok(None)
519        }
520        FeedbackCommand::Stop => {
521            client.stop();
522            Ok(Some(CarrierStop::Stopped))
523        }
524        FeedbackCommand::DropSlot { force, reply } => {
525            client.stop();
526            let _ = client.shutdown().await;
527            let result = admin::drop_slot(
528                &settings.postgres,
529                &settings.slot,
530                settings.slot_lifecycle,
531                force,
532            )
533            .await;
534            let stop = if result.is_ok() {
535                CarrierStop::DroppedSlot
536            } else {
537                CarrierStop::Stopped
538            };
539            let _ = reply.send(result);
540            Ok(Some(stop))
541        }
542    }
543}
544
545struct EventContext<'a> {
546    settings: &'a SourceSettings,
547    channel: &'a Channel<ChangeEvent>,
548    client: &'a mut ReplicationClient,
549    decoder: &'a mut PgOutputDecoder,
550    feedback: &'a mut FeedbackCoordinator,
551    tx: &'a mut Option<PendingTransaction>,
552    health: &'a HealthState,
553}
554
555async fn handle_event(
556    event: ReplicationEvent,
557    context: &mut EventContext<'_>,
558) -> CdcResult<Option<CarrierStop>> {
559    match event {
560        ReplicationEvent::KeepAlive { wal_end, .. } => {
561            context.health.record_server_wal_end(wal_end.into());
562        }
563        ReplicationEvent::Begin {
564            final_lsn,
565            xid,
566            commit_time_micros,
567        } => {
568            *context.tx = Some(PendingTransaction {
569                xid,
570                final_lsn: final_lsn.into(),
571                commit_time_micros,
572                changes: Vec::new(),
573            });
574        }
575        ReplicationEvent::XLogData {
576            wal_start,
577            wal_end,
578            data,
579            ..
580        } => {
581            context.health.record_received(wal_end.into());
582            let changes = context.decoder.decode(&data)?;
583            if !changes.is_empty() {
584                let pending = context.tx.as_mut().ok_or_else(|| {
585                    CdcError::Parse(format!(
586                        "pgoutput row data at WAL {wal_start} arrived outside a transaction"
587                    ))
588                })?;
589                pending.changes.extend(changes);
590            }
591        }
592        ReplicationEvent::Commit {
593            lsn,
594            end_lsn,
595            commit_time_micros,
596        } => {
597            let Some(mut pending) = context.tx.take() else {
598                return Ok(None);
599            };
600            let _begin_final_lsn = pending.final_lsn;
601            pending.commit_time_micros = commit_time_micros;
602            let commit_lsn = PgLsn::from(lsn);
603            let tx_end_lsn = PgLsn::from(end_lsn);
604            let event_count = u32::try_from(pending.changes.len()).map_err(|_| {
605                CdcError::Parse("transaction emitted more than u32::MAX events".into())
606            })?;
607            if event_count == 0 {
608                // An empty transaction has nothing to checkpoint downstream, so
609                // route it through the feedback coordinator: it may only advance
610                // PostgreSQL feedback once every earlier transaction is
611                // confirmed, otherwise the slot could skip an unconfirmed
612                // transaction on restart or reconnect (at-least-once violation).
613                if let Some(lsn) = context.feedback.note_empty_tx(tx_end_lsn)? {
614                    context.client.update_applied_lsn(lsn.into());
615                    context.health.record_feedback(lsn);
616                }
617                return Ok(None);
618            }
619            context.feedback.register_tx(tx_end_lsn, event_count)?;
620            for (index, change) in pending.changes.into_iter().enumerate() {
621                let event = build_change_event(
622                    context.settings,
623                    change,
624                    CommittedEventMeta {
625                        xid: pending.xid,
626                        commit_time_micros: pending.commit_time_micros,
627                        commit_lsn,
628                        tx_end_lsn,
629                        event_index: u32::try_from(index)
630                            .expect("event index bounded by event_count"),
631                        event_count,
632                    },
633                );
634                if context.channel.send(event).await.is_err() {
635                    return Ok(Some(CarrierStop::SourceClosed));
636                }
637                context.health.record_emitted();
638            }
639        }
640        ReplicationEvent::Message { .. } | ReplicationEvent::StoppedAt { .. } => {}
641    }
642    Ok(None)
643}
644
645#[derive(Debug, Clone, Copy)]
646struct CommittedEventMeta {
647    xid: u32,
648    commit_time_micros: i64,
649    commit_lsn: PgLsn,
650    tx_end_lsn: PgLsn,
651    event_index: u32,
652    event_count: u32,
653}
654
655fn build_change_event(
656    settings: &SourceSettings,
657    change: DecodedChange,
658    meta: CommittedEventMeta,
659) -> ChangeEvent {
660    let offset = CdcOffset {
661        slot: settings.slot.clone(),
662        tx_end_lsn: meta.tx_end_lsn,
663        commit_lsn: meta.commit_lsn,
664        xid: meta.xid,
665        event_index: meta.event_index,
666        event_count: meta.event_count,
667    };
668    let tx = TransactionMeta {
669        xid: meta.xid,
670        commit_time_micros: meta.commit_time_micros,
671        event_index: meta.event_index,
672        event_count: meta.event_count,
673    };
674    ChangeEvent {
675        source: settings.source.clone(),
676        schema: change.relation.schema.clone(),
677        table: change.relation.table.clone(),
678        op: change.op,
679        before: change.before,
680        after: change.after,
681        truncate: change.truncate,
682        lsn: offset,
683        tx,
684        relation: change.relation,
685    }
686}
687
688#[derive(Debug)]
689struct PendingTransaction {
690    xid: u32,
691    final_lsn: PgLsn,
692    commit_time_micros: i64,
693    changes: Vec<DecodedChange>,
694}
695
696#[cfg(test)]
697mod tests {
698    use super::*;
699    use crate::MemoryCheckpointStore;
700
701    #[test]
702    fn builder_requires_connection_slot_and_publication() {
703        let err = match CdcSource::postgres().build() {
704            Ok(_) => panic!("builder unexpectedly succeeded"),
705            Err(err) => err,
706        };
707        assert!(matches!(err, CdcError::Config(_)));
708    }
709
710    #[test]
711    fn builder_creates_context_source() {
712        let config = PostgresCdcConfig::from_url("postgresql://datum@127.0.0.1:5432/db").unwrap();
713        let source = CdcSource::postgres()
714            .connect(config)
715            .slot("slot")
716            .publication("pub")
717            .checkpoint_store(MemoryCheckpointStore::new())
718            .build_with_context();
719        assert!(source.is_ok());
720    }
721}
datum_cdc/source.rs

datum_cdc/
source.rs
