laminar_connectors/kafka/source/
lifecycle.rs

1//! Source contract, drain lifecycle, polling, checkpointing, and shutdown.
2
3use super::{
4    async_trait, info, join_background_task, reap_last_arc_off_runtime, warn, Arc, CommitMode,
5    ConnectorConfig, ConnectorError, ConnectorState, ConnectorTaskTracker, Consumer, Format,
6    KafkaReaderDrainCommand, KafkaSource, KafkaSourceConfig, KafkaSourceDrain, Notify,
7    OffsetTracker, Ordering, SchemaRef, SourceBatch, SourceCheckpoint, SourceConnector,
8    SourceConsistency, SourceContract, SourceDrainRequest, SourceDrainResolution, SourceInputMode,
9    SourceRowPositionCapability, SourceStart, SourceTopology, KAFKA_BACKGROUND_CLOSE_BUDGET,
10};
11
12#[async_trait]
13impl SourceConnector for KafkaSource {
14    fn supports_initialized_start(&self) -> bool {
15        true
16    }
17
18    async fn resolve_initial_position(
19        &mut self,
20        config: &ConnectorConfig,
21    ) -> Result<SourceCheckpoint, ConnectorError> {
22        self.resolve_initial_position_inner(config).await
23    }
24
25    async fn validate_initial_position(
26        &mut self,
27        config: &ConnectorConfig,
28        checkpoint: &SourceCheckpoint,
29    ) -> Result<(), ConnectorError> {
30        self.inspect_initial_position_inner(config, Some(checkpoint))
31            .await
32            .map(|_| ())
33    }
34
35    fn terminal_task_tracker(&self) -> Option<ConnectorTaskTracker> {
36        Some(self.task_tracker.clone())
37    }
38
39    fn recency(
40        &self,
41        config: &ConnectorConfig,
42    ) -> Result<Option<Arc<dyn crate::connector::SourceRecency>>, ConnectorError> {
43        Ok(super::recency::kafka_recency(
44            &KafkaSourceConfig::from_config(config)?,
45        ))
46    }
47
48    fn contract(&self, config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
49        let format = if config.properties().is_empty() {
50            self.config.validate()?;
51            self.config.format
52        } else {
53            KafkaSourceConfig::from_config(config)?.format
54        };
55        let input_mode = if format == Format::Debezium {
56            SourceInputMode::KeyedUpsert
57        } else {
58            SourceInputMode::AppendOnly
59        };
60        Ok(SourceContract::new(
61            SourceConsistency::Replayable,
62            SourceTopology::Splittable,
63            input_mode,
64        )
65        .with_row_positions(SourceRowPositionCapability::OrderedDeterministic)
66        .with_exact_delivery_certification())
67    }
68
69    fn set_vnode_assignment(
70        &mut self,
71        source_identity: &str,
72        registry: Arc<laminar_core::state::VnodeRegistry>,
73        self_id: laminar_core::state::NodeId,
74    ) -> Result<(), ConnectorError> {
75        if source_identity.is_empty() {
76            return Err(ConnectorError::ConfigurationError(
77                "Kafka vnode assignment requires a non-empty canonical source identity".into(),
78            ));
79        }
80        if self_id.is_unassigned() {
81            return Err(ConnectorError::ConfigurationError(
82                "Kafka vnode assignment requires a nonzero node identity".into(),
83            ));
84        }
85        info!(
86            source = source_identity,
87            self_id = self_id.0,
88            vnode_count = registry.vnode_count(),
89            "Kafka source: engine-controlled partition-to-vnode assignment enabled"
90        );
91        self.source_name = Arc::from(source_identity);
92        self.vnode_assignment = Some((registry, self_id));
93        Ok(())
94    }
95
96    fn begin_drain(
97        &mut self,
98        request: &SourceDrainRequest,
99        deadline: tokio::time::Instant,
100    ) -> Result<(), ConnectorError> {
101        self.check_reader_health("starting a global source drain")?;
102        if let Some(active) = self.source_drain.as_ref() {
103            if active.request != *request {
104                return Err(ConnectorError::InvalidState {
105                    expected: format!("active Kafka drain {:?}", active.request.round),
106                    actual: format!("conflicting Kafka drain {:?}", request.round),
107                });
108            }
109            // Retain the first preparation deadline. A caller retry carries a wait budget,
110            // not authority to extend or shorten work already in progress.
111            return Ok(());
112        }
113        if tokio::time::Instant::now() >= deadline {
114            return Err(ConnectorError::Internal(
115                "Kafka source drain began after its engine deadline".into(),
116            ));
117        }
118        let Some((registry, _)) = self.vnode_assignment.as_ref() else {
119            return Err(ConnectorError::InvalidState {
120                expected: "Kafka cluster assignment installed before source drain".into(),
121                actual: "embedded/single Kafka source".into(),
122            });
123        };
124        if registry.assignment_version() != request.round.predecessor_version {
125            return Err(ConnectorError::InvalidState {
126                expected: format!(
127                    "Kafka predecessor assignment {}",
128                    request.round.predecessor_version
129                ),
130                actual: registry.assignment_version().to_string(),
131            });
132        }
133        let reconciled = self.reconciled_assignment_version.load(Ordering::Acquire);
134        if reconciled != request.round.predecessor_version {
135            return Err(ConnectorError::InvalidState {
136                expected: format!(
137                    "reconciled Kafka predecessor assignment {}",
138                    request.round.predecessor_version
139                ),
140                actual: reconciled.to_string(),
141            });
142        }
143        self.ensure_reader_started();
144        let tx = self.reader_drain_tx.as_ref().ok_or_else(|| {
145            ConnectorError::Internal("Kafka cluster reader has no drain control channel".into())
146        })?;
147        tx.send(KafkaReaderDrainCommand::Begin {
148            request: request.clone(),
149            deadline,
150        })
151        .map_err(|_| ConnectorError::Internal("Kafka reader drain channel closed".into()))?;
152        self.source_drain = Some(KafkaSourceDrain {
153            request: request.clone(),
154            prepare_deadline: deadline,
155            boundary: None,
156            cut: None,
157            pending_resolution: None,
158        });
159        self.data_ready.notify_one();
160        Ok(())
161    }
162
163    fn poll_drain_ready(
164        &mut self,
165        round: laminar_core::checkpoint::AssignmentDrainId,
166    ) -> Result<bool, ConnectorError> {
167        self.check_reader_health("capturing a global source drain cut")?;
168        let Some(active) = self.source_drain.as_ref() else {
169            return Err(ConnectorError::InvalidState {
170                expected: format!("active Kafka drain {round:?}"),
171                actual: "no Kafka drain".into(),
172            });
173        };
174        if active.request.round != round {
175            return Err(ConnectorError::InvalidState {
176                expected: format!("active Kafka drain {:?}", active.request.round),
177                actual: format!("cut requested for {round:?}"),
178            });
179        }
180        if active.cut.is_some() {
181            return Ok(true);
182        }
183        if tokio::time::Instant::now() >= active.prepare_deadline {
184            return Err(ConnectorError::Internal(
185                "Kafka drain deadline expired before cursor capture".into(),
186            ));
187        }
188        let Some(boundary) = active.boundary.clone() else {
189            return Ok(false);
190        };
191        let cut = self.capture_drain_positions(&boundary.inputs, Some(active.prepare_deadline))?;
192        let active = self.source_drain.as_mut().expect("checked above");
193        active.cut = Some(cut);
194        Ok(true)
195    }
196
197    async fn finish_drain(
198        &mut self,
199        resolution: SourceDrainResolution,
200        deadline: tokio::time::Instant,
201    ) -> Result<(), ConnectorError> {
202        self.finish_drain_inner(resolution, deadline).await
203    }
204
205    async fn start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
206        if self.state != ConnectorState::Created {
207            return Err(ConnectorError::InvalidState {
208                expected: ConnectorState::Created.to_string(),
209                actual: self.state.to_string(),
210            });
211        }
212        let result = self.start_with_contract(request).await;
213        if result.is_err() && self.state == ConnectorState::Initializing {
214            self.fail_startup();
215        }
216        result
217    }
218    async fn resolve_schema(
219        &mut self,
220        config: &crate::config::ConnectorConfig,
221        explicit: Option<arrow_schema::SchemaRef>,
222    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
223        super::super::schema_resolution::resolve_source(config, explicit).await
224    }
225
226    async fn discover_schema(
227        &mut self,
228        properties: &std::collections::HashMap<String, String>,
229    ) -> Result<(), ConnectorError> {
230        let config = crate::config::ConnectorConfig::with_properties("kafka", properties.clone());
231        let binding = super::super::schema_resolution::resolve_source(&config, None).await?;
232        let parsed = KafkaSourceConfig::from_config(&config)?;
233        self.schema = super::super::schema_resolution::payload_schema(&binding, &parsed.metadata)?;
234        Ok(())
235    }
236
237    async fn poll_batch(
238        &mut self,
239        max_records: usize,
240    ) -> Result<Option<SourceBatch>, ConnectorError> {
241        self.poll_batch_inner(max_records).await
242    }
243
244    fn schema(&self) -> SchemaRef {
245        self.config.metadata.output_schema(&self.schema)
246    }
247
248    fn checkpoint_ready(&self) -> Result<bool, ConnectorError> {
249        self.check_reader_health("reconciling source ownership")?;
250        Ok(self.vnode_assignment.as_ref().is_none_or(|(registry, _)| {
251            let version = registry.assignment_version();
252            version != 0 && self.reconciled_assignment_version.load(Ordering::Acquire) == version
253        }))
254    }
255
256    fn drive_control_plane(&mut self) {
257        self.ensure_reader_started();
258    }
259
260    fn checkpoint(&self) -> SourceCheckpoint {
261        self.try_capture_checkpoint()
262            .ok()
263            .flatten()
264            .unwrap_or_default()
265    }
266
267    fn try_checkpoint(&self) -> Result<Option<SourceCheckpoint>, ConnectorError> {
268        self.validate_active_drain_cursor()?;
269        self.try_capture_checkpoint()
270    }
271
272    fn data_ready_notify(&self) -> Option<Arc<Notify>> {
273        Some(Arc::clone(&self.data_ready))
274    }
275
276    async fn notify_epoch_committed(
277        &mut self,
278        epoch: u64,
279        checkpoint: &SourceCheckpoint,
280    ) -> Result<(), ConnectorError> {
281        if !self.config.broker_commit_on_checkpoint || checkpoint.is_empty() {
282            return Ok(());
283        }
284        let tpl = OffsetTracker::try_from_checkpoint(checkpoint)?.to_topic_partition_list();
285        if tpl.count() == 0 {
286            return Ok(());
287        }
288        let Some(consumer) = self.consumer.as_ref() else {
289            // Engine recovery never uses broker-stored offsets. A missing consumer cannot
290            // invalidate the already-durable checkpoint, so report the observability failure
291            // without turning a committed epoch into a pipeline restart loop.
292            self.metrics.commit_failures.inc();
293            warn!(
294                epoch,
295                "Kafka progress commit skipped because the consumer is absent"
296            );
297            return Ok(());
298        };
299
300        // The engine checkpoint is the recovery authority; Kafka's group offset is an
301        // observability cursor only. Enqueue it without blocking the checkpoint/recovery path.
302        // The consumer context records the eventual broker acknowledgement or rejection.
303        if let Err(error) = consumer.commit(&tpl, CommitMode::Async) {
304            self.metrics.commit_failures.inc();
305            warn!(epoch, %error, "Kafka progress commit was not accepted for enqueue");
306        }
307        Ok(())
308    }
309
310    async fn close(&mut self) -> Result<(), ConnectorError> {
311        info!("closing Kafka source connector");
312
313        // Stop intake and give the reader plus final consumer reaper one shared cleanup budget
314        // below the coordinator's outer source-shutdown deadline.
315        if let Some(tx) = self.reader_shutdown.take() {
316            let _ = tx.send(true);
317        }
318        if let Some(progress) = &mut self.progress {
319            progress.stop();
320        }
321        // Wake assignment and poll work before joining. Any advisory async commit cleanup remains
322        // librdkafka-owned and is not allowed to extend the engine's source-shutdown deadline.
323        if let Some(ref consumer) = self.consumer {
324            consumer.unsubscribe();
325        }
326        let deadline = tokio::time::Instant::now() + KAFKA_BACKGROUND_CLOSE_BUDGET;
327        join_background_task(&mut self.reader_handle, deadline, "reader").await;
328        if let Some(mut progress) = self.progress.take() {
329            progress.close(deadline).await;
330        }
331        self.msg_rx = None;
332        self.reader_drain_tx = None;
333        self.source_drain = None;
334        self.channel_len.store(0, Ordering::Release);
335        if let Some(consumer) = self.consumer.take() {
336            reap_last_arc_off_runtime(&self.blocking_tasks, consumer, deadline, "consumer").await;
337        }
338        if !self.blocking_tasks.join_until(deadline).await {
339            self.blocking_tasks.ensure_reaper();
340        }
341        self.state = ConnectorState::Closed;
342        info!("Kafka source connector closed");
343        Ok(())
344    }
345}
laminar_connectors/kafka/source/lifecycle.rs

laminar_connectors/kafka/source/
lifecycle.rs
