laminar_connectors/postgres/cdc/source/
lifecycle.rs

1//! Source contract, startup, polling, durable feedback, and shutdown ownership.
2
3use std::collections::BTreeMap;
4
5use arrow_schema::SchemaRef;
6use async_trait::async_trait;
7
8use crate::checkpoint::SourceCheckpoint;
9use crate::config::{ConnectorConfig, ConnectorState};
10use crate::connector::{
11    ConnectorTaskTracker, SourceBatch, SourceConnector, SourceConsistency, SourceContract,
12    SourceInputMode, SourcePosition, SourceRowPositionCapability, SourceStart, SourceTopology,
13};
14use crate::error::ConnectorError;
15
16use super::super::config::{OutputMode, PostgresCdcConfig};
17use super::super::schema_resolution::{committed_relation, declared_primary_key, declared_schema};
18use super::super::typed_rows::RowBuilder;
19use super::checkpoint::{parse_cursor, validate_live_binding, write_cursor, CursorPhase};
20use super::claim::incarnation;
21use super::decoding::Decoded;
22use super::startup::{prepare, ReaderRuntime, StartInputs, StartPhase};
23use super::{reap_postgres_reader, Arc, ChunkedSnapshot, Lsn, Notify, Phase, PostgresCdcSource};
24
25/// Minimum spacing of live publication/table revalidation at checkpoint commits.
26const CONTRACT_CHECK_INTERVAL: std::time::Duration = std::time::Duration::from_secs(30);
27
28fn parsed_config(
29    current: &PostgresCdcConfig,
30    config: &ConnectorConfig,
31) -> Result<PostgresCdcConfig, ConnectorError> {
32    if config.properties().is_empty() {
33        current.validate()?;
34        Ok(current.clone())
35    } else {
36        PostgresCdcConfig::from_config(config)
37    }
38}
39
40impl PostgresCdcSource {
41    /// Decode queued WAL within one bounded work quantum and emit whole committed transactions.
42    fn poll_streaming(
43        &mut self,
44        max_records: usize,
45    ) -> Result<Option<SourceBatch>, ConnectorError> {
46        self.decode_queued(max_records)?;
47        let batch = self.drain_committed(max_records)?;
48        if max_records > 0 && (!self.pending_payloads.is_empty() || !self.committed.is_empty()) {
49            self.data_ready.notify_one();
50        }
51        self.metrics
52            .set_replication_lag_bytes(self.replication_lag_bytes());
53        Ok(batch)
54    }
55
56    /// Decode queued WAL within one bounded work quantum; `true` when any payload was decoded.
57    pub(super) fn decode_queued(&mut self, max_records: usize) -> Result<bool, ConnectorError> {
58        self.fail_on_terminal_wal_error()?;
59        let payload_budget = max_records.max(1);
60        let mut processed = 0_usize;
61        let mut reader_closed = false;
62        while processed < payload_budget {
63            let payload = if let Some(payload) = self.pending_payloads.pop_front() {
64                payload
65            } else {
66                match self.wal_rx.as_ref().map(crossfire::AsyncRx::try_recv) {
67                    Some(Ok(payload)) => payload,
68                    Some(Err(crossfire::TryRecvError::Empty)) | None => break,
69                    Some(Err(crossfire::TryRecvError::Disconnected)) => {
70                        reader_closed = true;
71                        break;
72                    }
73                }
74            };
75            match self.process_owned_wal_payload(payload) {
76                Ok(Decoded::Applied) => processed += 1,
77                Ok(Decoded::Deferred(payload)) => {
78                    self.pending_payloads.push_front(payload);
79                    break;
80                }
81                Err(error) => return Err(self.fail(error)),
82            }
83        }
84        // A full quantum may hide queued work behind a coalesced notification; keep one
85        // payload so the next poll is self-notified instead of waiting on the reader.
86        if processed == payload_budget && self.pending_payloads.is_empty() {
87            match self.wal_rx.as_ref().map(crossfire::AsyncRx::try_recv) {
88                Some(Ok(payload)) => self.pending_payloads.push_back(payload),
89                Some(Err(crossfire::TryRecvError::Disconnected)) => reader_closed = true,
90                Some(Err(crossfire::TryRecvError::Empty)) | None => {}
91            }
92        }
93        #[cfg(test)]
94        self.process_pending_messages()?;
95        self.fail_on_terminal_wal_error()?;
96        if reader_closed && self.committed.is_empty() && self.pending_payloads.is_empty() {
97            return Err(self.fail(ConnectorError::ReadError(
98                "WAL reader task terminated unexpectedly — replication stream lost".to_string(),
99            )));
100        }
101        Ok(processed > 0)
102    }
103
104    /// Launch the replication reader of the owned slot at `start_lsn`.
105    pub(super) async fn begin_streaming(&mut self, start_lsn: Lsn) -> Result<(), ConnectorError> {
106        let (Some(binding), Some(claim)) = (&self.checkpoint_binding, &self.claim) else {
107            return Err(self.fail(ConnectorError::Internal(
108                "PostgreSQL CDC streams without a binding or a slot claim".into(),
109            )));
110        };
111        let launched = super::startup::launch_reader(
112            &self.task_owner,
113            Arc::clone(&self.data_ready),
114            &self.config,
115            claim,
116            &self.incarnation,
117            binding,
118            start_lsn,
119        )
120        .await;
121        match launched {
122            Ok(runtime) => {
123                self.install_reader(start_lsn, runtime);
124                Ok(())
125            }
126            Err(error) => Err(self.fail(error)),
127        }
128    }
129
130    fn install_reader(&mut self, start_lsn: Lsn, runtime: ReaderRuntime) {
131        self.wal_rx = Some(runtime.wal_rx);
132        self.wal_byte_budget = Some(runtime.wal_byte_budget);
133        self.wal_terminal_error = Some(runtime.terminal_error);
134        self.reader_handle = Some(runtime.reader_handle);
135        self.reader_shutdown = Some(runtime.shutdown_tx);
136        self.applied_lsn = Some(runtime.applied_lsn);
137        self.polled_lsn = start_lsn;
138        self.write_lsn = self.write_lsn.max(start_lsn);
139        self.confirmed_flush_lsn = start_lsn;
140        self.metrics.set_confirmed_flush_lsn(start_lsn.as_u64());
141        self.data_ready.notify_one();
142    }
143
144    /// The resumable cursor of the current phase: always present once started, so barriers
145    /// keep flowing while the claim, the slot creation, and the snapshot or stream gate are
146    /// pending.
147    pub(super) fn cursor(&self) -> SourceCheckpoint {
148        // A source that failed while installing its slot still reports where it stands: a
149        // barrier captured between a retried poll error and the fault may commit.
150        let phase = match (&self.phase, self.consistent_point) {
151            (Phase::Claiming(_), _) | (Phase::Idle, None) => CursorPhase::Claimed,
152            (Phase::Snapshot(snapshot), Some(consistent_point)) => {
153                CursorPhase::Snapshot(snapshot.position(consistent_point, self.polled_lsn))
154            }
155            (
156                Phase::AwaitingStream { .. } | Phase::Streaming | Phase::Idle,
157                Some(consistent_point),
158            ) => CursorPhase::Streaming {
159                consistent_point,
160                lsn: self.polled_lsn,
161            },
162            (Phase::AwaitingStream { .. } | Phase::Snapshot(_) | Phase::Streaming, None) => {
163                return write_cursor(&self.config, None);
164            }
165        };
166        let position = match (&self.claim, &self.checkpoint_binding) {
167            (Some(claim), Some(binding)) => Some((claim, binding, phase)),
168            _ => None,
169        };
170        write_cursor(&self.config, position)
171    }
172
173    /// Re-read the live contract and report orphaned slots at most every
174    /// [`CONTRACT_CHECK_INTERVAL`].
175    async fn revalidate_contract(&mut self) -> Result<(), ConnectorError> {
176        let now = tokio::time::Instant::now();
177        if self.next_contract_check.is_some_and(|due| now < due) {
178            return Ok(());
179        }
180        let (Some(binding), Some(relation), Some(claim)) =
181            (&self.checkpoint_binding, &self.relation, &self.claim)
182        else {
183            return Ok(());
184        };
185        self.next_contract_check = Some(now + CONTRACT_CHECK_INTERVAL);
186        let checked = super::startup::revalidate(
187            &self.task_owner,
188            &self.config,
189            claim,
190            &self.incarnation,
191            binding,
192            relation,
193        )
194        .await;
195        match checked {
196            Ok(orphans) => {
197                if let Some(orphans) = orphans {
198                    self.metrics.set_orphaned_slots(orphans);
199                }
200                Ok(())
201            }
202            // An unreachable control plane is not drift; replication itself is still healthy.
203            Err(ConnectorError::ConnectionFailed(reason)) => {
204                tracing::warn!(%reason, "PostgreSQL CDC contract revalidation deferred");
205                Ok(())
206            }
207            Err(error) => Err(error),
208        }
209    }
210}
211
212#[async_trait]
213impl SourceConnector for PostgresCdcSource {
214    fn terminal_task_tracker(&self) -> Option<ConnectorTaskTracker> {
215        Some(self.task_tracker.clone())
216    }
217
218    fn recovery_identity_options(
219        &self,
220        config: &ConnectorConfig,
221    ) -> Result<Option<BTreeMap<String, String>>, ConnectorError> {
222        let parsed = parsed_config(&self.config, config)?;
223        Ok(Some(BTreeMap::from([
224            ("database".into(), parsed.database),
225            ("publication".into(), parsed.publication),
226            ("slot.name".into(), parsed.slot_name),
227            ("table".into(), parsed.table.to_string()),
228            ("output.mode".into(), parsed.output_mode.to_string()),
229            ("wire.protocol".into(), "pgoutput-v1-typed".into()),
230        ])))
231    }
232
233    async fn resolve_schema(
234        &mut self,
235        config: &ConnectorConfig,
236        explicit: Option<SchemaRef>,
237    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
238        let guard = self.task_owner.track().ok_or_else(|| {
239            ConnectorError::Internal("PostgreSQL CDC schema owner is retired".into())
240        })?;
241        super::super::schema_resolution::resolve(config, explicit, guard).await
242    }
243
244    async fn start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
245        if self.state != ConnectorState::Created {
246            return Err(ConnectorError::InvalidState {
247                expected: ConnectorState::Created.to_string(),
248                actual: self.state.to_string(),
249            });
250        }
251        let (config, position, _) = request.into_parts();
252        let parsed = parsed_config(&self.config, &config)?;
253        let declared = declared_schema(&config, None)?;
254        let primary_key = declared_primary_key(&config);
255        let committed = committed_relation(config.schema_binding(), &parsed)?;
256        let cursor = match position {
257            SourcePosition::Initial => None,
258            SourcePosition::Initialized { .. } => {
259                return Err(ConnectorError::ConfigurationError(
260                    "PostgreSQL CDC has no sealed topology startup contract".into(),
261                ));
262            }
263            SourcePosition::Resume {
264                attempt,
265                checkpoint,
266            } => Some(parse_cursor(
267                &checkpoint,
268                &parsed,
269                &format!("checkpoint {attempt:?}"),
270            )?),
271        };
272
273        let incarnation = incarnation();
274        let prepared = prepare(
275            &self.task_owner,
276            StartInputs {
277                data_ready: &self.data_ready,
278                config: &parsed,
279                declared: &declared,
280                primary_key: &primary_key,
281                committed_relation: committed.as_ref(),
282                incarnation: &incarnation,
283            },
284            cursor,
285        )
286        .await?;
287        // Publish the runtime only after all fallible network preparation succeeded.
288        self.open_rows = Some(RowBuilder::new(&prepared.layout));
289        self.layout = Some(prepared.layout);
290        self.relation = Some(prepared.relation);
291        self.checkpoint_binding = Some(prepared.binding);
292        self.copy_plan = prepared.copy_plan.map(Arc::new);
293        self.incarnation = incarnation;
294        self.config = parsed;
295        self.schema = declared;
296        if let Some(orphans) = prepared.orphans {
297            self.metrics.set_orphaned_slots(orphans);
298        }
299        match prepared.phase {
300            StartPhase::Claim { claim, committed } => self.begin_claim(claim, committed)?,
301            StartPhase::AwaitStream {
302                claim,
303                consistent_point,
304            } => {
305                self.claim = Some(claim);
306                self.consistent_point = Some(consistent_point);
307                self.polled_lsn = consistent_point;
308                self.phase = Phase::AwaitingStream { released: false };
309            }
310            StartPhase::Snapshot {
311                claim,
312                position,
313                runtime,
314            } => {
315                let plan = self.copy_plan.clone().ok_or_else(|| {
316                    ConnectorError::Internal("PostgreSQL CDC resumes a copy without a plan".into())
317                })?;
318                let (consistent_point, lsn) = (position.consistent_point, position.lsn);
319                let snapshot = ChunkedSnapshot::resume(
320                    plan,
321                    claim.application_name(&self.incarnation),
322                    position,
323                )?;
324                self.consistent_point = Some(consistent_point);
325                self.install_reader(lsn, runtime);
326                self.claim = Some(claim);
327                self.phase = Phase::Snapshot(Box::new(snapshot));
328            }
329            StartPhase::Stream {
330                claim,
331                consistent_point,
332                lsn,
333                runtime,
334            } => {
335                self.claim = Some(claim);
336                self.consistent_point = Some(consistent_point);
337                self.install_reader(lsn, runtime);
338                self.phase = Phase::Streaming;
339            }
340        }
341        self.state = ConnectorState::Running;
342        tracing::info!(
343            table = %self.config.table,
344            slot = self.slot_name()?,
345            output_mode = %self.config.output_mode,
346            snapshot_mode = ?self.config.snapshot_mode,
347            "PostgreSQL CDC source opened"
348        );
349        Ok(())
350    }
351
352    async fn poll_batch(
353        &mut self,
354        max_records: usize,
355    ) -> Result<Option<SourceBatch>, ConnectorError> {
356        if self.state != ConnectorState::Running {
357            // Repeat the cause as a terminal error: the engine retries transient errors, and a
358            // failed source never recovers in place.
359            let actual = match (self.state, self.failure.as_deref()) {
360                (ConnectorState::Failed, Some(cause)) => format!("Failed: {cause}"),
361                (state, _) => state.to_string(),
362            };
363            return Err(ConnectorError::InvalidState {
364                expected: "Running".to_string(),
365                actual,
366            });
367        }
368        match self.phase {
369            Phase::Claiming(_) => self.poll_claim().await,
370            Phase::AwaitingStream { released: false } => Ok(None),
371            Phase::AwaitingStream { released: true } => {
372                let start = self.consistent_point.ok_or_else(|| {
373                    ConnectorError::Internal(
374                        "PostgreSQL CDC awaits streaming without a consistent point".into(),
375                    )
376                })?;
377                self.begin_streaming(start).await?;
378                self.phase = Phase::Streaming;
379                Ok(None)
380            }
381            Phase::Snapshot(_) => self.poll_snapshot(max_records).await,
382            Phase::Streaming => self.poll_streaming(max_records),
383            Phase::Idle => Err(ConnectorError::Internal(
384                "PostgreSQL CDC is running without a read phase".into(),
385            )),
386        }
387    }
388
389    fn schema(&self) -> SchemaRef {
390        Arc::clone(&self.schema)
391    }
392
393    fn checkpoint(&self) -> SourceCheckpoint {
394        self.cursor()
395    }
396
397    async fn notify_epoch_committed(
398        &mut self,
399        epoch: u64,
400        checkpoint: &SourceCheckpoint,
401    ) -> Result<(), ConnectorError> {
402        // Empty cursors (another source's commit) and unstarted ones carry no position.
403        if checkpoint.get_metadata("phase").is_none() {
404            return Ok(());
405        }
406        let context = format!("committed epoch {epoch} checkpoint");
407        let committed = parse_cursor(checkpoint, &self.config, &context)?;
408        let (Some(active), Some(claim)) = (&self.checkpoint_binding, &self.claim) else {
409            return Err(ConnectorError::InvalidState {
410                expected: "running PostgreSQL CDC slot claim".into(),
411                actual: "the source has no slot claim".into(),
412            });
413        };
414        validate_live_binding(&committed.binding, active, &context)?;
415        if committed.claim != *claim {
416            // A claim replaced after its slot proved unusable can still commit; it emitted
417            // nothing, so it is not feedback.
418            return match committed.phase {
419                CursorPhase::Claimed => Ok(()),
420                CursorPhase::Snapshot { .. } | CursorPhase::Streaming { .. } => {
421                    Err(ConnectorError::ConfigurationError(format!(
422                        "committed PostgreSQL CDC epoch {epoch} names slot '{}', not this \
423                         source's slot '{}'",
424                        committed.claim.slot(),
425                        claim.slot()
426                    )))
427                }
428            };
429        }
430        let lsn = match (committed.phase, &mut self.phase) {
431            (CursorPhase::Claimed, Phase::Claiming(task)) => {
432                task.commit();
433                return Ok(());
434            }
435            (CursorPhase::Snapshot(_), Phase::Snapshot(snapshot)) if !snapshot.released() => {
436                snapshot.release();
437                return Ok(());
438            }
439            (CursorPhase::Streaming { .. }, Phase::AwaitingStream { released }) => {
440                *released = true;
441                return Ok(());
442            }
443            (CursorPhase::Streaming { lsn, .. }, _) => lsn,
444            // Released but the stream not opened yet: nothing past the consistent point was read.
445            (CursorPhase::Snapshot(_), _) if self.applied_lsn.is_none() => return Ok(()),
446            (CursorPhase::Snapshot(position), _) => position.lsn,
447            (CursorPhase::Claimed, _) => return Ok(()),
448        };
449        // Advance the slot only after the epoch is durably committed (manifest persisted and
450        // sinks committed), so PostgreSQL never reclaims WAL for rows still in the pipeline.
451        if lsn > self.polled_lsn {
452            return Err(ConnectorError::ConfigurationError(format!(
453                "committed PostgreSQL CDC epoch {epoch} LSN {lsn} is ahead of the source's polled LSN {}; refusing irreversible slot feedback",
454                self.polled_lsn
455            )));
456        }
457        if lsn < self.confirmed_flush_lsn {
458            return Ok(());
459        }
460        self.revalidate_contract().await?;
461        let applied = self
462            .applied_lsn
463            .as_ref()
464            .ok_or_else(|| ConnectorError::InvalidState {
465                expected: "running PostgreSQL CDC replication feedback".into(),
466                actual: "replication feedback handle is missing".into(),
467            })?;
468        applied.update(pgwire_replication::Lsn::from_u64(lsn.as_u64()));
469        self.confirmed_flush_lsn = lsn;
470        self.metrics.set_confirmed_flush_lsn(lsn.as_u64());
471        Ok(())
472    }
473
474    fn contract(&self, config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
475        let parsed = parsed_config(&self.config, config)?;
476        // Slot WAL is reclaimed only as durable commits advance the confirmed-flush LSN, so the
477        // source is commit-coupled. Exact delivery is not certified: feedback and sink commits
478        // are not one atomic protocol.
479        let input_mode = match parsed.output_mode {
480            OutputMode::Upsert => SourceInputMode::KeyedUpsert,
481            OutputMode::Changelog => SourceInputMode::FullChangelog,
482        };
483        Ok(SourceContract::new(
484            SourceConsistency::CommitCoupled,
485            SourceTopology::Singleton,
486            input_mode,
487        )
488        .with_row_positions(SourceRowPositionCapability::OrderedDeterministic))
489    }
490
491    fn data_ready_notify(&self) -> Option<Arc<Notify>> {
492        Some(Arc::clone(&self.data_ready))
493    }
494
495    async fn close(&mut self) -> Result<(), ConnectorError> {
496        // Keep both fields installed while awaiting. If this close future is
497        // cancelled, the same instance still owns the reader and can retry.
498        if let Some(tx) = self.reader_shutdown.as_ref() {
499            tx.send_replace(true);
500        }
501        let detach_reader = if let Some(handle) = self.reader_handle.as_mut() {
502            tokio::time::timeout(std::time::Duration::from_secs(5), &mut *handle)
503                .await
504                .is_err()
505        } else {
506            false
507        };
508        if detach_reader {
509            tracing::warn!(
510                "PostgreSQL CDC reader did not stop before the close deadline; its tracked reaper retains shutdown ownership"
511            );
512            if let Some(handle) = self.reader_handle.take() {
513                reap_postgres_reader(handle, &self.task_owner);
514            }
515        }
516        match std::mem::replace(&mut self.phase, Phase::Idle) {
517            Phase::Claiming(task) => task.cancel().await,
518            Phase::Snapshot(snapshot) => drop(snapshot),
519            Phase::Idle | Phase::AwaitingStream { .. } | Phase::Streaming => {}
520        }
521        self.reader_handle = None;
522        self.reader_shutdown = None;
523        self.wal_rx = None;
524        self.applied_lsn = None;
525        self.pending_payloads.clear();
526        self.wal_byte_budget = None;
527        self.wal_terminal_error = None;
528        self.state = ConnectorState::Closed;
529        self.committed.clear();
530        self.committed_bytes = 0;
531        self.open_transaction = None;
532        self.open_mutations.clear();
533        if let (Some(rows), Some(layout)) = (self.open_rows.as_mut(), self.layout.as_ref()) {
534            drop(rows.finish(layout));
535        }
536        #[cfg(test)]
537        self.pending_messages.clear();
538        Ok(())
539    }
540}
laminar_connectors/postgres/cdc/source/lifecycle.rs

laminar_connectors/postgres/cdc/source/
lifecycle.rs
