laminar_connectors/kafka/source/
mod.rs

1//! Kafka source connector: consumes topics via rdkafka's `StreamConsumer`,
2//! deserializes with pluggable formats, and yields Arrow `RecordBatch`es.
3
4use arrow_array::builder::BinaryBuilder;
5use arrow_array::{Array, RecordBatch, RecordBatchOptions, StringArray, UInt32Array};
6use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit};
7use async_trait::async_trait;
8use rdkafka::consumer::{CommitMode, Consumer, StreamConsumer};
9use rdkafka::error::KafkaError;
10use rdkafka::message::Message;
11use rdkafka::ClientConfig;
12use rdkafka::TopicPartitionList;
13use std::collections::BTreeSet;
14use std::num::NonZeroU64;
15use std::sync::atomic::{AtomicBool, AtomicU64, AtomicU8, AtomicUsize, Ordering};
16use std::sync::{Arc, Mutex};
17use tokio::sync::{Mutex as AsyncMutex, Notify};
18use tracing::{debug, info, warn};
19
20use crate::checkpoint::{SourceCheckpoint, SourceCheckpointDelta};
21use crate::config::{ConnectorConfig, ConnectorState};
22use crate::connector::{
23    ConnectorTaskGuard, ConnectorTaskOwner, ConnectorTaskTracker, DeliveryGuarantee, SourceBatch,
24    SourceConnector, SourceConsistency, SourceContract, SourceDrainOutcome, SourceDrainRequest,
25    SourceDrainResolution, SourceInputMode, SourceMutation, SourcePosition,
26    SourceRowPositionCapability, SourceRowPositions, SourceStart, SourceTopology,
27};
28use crate::error::{ConnectorError, SerdeError};
29use crate::serde::{self, Format, RecordDeserializer};
30
31use super::avro::AvroDeserializer;
32use super::config::{
33    resolve_value_subject, KafkaSourceConfig, OffsetReset, SchemaEvolutionStrategy, StartupMode,
34    TopicSubscription,
35};
36use super::metadata_error::{fetch_error, invalid_response, topic_error};
37use super::metrics::KafkaSourceMetrics;
38use super::offsets::{OffsetTracker, KAFKA_PARTITION_BASELINE_PREFIX};
39use super::rebalance::RebalanceState;
40use super::schema_registry::SchemaRegistryClient;
41
42use super::rebalance::LaminarConsumerContext;
43use crate::schema::evolution::SchemaEvolution;
44use crate::schema::traits::{CompatibilityMode, EvolutionVerdict};
45
46mod background;
47mod checkpoint;
48mod debezium;
49mod decoding;
50mod drain;
51mod lifecycle;
52mod metadata;
53mod polling;
54mod reader;
55mod startup;
56
57use background::{
58    ensure_background_task_reaper, join_background_task, reap_last_arc_off_runtime,
59    KafkaBlockingTaskError, KafkaBlockingTasks,
60};
61use checkpoint::{
62    acquired_numeric_position, assignment_seek_tpl, build_vnode_assignment_tpl,
63    consumer_creation_error, decode_partition_baselines, deterministic_initial_offset,
64    kafka_input_channels, kafka_output_schema, kafka_reader_error_is_transient,
65    kafka_row_positions, retire_accepted_rotation_baselines, rotation_baselines_len,
66    rotation_partition_baseline, startup_default_offset, tpl_of, update_rotation_baselines,
67    validate_kafka_output_schema, validate_partition_baselines, validate_positions_not_expired,
68    validate_resume_input_channels, vnode_payload_is_current, NormalizedDebeziumBatch,
69};
70use debezium::normalize_kafka_debezium_batch;
71use drain::{
72    cached_partition_vnode, kafka_bootstrap_is_unassigned, kafka_drain_partitions,
73    kafka_drain_target_ready, kafka_owned_partition_sets, kafka_partition_routes,
74    kafka_partition_set, lock_or_recover, publish_reader_fault, resolve_kafka_reader_drain,
75    terminalize_guaranteed_poll_error, try_capture_at_assignment_fence, validate_kafka_assignment,
76    validate_kafka_partition_results, KafkaAssignmentPublication, KafkaDrainBoundary,
77    KafkaDrainPartition, KafkaDrainPosition, KafkaPartitionBaselines, KafkaPartitionRoutes,
78    KafkaPartitionSet, KafkaPayload, KafkaReaderDrain, KafkaReaderDrainCommand, KafkaReaderItem,
79    KafkaReaderRx, KafkaRotationBaselines, KafkaSourceDrain, KafkaStartPlan,
80    KAFKA_BACKGROUND_CLOSE_BUDGET, KAFKA_POSITION_LOOKUP_BUDGET, KAFKA_POSITION_LOOKUP_CONCURRENCY,
81};
82use metadata::{
83    fetch_explicit_topic_metadata, fetch_partition_low_watermarks, fetch_partition_watermarks,
84    resolve_timestamp_offsets,
85};
86
87/// Kafka source connector that consumes messages and produces Arrow batches.
88///
89/// Operates in Ring 1 (background) and pushes deserialized `RecordBatch`
90/// data to Ring 0 via the streaming `Source<T>` API.
91///
92/// # Lifecycle
93///
94/// 1. Create with [`KafkaSource::new`] or [`KafkaSource::with_schema_registry`]
95/// 2. Call `start()` with the initial or exact recovered position
96/// 3. Call `poll_batch()` in a loop to consume messages
97/// 4. Call `checkpoint()` for fault tolerance
98/// 5. Call `close()` for clean shutdown
99///
100/// Guaranteed delivery uses engine-owned manual assignment: the complete explicit topic inventory
101/// in embedded/single-node mode or the owned vnode subset in cluster mode. Broker consumer-group
102/// ownership can revoke a partition before an asynchronous engine checkpoint captures its final
103/// offset, so dynamic group ownership is intentionally limited to `BestEffort`.
104pub struct KafkaSource {
105    consumer: Option<Arc<StreamConsumer<LaminarConsumerContext>>>,
106    config: KafkaSourceConfig,
107    /// Delivery contract selected by the engine for this connector generation.
108    /// Only `BestEffort` may discard records that fail deserialization.
109    delivery: DeliveryGuarantee,
110    deserializer: Box<dyn RecordDeserializer>,
111    offsets: OffsetTracker,
112    state: ConnectorState,
113    metrics: KafkaSourceMetrics,
114    schema: SchemaRef,
115    channel_len: Arc<AtomicUsize>,
116    rebalance_state: Arc<Mutex<RebalanceState>>,
117    /// Shared rebalance counter bridging `LaminarConsumerContext` → `KafkaSourceMetrics`.
118    rebalance_counter: Arc<AtomicU64>,
119    /// Bumped on each partition revoke; `poll_batch` compares it lock-free to
120    /// detect a revoke and purge the lost partitions' offsets.
121    revoke_generation: Arc<AtomicU64>,
122    /// Bumped on each partition assign; the reader loop seeks the newly-assigned
123    /// partitions on change (see the seek block in `ensure_reader_started`).
124    assign_generation: Arc<AtomicU64>,
125    last_seen_revoke_gen: u64,
126    schema_registry: Option<Arc<SchemaRegistryClient>>,
127    data_ready: Arc<Notify>,
128    msg_rx: Option<KafkaReaderRx>,
129    reader_handle: Option<tokio::task::JoinHandle<()>>,
130    reader_shutdown: Option<tokio::sync::watch::Sender<bool>>,
131    /// Sole admission authority for detached work in this connector generation.
132    task_owner: ConnectorTaskOwner,
133    /// Cloneable terminal observer returned to the connector runtime.
134    task_tracker: ConnectorTaskTracker,
135    /// Blocking native calls owned and fenced by this connector generation.
136    blocking_tasks: KafkaBlockingTasks,
137    /// First terminal reader reason, published before an assignment wait can hide its exit.
138    reader_fault: Arc<Mutex<Option<Arc<str>>>>,
139    /// Allocated only for a cluster-assigned instance; local readers have no rotation path.
140    reader_drain_tx: Option<tokio::sync::mpsc::UnboundedSender<KafkaReaderDrainCommand>>,
141    source_drain: Option<KafkaSourceDrain>,
142    /// Offset snapshot for the rebalance callback's seek-on-assign, refreshed
143    /// once per `poll_batch()` cycle.
144    offset_snapshot: Arc<Mutex<OffsetTracker>>,
145    /// When set, every partition absent from the installed engine checkpoint is
146    /// explicitly positioned at the configured deterministic start. This must
147    /// remain armed for the connector lifetime: a later rebalance can introduce
148    /// a partition that was not present in the startup checkpoint, and using the
149    /// broker's stored group offset there would cross engine timelines.
150    deterministic_unrecorded_position: Arc<AtomicBool>,
151
152    /// Stable catalog source identity for namespaced cluster handoff offsets.
153    source_name: Arc<str>,
154
155    /// Cluster vnode assignment: when set, `start()` manually `assign()`s this
156    /// source's vnode-owned partitions instead of `subscribe()`, and the reader
157    /// re-binds on version rotation.
158    vnode_assignment: Option<(
159        Arc<laminar_core::state::VnodeRegistry>,
160        laminar_core::state::NodeId,
161    )>,
162    /// Precomputed source/topic/partition routes used by the per-record stale-owner fence.
163    vnode_partition_routes: KafkaPartitionRoutes,
164    /// Canonical non-vnode manual assignment used by local guaranteed, specific, and timestamp
165    /// starts. Checkpointing trusts this engine-owned inventory, never a fallible broker query.
166    manual_topic_partitions: KafkaPartitionSet,
167    /// Cached physical-channel inventory for stable engine-owned local assignments.
168    manual_input_channels: Arc<[Vec<u8>]>,
169    /// Concrete next-to-read position captured before intake for every engine-owned partition.
170    /// It remains authoritative until that partition has an accepted-record offset.
171    manual_partition_baselines: KafkaPartitionBaselines,
172    /// Exact Kafka ownership and durable handoff positions for one assignment version.
173    assignment_publication: Arc<Mutex<Arc<KafkaAssignmentPublication>>>,
174    /// Zero-cost steady-state fast path: avoids locking the rotation snapshot
175    /// when no acquired partition is waiting for its first accepted record.
176    rotation_partition_baseline_count: Arc<AtomicUsize>,
177    /// Rotation publication whose stale local offsets have already been removed.
178    /// Baseline publication is immutable within one assignment version, so this
179    /// keeps that cleanup off the steady-state poll path.
180    applied_rotation_baseline_version: Option<u64>,
181    /// Assignment whose complete batch cursor has already been produced.
182    batch_cursor_assignment_version: Option<u64>,
183    /// Assignment version whose Kafka consumer rebind and durable cursor
184    /// validation have completed. Barriers stay fenced while this trails the
185    /// registry's atomically published version.
186    reconciled_assignment_version: Arc<AtomicU64>,
187
188    /// Previous Avro writer schema, diffed against the next for evolution detection.
189    last_avro_schema: Option<SchemaRef>,
190
191    // Reusable poll_batch buffers — cleared each cycle, capacity retained.
192    poll_payloads: Vec<KafkaPayload>,
193    poll_payload_buf: Vec<u8>,
194    poll_payload_offsets: Vec<(usize, usize)>,
195    poll_meta_partitions: Vec<i32>,
196    poll_meta_offsets: Vec<i64>,
197    poll_meta_timestamps: Vec<Option<i64>>,
198    poll_meta_headers: Vec<Option<String>>,
199    /// This poll's (topic, partition, offset) triples, folded into `offsets` only after the complete
200    /// output batch is constructed — so decode or finalization failure cannot advance beyond data
201    /// that was never emitted. Reused across polls.
202    poll_staged_offsets: Vec<(Arc<str>, i32, i64)>,
203}
204
205impl KafkaSource {
206    /// Creates a new Kafka source connector with explicit schema.
207    #[must_use]
208    pub fn new(
209        schema: SchemaRef,
210        config: KafkaSourceConfig,
211        registry: Option<&prometheus::Registry>,
212    ) -> Self {
213        Self::build_base(schema, config, select_deserializer, None, registry)
214    }
215
216    /// Creates a new Kafka source connector with Schema Registry.
217    #[must_use]
218    pub fn with_schema_registry(
219        schema: SchemaRef,
220        config: KafkaSourceConfig,
221        sr_client: SchemaRegistryClient,
222    ) -> Self {
223        let sr = Arc::new(sr_client);
224        let sr_clone = Arc::clone(&sr);
225        let deser_factory = move |format: Format| -> Box<dyn RecordDeserializer> {
226            if format == Format::Avro {
227                Box::new(AvroDeserializer::with_schema_registry(sr_clone))
228            } else {
229                select_deserializer(format)
230            }
231        };
232        Self::build_base(schema, config, deser_factory, Some(sr), None)
233    }
234
235    /// Build a Schema Registry client from the parsed config, or
236    /// `Ok(None)` when `schema.registry.url` is not set.
237    fn build_sr_client(
238        config: &KafkaSourceConfig,
239    ) -> Result<Option<SchemaRegistryClient>, ConnectorError> {
240        let Some(sr_url) = config.schema_registry_url.as_ref() else {
241            return Ok(None);
242        };
243        let client = if let Some(ca) = config.schema_registry_ssl_ca_location.as_deref() {
244            SchemaRegistryClient::with_tls_mtls(
245                sr_url.clone(),
246                config.schema_registry_auth.clone(),
247                ca,
248                config.schema_registry_ssl_certificate_location.as_deref(),
249                config.schema_registry_ssl_key_location.as_deref(),
250            )?
251        } else {
252            SchemaRegistryClient::new(sr_url.clone(), config.schema_registry_auth.clone())?
253        };
254        Ok(Some(client))
255    }
256
257    fn build_base(
258        schema: SchemaRef,
259        config: KafkaSourceConfig,
260        deser_factory: impl FnOnce(Format) -> Box<dyn RecordDeserializer>,
261        schema_registry: Option<Arc<SchemaRegistryClient>>,
262        registry: Option<&prometheus::Registry>,
263    ) -> Self {
264        let deserializer = deser_factory(config.format);
265        let channel_len = Arc::new(AtomicUsize::new(0));
266        let (task_owner, task_tracker) = ConnectorTaskOwner::new();
267        let blocking_guard = task_owner
268            .track()
269            .expect("new Kafka source task generation must be live");
270
271        Self {
272            consumer: None,
273            config,
274            delivery: DeliveryGuarantee::BestEffort,
275            deserializer,
276            offsets: OffsetTracker::new(),
277            state: ConnectorState::Created,
278            metrics: KafkaSourceMetrics::new(registry),
279            schema,
280            channel_len,
281            rebalance_state: Arc::new(Mutex::new(RebalanceState::new())),
282            rebalance_counter: Arc::new(AtomicU64::new(0)),
283            revoke_generation: Arc::new(AtomicU64::new(0)),
284            assign_generation: Arc::new(AtomicU64::new(0)),
285            last_seen_revoke_gen: 0,
286            schema_registry,
287            data_ready: Arc::new(Notify::new()),
288            msg_rx: None,
289            reader_handle: None,
290            reader_shutdown: None,
291            task_owner,
292            task_tracker,
293            blocking_tasks: KafkaBlockingTasks::new(blocking_guard),
294            reader_fault: Arc::new(Mutex::new(None)),
295            reader_drain_tx: None,
296            source_drain: None,
297            offset_snapshot: Arc::new(Mutex::new(OffsetTracker::new())),
298            deterministic_unrecorded_position: Arc::new(AtomicBool::new(false)),
299            source_name: Arc::from(""),
300            vnode_assignment: None,
301            vnode_partition_routes: KafkaPartitionRoutes::new(),
302            manual_topic_partitions: std::collections::HashSet::new(),
303            manual_input_channels: Arc::from([]),
304            manual_partition_baselines: std::collections::HashMap::new(),
305            assignment_publication: Arc::new(Mutex::new(Arc::new(
306                KafkaAssignmentPublication::default(),
307            ))),
308            rotation_partition_baseline_count: Arc::new(AtomicUsize::new(0)),
309            applied_rotation_baseline_version: None,
310            batch_cursor_assignment_version: None,
311            reconciled_assignment_version: Arc::new(AtomicU64::new(0)),
312            last_avro_schema: None,
313            poll_payloads: Vec::new(),
314            poll_payload_buf: Vec::new(),
315            poll_payload_offsets: Vec::new(),
316            poll_meta_partitions: Vec::new(),
317            poll_meta_offsets: Vec::new(),
318            poll_meta_timestamps: Vec::new(),
319            poll_meta_headers: Vec::new(),
320            poll_staged_offsets: Vec::new(),
321        }
322    }
323
324    /// Lifecycle state (Created → Initializing → Running → Closed).
325    #[must_use]
326    pub fn state(&self) -> ConnectorState {
327        self.state
328    }
329
330    /// Shared backpressure fill counter for downstream wiring.
331    #[must_use]
332    pub fn channel_len(&self) -> Arc<AtomicUsize> {
333        Arc::clone(&self.channel_len)
334    }
335
336    /// Shared partition assignment state (updated by rebalance callbacks).
337    #[must_use]
338    pub fn rebalance_state(&self) -> Arc<Mutex<RebalanceState>> {
339        Arc::clone(&self.rebalance_state)
340    }
341
342    /// Whether a Schema Registry client is configured.
343    #[must_use]
344    pub fn has_schema_registry(&self) -> bool {
345        self.schema_registry.is_some()
346    }
347
348    fn fail_startup(&mut self) {
349        self.state = ConnectorState::Failed;
350        self.blocking_tasks.retire();
351        self.blocking_tasks.ensure_reaper();
352        if let Some(consumer) = self.consumer.take() {
353            consumer.unsubscribe();
354            self.blocking_tasks
355                .spawn_final_drop(consumer, "consumer after failed startup");
356        }
357    }
358
359    fn capture_drain_positions(
360        &self,
361        inputs: &[KafkaDrainPartition],
362        prepare_deadline: Option<tokio::time::Instant>,
363    ) -> Result<Arc<[KafkaDrainPosition]>, ConnectorError> {
364        if prepare_deadline.is_some_and(|deadline| tokio::time::Instant::now() >= deadline) {
365            return Err(ConnectorError::Internal(
366                "Kafka drain deadline expired before cursor capture".into(),
367            ));
368        }
369        let publication = Arc::clone(&lock_or_recover(&self.assignment_publication));
370        let mut cut = Vec::with_capacity(inputs.len());
371        for input in inputs {
372            if prepare_deadline.is_some_and(|deadline| tokio::time::Instant::now() >= deadline) {
373                return Err(ConnectorError::Internal(
374                    "Kafka drain deadline expired during cursor capture".into(),
375                ));
376            }
377            let next_offset =
378                if let Some(offset) = self.offsets.get(input.topic.as_ref(), input.partition) {
379                    offset.checked_add(1).ok_or_else(|| {
380                        ConnectorError::Internal(format!(
381                            "Kafka drain offset overflow for '{}-{}'",
382                            input.topic, input.partition
383                        ))
384                    })?
385                } else if let Some(next) = rotation_partition_baseline(
386                    &publication.baselines,
387                    input.topic.as_ref(),
388                    input.partition,
389                ) {
390                    next
391                } else if let Some(next) = self
392                    .manual_partition_baselines
393                    .get(&(input.topic.to_string(), input.partition))
394                {
395                    *next
396                } else {
397                    return Err(ConnectorError::InvalidState {
398                        expected: format!(
399                            "numeric next-to-read position for drained Kafka input '{}-{}'",
400                            input.topic, input.partition
401                        ),
402                        actual: "no accepted offset or deterministic baseline".into(),
403                    });
404                };
405            cut.push(KafkaDrainPosition {
406                topic: Arc::clone(&input.topic),
407                partition: input.partition,
408                next_offset,
409            });
410        }
411        if prepare_deadline.is_some_and(|deadline| tokio::time::Instant::now() >= deadline) {
412            return Err(ConnectorError::Internal(
413                "Kafka drain deadline expired during cursor capture".into(),
414            ));
415        }
416        Ok(cut.into())
417    }
418
419    fn check_reader_health(&self, operation: &str) -> Result<(), ConnectorError> {
420        if let Some(reason) = lock_or_recover(&self.reader_fault).clone() {
421            return Err(ConnectorError::Internal(format!(
422                "Kafka reader terminated while {operation}: {reason}"
423            )));
424        }
425        if self
426            .reader_handle
427            .as_ref()
428            .is_some_and(tokio::task::JoinHandle::is_finished)
429        {
430            return Err(ConnectorError::Internal(format!(
431                "Kafka reader task exited while {operation}"
432            )));
433        }
434        Ok(())
435    }
436
437    fn validate_active_drain_cursor(&self) -> Result<(), ConnectorError> {
438        self.check_reader_health("validating a drain cursor")?;
439        let Some(active) = self.source_drain.as_ref() else {
440            return Ok(());
441        };
442        let Some(expected) = active.cut.as_deref() else {
443            return Ok(());
444        };
445        let Some(boundary) = active.boundary.as_ref() else {
446            return Err(ConnectorError::Internal(
447                "Kafka drain cut exists without its FIFO boundary".into(),
448            ));
449        };
450        // The certified cut remains valid while terminal materialization is pending; the
451        // preparation deadline no longer applies after the cut has been captured.
452        let current = self.capture_drain_positions(&boundary.inputs, None)?;
453        if current.as_ref() != expected {
454            return Err(ConnectorError::Internal(
455                "Kafka input cursor advanced after its drain receipt".into(),
456            ));
457        }
458        Ok(())
459    }
460}
461
462impl Drop for KafkaSource {
463    fn drop(&mut self) {
464        if let Some(shutdown) = self.reader_shutdown.take() {
465            let _ = shutdown.send(true);
466        }
467        if let Some(consumer) = self.consumer.as_ref() {
468            consumer.unsubscribe();
469        }
470        self.blocking_tasks.retire();
471        if let Some(handle) = self.reader_handle.take() {
472            ensure_background_task_reaper(handle, &self.task_owner, "reader");
473        }
474        self.blocking_tasks.ensure_reaper();
475        if let Some(consumer) = self.consumer.take() {
476            self.blocking_tasks.spawn_final_drop(consumer, "consumer");
477        }
478    }
479}
480
481impl std::fmt::Debug for KafkaSource {
482    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
483        f.debug_struct("KafkaSource")
484            .field("state", &self.state)
485            .field("subscription", &self.config.subscription)
486            .field("group_id", &self.config.group_id)
487            .field("format", &self.config.format)
488            .field("delivery", &self.delivery)
489            .field("partitions", &self.offsets.partition_count())
490            .finish_non_exhaustive()
491    }
492}
493
494/// Warn if the CREATE-SOURCE catalog schema has drifted from the live
495/// Schema Registry schema. Empty `declared` means nothing was declared.
496fn log_schema_drift(declared: &arrow_schema::Schema, live: &arrow_schema::Schema, subject: &str) {
497    if declared.fields().is_empty() || declared.fields() == live.fields() {
498        return;
499    }
500    let decl: BTreeSet<&str> = declared
501        .fields()
502        .iter()
503        .map(|f| f.name().as_str())
504        .collect();
505    let lv: BTreeSet<&str> = live.fields().iter().map(|f| f.name().as_str()).collect();
506    warn!(
507        %subject,
508        missing_in_sr = ?decl.difference(&lv).collect::<Vec<_>>(),
509        added_in_sr = ?lv.difference(&decl).collect::<Vec<_>>(),
510        "schema drift: re-apply CREATE SOURCE DDL to pick up the current SR schema"
511    );
512}
513
514fn select_deserializer(format: Format) -> Box<dyn RecordDeserializer> {
515    match format {
516        Format::Avro => Box::new(AvroDeserializer::new()),
517        other => serde::create_deserializer(other).unwrap_or_else(|_| {
518            warn!(format = %other, "unsupported format, falling back to JSON");
519            Box::new(serde::json::JsonDeserializer::new())
520        }),
521    }
522}
523
524#[cfg(test)]
525mod tests;
laminar_connectors/kafka/source/mod.rs

laminar_connectors/kafka/source/
mod.rs
