laminar_connectors/postgres/cdc/source/
lifecycle.rs

1//! Source startup, polling, durable feedback, and shutdown ownership.
2
3#[cfg(not(test))]
4use super::startup::{prepare_reader_runtime, PreparedReaderRuntime};
5use super::{
6    async_trait, reap_postgres_reader, validate_checkpoint_identity, validate_live_binding,
7    write_checkpoint_binding, Arc, BTreeMap, ConnectorConfig, ConnectorError, ConnectorState,
8    ConnectorTaskTracker, Lsn, Notify, PostgresCdcConfig, PostgresCdcSource,
9    PostgresCheckpointBinding, SchemaRef, SourceBatch, SourceCheckpoint, SourceConnector,
10    SourceContract, SourcePosition, SourceStart, INITIAL_BOOTSTRAP_NOT_ADMITTED,
11};
12
13struct PreparedSourceStart {
14    config: PostgresCdcConfig,
15    start_lsn: Lsn,
16    checkpoint_binding: PostgresCheckpointBinding,
17}
18
19fn prepare_source_start(
20    current_config: &PostgresCdcConfig,
21    config: &ConnectorConfig,
22    position: SourcePosition,
23) -> Result<PreparedSourceStart, ConnectorError> {
24    let mut config = if config.properties().is_empty() {
25        current_config.clone()
26    } else {
27        PostgresCdcConfig::from_config(config)?
28    };
29    config.normalize_table_filters();
30    config.validate()?;
31
32    let SourcePosition::Resume {
33        attempt,
34        checkpoint,
35    } = position
36    else {
37        return Err(ConnectorError::ConfigurationError(
38            INITIAL_BOOTSTRAP_NOT_ADMITTED.into(),
39        ));
40    };
41    let context = format!("checkpoint {attempt:?}");
42    let checkpoint_binding = validate_checkpoint_identity(&checkpoint, &config, &context)?;
43    let lsn_str = checkpoint.get_offset("lsn").ok_or_else(|| {
44        ConnectorError::ConfigurationError(format!(
45            "PostgreSQL CDC checkpoint {attempt:?} is missing required 'lsn' offset"
46        ))
47    })?;
48    let start_lsn = lsn_str.parse::<Lsn>().map_err(|error| {
49        ConnectorError::ConfigurationError(format!(
50            "invalid LSN '{lsn_str}' in PostgreSQL CDC checkpoint {attempt:?}: {error}"
51        ))
52    })?;
53    Ok(PreparedSourceStart {
54        config,
55        start_lsn,
56        checkpoint_binding,
57    })
58}
59
60#[cfg(not(test))]
61fn install_reader_runtime(source: &mut PostgresCdcSource, runtime: PreparedReaderRuntime) {
62    source.wal_rx = Some(runtime.wal_rx);
63    source.wal_byte_budget = Some(runtime.wal_byte_budget);
64    source.wal_terminal_error = Some(runtime.terminal_error);
65    source.reader_handle = Some(runtime.reader_handle);
66    source.reader_shutdown = Some(runtime.shutdown_tx);
67    source.confirmed_lsn_tx = Some(runtime.confirmed_lsn_tx);
68}
69
70#[async_trait]
71impl SourceConnector for PostgresCdcSource {
72    fn terminal_task_tracker(&self) -> Option<ConnectorTaskTracker> {
73        Some(self.task_tracker.clone())
74    }
75
76    fn recovery_identity_options(
77        &self,
78        config: &ConnectorConfig,
79    ) -> Result<Option<BTreeMap<String, String>>, ConnectorError> {
80        let mut parsed = if config.properties().is_empty() {
81            self.config.clone()
82        } else {
83            PostgresCdcConfig::from_config(config)?
84        };
85        parsed.normalize_table_filters();
86        parsed.validate()?;
87
88        Ok(Some(BTreeMap::from([
89            ("database".into(), parsed.database),
90            ("publication".into(), parsed.publication),
91            ("slot.name".into(), parsed.slot_name),
92            ("table.exclude".into(), parsed.table_exclude.join(",")),
93            ("table.include".into(), parsed.table_include.join(",")),
94            ("wire.protocol".into(), "pgoutput-v1".into()),
95        ])))
96    }
97
98    async fn start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
99        if self.state != ConnectorState::Created {
100            return Err(ConnectorError::InvalidState {
101                expected: ConnectorState::Created.to_string(),
102                actual: self.state.to_string(),
103            });
104        }
105        let (config, position, _) = request.into_parts();
106        let prepared = prepare_source_start(&self.config, &config, position)?;
107
108        #[cfg(not(test))]
109        {
110            let runtime = prepare_reader_runtime(
111                self,
112                &prepared.config,
113                &prepared.checkpoint_binding,
114                prepared.start_lsn,
115            )
116            .await?;
117            install_reader_runtime(self, runtime);
118        }
119
120        // Publish the new runtime only after all fallible startup work has
121        // succeeded. A failed start remains a clean Created connector.
122        self.config = prepared.config;
123        self.confirmed_flush_lsn = prepared.start_lsn;
124        self.write_lsn = prepared.start_lsn;
125        self.polled_lsn = prepared.start_lsn;
126        self.checkpoint_binding = Some(prepared.checkpoint_binding);
127        self.metrics
128            .set_confirmed_flush_lsn(prepared.start_lsn.as_u64());
129        self.state = ConnectorState::Running;
130        Ok(())
131    }
132
133    async fn poll_batch(
134        &mut self,
135        max_records: usize,
136    ) -> Result<Option<SourceBatch>, ConnectorError> {
137        if self.state != ConnectorState::Running {
138            return Err(ConnectorError::InvalidState {
139                expected: "Running".to_string(),
140                actual: self.state.to_string(),
141            });
142        }
143
144        // Backpressure: stop draining raw WAL before the decoded-stage hard limit. The raw byte
145        // budget then propagates pressure to the replication reader and PostgreSQL.
146        {
147            self.fail_on_terminal_wal_error()?;
148            let high_watermark = self.config.decoded_high_watermark_bytes();
149            let decoded_retained_bytes = self.decoded_retained_bytes().inspect_err(|_error| {
150                self.state = ConnectorState::Failed;
151            })?;
152            let mut reader_closed = false;
153            let must_finish_transaction = self.current_txn.is_some();
154            let payload_budget = max_records.max(1);
155            let drain_reader = must_finish_transaction || decoded_retained_bytes < high_watermark;
156            if !drain_reader && self.pending_payloads.is_empty() {
157                tracing::debug!(
158                    retained_bytes = decoded_retained_bytes,
159                    high_watermark,
160                    "CDC backpressure active — pausing WAL reader drain"
161                );
162            }
163
164            let mut processed_payloads = 0_usize;
165            while processed_payloads < payload_budget {
166                let payload = if let Some(payload) = self.pending_payloads.pop_front() {
167                    Some(payload)
168                } else if drain_reader {
169                    match self.wal_rx.as_mut().map(|receiver| receiver.try_recv()) {
170                        Some(Ok(payload)) => Some(payload),
171                        Some(Err(crossfire::TryRecvError::Empty)) | None => None,
172                        Some(Err(crossfire::TryRecvError::Disconnected)) => {
173                            reader_closed = true;
174                            None
175                        }
176                    }
177                } else {
178                    None
179                };
180                let Some(payload) = payload else {
181                    break;
182                };
183                if let Err(e) = self.process_owned_wal_payload(payload) {
184                    self.state = ConnectorState::Failed;
185                    return Err(e);
186                }
187                processed_payloads = processed_payloads.checked_add(1).ok_or_else(|| {
188                    self.state = ConnectorState::Failed;
189                    ConnectorError::Internal(
190                        "PostgreSQL CDC poll payload-count accounting overflow".into(),
191                    )
192                })?;
193            }
194
195            // Notify ourselves only when a bounded drain demonstrably left work queued.
196            // Retaining one item avoids both a lost coalesced notification and an
197            // open-transaction busy loop while the server is genuinely idle.
198            let reached_payload_budget = processed_payloads == payload_budget;
199            let may_drain_more = if reached_payload_budget && self.current_txn.is_none() {
200                self.decoded_retained_bytes().inspect_err(|_error| {
201                    self.state = ConnectorState::Failed;
202                })? < high_watermark
203            } else {
204                reached_payload_budget
205            };
206            if may_drain_more && !reader_closed {
207                if let Some(ref mut rx) = self.wal_rx {
208                    match rx.try_recv() {
209                        Ok(payload) => self.pending_payloads.push_back(payload),
210                        Err(crossfire::TryRecvError::Empty) => {}
211                        Err(crossfire::TryRecvError::Disconnected) => reader_closed = true,
212                    }
213                }
214            }
215            if !self.pending_payloads.is_empty() {
216                self.data_ready.notify_one();
217            }
218            self.fail_on_terminal_wal_error()?;
219            if reader_closed && self.committed_transactions.is_empty() {
220                self.state = ConnectorState::Failed;
221                return Err(ConnectorError::ReadError(
222                    "WAL reader task terminated unexpectedly — replication stream lost".to_string(),
223                ));
224            }
225        }
226
227        #[cfg(test)]
228        self.process_pending_messages()?;
229
230        // Drain buffered events into a RecordBatch.
231        // Configured Arrow-column extractors derive event-time watermarks from `_ts_ms`.
232        let result = match self.drain_events(max_records)? {
233            Some(batch) => {
234                self.metrics
235                    .set_confirmed_flush_lsn(self.confirmed_flush_lsn.as_u64());
236                self.metrics
237                    .set_replication_lag_bytes(self.replication_lag_bytes());
238
239                Ok(Some(SourceBatch::new(batch)))
240            }
241            None => Ok(None),
242        };
243        let emitted_batch = matches!(&result, Ok(Some(_)));
244        if max_records > 0 && (emitted_batch || !self.committed_transactions.is_empty()) {
245            self.data_ready.notify_one();
246        }
247        result
248    }
249
250    fn schema(&self) -> SchemaRef {
251        Arc::clone(&self.schema)
252    }
253
254    fn checkpoint(&self) -> SourceCheckpoint {
255        let mut cp = SourceCheckpoint::new();
256        // polled_lsn = latest position drained into a batch — the resumable point recorded in the
257        // manifest. The PG slot is NOT advanced here: doing so per poll lets PG reclaim WAL for
258        // data that is only in-pipeline, so a crash loses an LSN range recovery still needs.
259        // Slot feedback is deferred to notify_epoch_committed (durable-commit only).
260        cp.set_offset("lsn", self.polled_lsn.to_string());
261        cp.set_metadata("slot_name", &self.config.slot_name);
262        cp.set_metadata("publication", &self.config.publication);
263        cp.set_metadata("database", &self.config.database);
264        if let Some(binding) = &self.checkpoint_binding {
265            write_checkpoint_binding(&mut cp, binding);
266        }
267        cp
268    }
269
270    async fn notify_epoch_committed(
271        &mut self,
272        epoch: u64,
273        checkpoint: &SourceCheckpoint,
274    ) -> Result<(), ConnectorError> {
275        // Advance the PG replication slot only after the epoch is durably committed (manifest
276        // persisted + sinks committed), so PG never reclaims WAL for data still in-pipeline.
277        // The checkpoint carries the exact LSN persisted for this epoch; a timer-driven empty
278        // checkpoint has no "lsn" offset and is a no-op.
279        let Some(lsn_str) = checkpoint.get_offset("lsn") else {
280            return Ok(());
281        };
282        let lsn = lsn_str.parse::<Lsn>().map_err(|error| {
283            ConnectorError::ConfigurationError(format!(
284                "committed PostgreSQL CDC epoch {epoch} contains invalid LSN '{lsn_str}': {error}"
285            ))
286        })?;
287        let context = format!("committed epoch {epoch} checkpoint");
288        let committed_binding = validate_checkpoint_identity(checkpoint, &self.config, &context)?;
289        let active_binding =
290            self.checkpoint_binding
291                .as_ref()
292                .ok_or_else(|| ConnectorError::InvalidState {
293                    expected: "running PostgreSQL CDC checkpoint binding".into(),
294                    actual: "checkpoint binding is missing".into(),
295                })?;
296        validate_live_binding(&committed_binding, active_binding, &context)?;
297        if lsn.as_u64() > self.polled_lsn.as_u64() {
298            return Err(ConnectorError::ConfigurationError(format!(
299                "committed PostgreSQL CDC epoch {epoch} LSN {lsn} is ahead of the source's polled LSN {}; refusing irreversible slot feedback",
300                self.polled_lsn
301            )));
302        }
303        // A strictly stale notification is already satisfied and must never regress either cursor.
304        // An equal notification is handed off again: that is idempotent and repairs feedback after
305        // a reader restart whose local cursor was restored before its channel was created.
306        if lsn.as_u64() < self.confirmed_flush_lsn.as_u64() {
307            return Ok(());
308        }
309
310        let tx = self
311            .confirmed_lsn_tx
312            .as_ref()
313            .ok_or_else(|| ConnectorError::InvalidState {
314                expected: "running PostgreSQL CDC confirmed-LSN feedback channel".into(),
315                actual: "feedback channel is missing".into(),
316            })?;
317        tx.send(lsn.as_u64()).map_err(|_| {
318            ConnectorError::ConnectionFailed(
319                "PostgreSQL CDC confirmed-LSN feedback channel is closed".into(),
320            )
321        })?;
322        // The local cursor is authoritative only after the reader accepted the handoff.
323        self.confirmed_flush_lsn = lsn;
324        self.metrics.set_confirmed_flush_lsn(lsn.as_u64());
325        Ok(())
326    }
327
328    fn contract(&self, config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
329        // The replication slot's WAL is reclaimed only as the confirmed-flush LSN advances, which
330        // happens on durable commit. Without checkpointing the slot never advances and the source
331        // database's WAL fills without bound, so this source is commit-coupled.
332        if config.properties().is_empty() {
333            self.config.validate()?;
334        } else {
335            PostgresCdcConfig::from_config(config)?.validate()?;
336        }
337        Err(ConnectorError::ConfigurationError(
338            "PostgreSQL CDC emits a raw JSON change envelope; canonical primary-keyed row/delete records are required"
339                .into(),
340        ))
341    }
342
343    fn data_ready_notify(&self) -> Option<Arc<Notify>> {
344        Some(Arc::clone(&self.data_ready))
345    }
346
347    async fn close(&mut self) -> Result<(), ConnectorError> {
348        // Keep both fields installed while awaiting. If this close future is
349        // cancelled, the same instance still owns the reader and can retry.
350        if let Some(tx) = self.reader_shutdown.as_ref() {
351            tx.send_replace(true);
352        }
353        let detach_reader = if let Some(handle) = self.reader_handle.as_mut() {
354            tokio::time::timeout(std::time::Duration::from_secs(5), &mut *handle)
355                .await
356                .is_err()
357        } else {
358            false
359        };
360        if detach_reader {
361            tracing::warn!(
362                "PostgreSQL CDC reader did not stop before the close deadline; its tracked reaper retains shutdown ownership"
363            );
364            let handle = self
365                .reader_handle
366                .take()
367                .expect("reader handle was present while awaiting it");
368            reap_postgres_reader(handle, &self.task_owner);
369        }
370        self.reader_handle = None;
371        self.reader_shutdown = None;
372        self.wal_rx = None;
373        self.confirmed_lsn_tx = None;
374        self.pending_payloads.clear();
375        self.wal_byte_budget = None;
376        self.wal_terminal_error = None;
377
378        self.state = ConnectorState::Closed;
379        self.committed_transactions.clear();
380        self.current_txn = None;
381        self.relation_cache.clear();
382        self.buffered_event_count = 0;
383        self.buffered_event_bytes = 0;
384        #[cfg(test)]
385        self.pending_messages.clear();
386        Ok(())
387    }
388}
laminar_connectors/postgres/cdc/source/lifecycle.rs

laminar_connectors/postgres/cdc/source/
lifecycle.rs
