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

laminar_connectors/kafka/source/
lifecycle.rs
