risingwave_connector/source/kafka/source/
reader.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::hash_map::Entry;
16use std::collections::{HashMap, HashSet};
17use std::mem::swap;
18use std::time::Duration;
19
20use anyhow::Context;
21use async_trait::async_trait;
22use futures::{StreamExt, TryStreamExt};
23use futures_async_stream::try_stream;
24use rdkafka::consumer::{Consumer, StreamConsumer};
25use rdkafka::error::KafkaError;
26use rdkafka::{ClientConfig, Message, Offset, TopicPartitionList};
27use risingwave_common::metrics::{
28    LabelGuardedIntCounter, LabelGuardedIntGauge, RelabeledAggregatedIntGauge,
29};
30use risingwave_pb::plan_common::additional_column::ColumnType as AdditionalColumnType;
31
32use crate::connector_common::read_kafka_log_level;
33use crate::error::ConnectorResult as Result;
34use crate::parser::ParserConfig;
35use crate::source::base::SourceMessage;
36use crate::source::kafka::{
37    KAFKA_ISOLATION_LEVEL, KafkaContextCommon, KafkaProperties, KafkaSplit, RwConsumerContext,
38};
39use crate::source::{
40    BackfillInfo, BoxSourceChunkStream, BoxSourceReaderEventStream, Column, SourceContextRef,
41    SourceMessageEvent, SplitId, SplitImpl, SplitMetaData, SplitReader, into_chunk_event_stream,
42    into_chunk_stream,
43};
44
45pub struct KafkaSplitReader {
46    consumer: StreamConsumer<RwConsumerContext>,
47    offsets: HashMap<SplitId, (Option<i64>, Option<i64>)>,
48    backfill_info: HashMap<SplitId, BackfillInfo>,
49    /// The furthest inclusive offset that is known to exist when the reader observes EOF for a
50    /// split, even if that offset does not correspond to a visible data record.
51    known_eof_offsets: HashMap<SplitId, i64>,
52    splits: Vec<KafkaSplit>,
53    sync_call_timeout: Duration,
54    bytes_per_second: usize,
55    max_num_messages: usize,
56    parser_config: ParserConfig,
57    source_ctx: SourceContextRef,
58}
59
60#[async_trait]
61impl SplitReader for KafkaSplitReader {
62    type Properties = KafkaProperties;
63    type Split = KafkaSplit;
64
65    async fn new(
66        properties: KafkaProperties,
67        splits: Vec<KafkaSplit>,
68        parser_config: ParserConfig,
69        source_ctx: SourceContextRef,
70        _columns: Option<Vec<Column>>,
71    ) -> Result<Self> {
72        let mut config = ClientConfig::new();
73
74        let bootstrap_servers = &properties.connection.brokers;
75        let broker_rewrite_map = properties.privatelink_common.broker_rewrite_map.clone();
76
77        // Enable partition EOF to emit split progress updates when the fetched data ends with
78        // non-deliverable records, e.g. transactional control records in read-committed mode.
79        config.set("enable.partition.eof", "true");
80        config.set("auto.offset.reset", "smallest");
81        config.set("isolation.level", KAFKA_ISOLATION_LEVEL);
82        config.set("bootstrap.servers", bootstrap_servers);
83
84        properties.connection.set_security_properties(&mut config);
85        properties.set_client(&mut config);
86
87        config.set("group.id", properties.group_id(source_ctx.fragment_id));
88
89        let ctx_common = KafkaContextCommon::new(
90            broker_rewrite_map,
91            Some(format!(
92                "fragment-{}-source-{}-actor-{}",
93                source_ctx.fragment_id, source_ctx.source_id, source_ctx.actor_id
94            )),
95            // thread consumer will keep polling in the background, we don't need to call `poll`
96            // explicitly
97            Some(source_ctx.metrics.rdkafka_native_metric.clone()),
98            properties.aws_auth_props,
99            properties.connection.is_aws_msk_iam(),
100        )
101        .await?;
102
103        let client_ctx = RwConsumerContext::new(ctx_common);
104
105        if let Some(log_level) = read_kafka_log_level() {
106            config.set_log_level(log_level);
107        }
108        let consumer: StreamConsumer<RwConsumerContext> = config
109            .create_with_context(client_ctx)
110            .await
111            .context("failed to create kafka consumer")?;
112        // Skip the final offset commit on drop. The consumer is closed in the background, so that
113        // commit could land after meta has deleted the group and recreate it. Committed offsets
114        // only expose progress to external tools.
115        consumer.set_no_consumer_close_on_drop();
116
117        let mut tpl = TopicPartitionList::with_capacity(splits.len());
118
119        let mut offsets = HashMap::new();
120        let mut backfill_info = HashMap::new();
121        let mut known_eof_offsets = HashMap::new();
122        for split in splits.clone() {
123            offsets.insert(split.id(), (split.start_offset, split.stop_offset));
124
125            if let Some(offset) = split.start_offset {
126                tpl.add_partition_offset(
127                    split.topic.as_str(),
128                    split.partition,
129                    Offset::Offset(offset + 1),
130                )?;
131            } else {
132                tpl.add_partition(split.topic.as_str(), split.partition);
133            }
134
135            let (low, high) = consumer
136                .fetch_watermarks(
137                    split.topic.as_str(),
138                    split.partition,
139                    properties.common.sync_call_timeout,
140                )
141                .await?;
142            tracing::info!("fetch kafka watermarks: low: {low}, high: {high}, split: {split:?}");
143            // note: low is inclusive, high is exclusive, start_offset is exclusive
144            let has_data_to_backfill =
145                low != high && split.start_offset.is_none_or(|offset| offset + 1 < high);
146            let watermark_offset = if has_data_to_backfill {
147                debug_assert!(high > 0);
148                Some(high - 1)
149            } else {
150                None
151            };
152            if !has_data_to_backfill {
153                backfill_info.insert(split.id(), BackfillInfo::NoDataToBackfill);
154            } else {
155                backfill_info.insert(
156                    split.id(),
157                    BackfillInfo::HasDataToBackfill {
158                        latest_offset: (high - 1).to_string(),
159                    },
160                );
161            }
162
163            let stop_offset = split.stop_offset.and_then(Self::offset_before);
164            if let Some(known_eof_offset) = Self::max_known_offset(watermark_offset, stop_offset) {
165                known_eof_offsets.insert(split.id(), known_eof_offset);
166            }
167        }
168        tracing::info!(
169            topic = properties.common.topic,
170            source_name = source_ctx.source_name,
171            fragment_id = %source_ctx.fragment_id,
172            source_id = %source_ctx.source_id,
173            actor_id = %source_ctx.actor_id,
174            "backfill_info: {:?}",
175            backfill_info
176        );
177
178        consumer.assign(&tpl)?;
179
180        // The two parameters below are only used by developers for performance testing purposes,
181        // so we panic here on purpose if the input is not correctly recognized.
182        let bytes_per_second = match properties.bytes_per_second {
183            None => usize::MAX,
184            Some(number) => number
185                .parse::<usize>()
186                .expect("bytes.per.second expect usize"),
187        };
188        let max_num_messages = match properties.max_num_messages {
189            None => usize::MAX,
190            Some(number) => number
191                .parse::<usize>()
192                .expect("max.num.messages expect usize"),
193        };
194
195        Ok(Self {
196            consumer,
197            offsets,
198            splits,
199            backfill_info,
200            known_eof_offsets,
201            bytes_per_second,
202            sync_call_timeout: properties.common.sync_call_timeout,
203            max_num_messages,
204            parser_config,
205            source_ctx,
206        })
207    }
208
209    fn into_stream(self) -> BoxSourceChunkStream {
210        let parser_config = self.parser_config.clone();
211        let source_context = self.source_ctx.clone();
212        let data_stream = self
213            .into_data_event_stream()
214            .try_filter_map(|event| async move {
215                Ok(match event {
216                    SourceMessageEvent::Data(batch) => Some(batch),
217                    SourceMessageEvent::SplitProgress(_) => None,
218                })
219            });
220        into_chunk_stream(data_stream, parser_config, source_context)
221    }
222
223    fn into_event_stream(self) -> BoxSourceReaderEventStream {
224        let parser_config = self.parser_config.clone();
225        let source_context = self.source_ctx.clone();
226        into_chunk_event_stream(self.into_data_event_stream(), parser_config, source_context)
227    }
228
229    fn backfill_info(&self) -> HashMap<SplitId, BackfillInfo> {
230        self.backfill_info.clone()
231    }
232
233    async fn seek_to_latest(&mut self) -> Result<Vec<SplitImpl>> {
234        let mut latest_splits: Vec<SplitImpl> = Vec::new();
235        let mut tpl = TopicPartitionList::with_capacity(self.splits.len());
236        for mut split in self.splits.clone() {
237            // we can't get latest offset if we use Offset::End, so we just fetch watermark here.
238            let (_low, high) = self
239                .consumer
240                .fetch_watermarks(
241                    split.topic.as_str(),
242                    split.partition,
243                    self.sync_call_timeout,
244                )
245                .await?;
246            tpl.add_partition_offset(split.topic.as_str(), split.partition, Offset::Offset(high))?;
247            split.start_offset = Some(high - 1);
248            latest_splits.push(split.into());
249        }
250        // replace the previous assignment
251        self.consumer.assign(&tpl)?;
252        Ok(latest_splits)
253    }
254}
255
256impl KafkaSplitReader {
257    fn offset_before(next_offset: i64) -> Option<i64> {
258        next_offset.checked_sub(1).filter(|offset| *offset >= 0)
259    }
260
261    fn max_known_offset(current_offset: Option<i64>, candidate_offset: Option<i64>) -> Option<i64> {
262        match (current_offset, candidate_offset) {
263            (Some(current_offset), Some(candidate_offset)) => {
264                Some(current_offset.max(candidate_offset))
265            }
266            (Some(current_offset), None) => Some(current_offset),
267            (None, Some(candidate_offset)) => Some(candidate_offset),
268            (None, None) => None,
269        }
270    }
271
272    fn record_split_progress(
273        progress: &mut HashMap<SplitId, String>,
274        latest_progress_offsets: &mut HashMap<SplitId, i64>,
275        split_id: &SplitId,
276        inclusive_offset: i64,
277    ) {
278        let should_emit = latest_progress_offsets
279            .get(split_id)
280            .is_none_or(|offset| *offset < inclusive_offset);
281        if should_emit {
282            latest_progress_offsets.insert(split_id.clone(), inclusive_offset);
283            progress.insert(split_id.clone(), inclusive_offset.to_string());
284        }
285    }
286
287    fn drain_stop_offsets_with_progress(
288        stop_offsets: &mut HashMap<SplitId, i64>,
289        progress: &HashMap<SplitId, String>,
290    ) {
291        let mut finished_splits = Vec::new();
292
293        for (split_id, progress_offset) in progress {
294            let Some(stop_offset) = stop_offsets.get(split_id).copied() else {
295                continue;
296            };
297            let Ok(progress_offset) = progress_offset.parse::<i64>() else {
298                tracing::warn!(
299                    split_id = split_id.as_ref(),
300                    progress_offset = progress_offset.as_str(),
301                    "invalid split progress offset from kafka reader"
302                );
303                continue;
304            };
305
306            // `stop_offset` is exclusive while progress is inclusive.
307            if progress_offset >= stop_offset - 1 {
308                tracing::debug!(
309                    split_id = split_id.as_ref(),
310                    stop_offset,
311                    progress_offset,
312                    "stop offset reached by split progress"
313                );
314                finished_splits.push(split_id.clone());
315            }
316        }
317
318        for split_id in finished_splits {
319            stop_offsets.remove(&split_id);
320        }
321    }
322
323    fn snapshot_split_progress(
324        eof_offsets: &HashMap<SplitId, i64>,
325        latest_progress_offsets: &mut HashMap<SplitId, i64>,
326    ) -> Option<HashMap<SplitId, String>> {
327        let mut progress = HashMap::new();
328
329        for (split_id, inclusive_offset) in eof_offsets {
330            Self::record_split_progress(
331                &mut progress,
332                latest_progress_offsets,
333                split_id,
334                *inclusive_offset,
335            );
336        }
337
338        (!progress.is_empty()).then_some(progress)
339    }
340
341    fn resolve_eof_offsets(&self, eof_partitions: &HashSet<i32>) -> Result<HashMap<SplitId, i64>> {
342        let positions = self.consumer.position()?;
343        let mut eof_offsets = HashMap::new();
344
345        for split in &self.splits {
346            if !eof_partitions.contains(&split.partition) {
347                continue;
348            }
349
350            let split_id = split.id();
351            let position_offset = positions
352                .find_partition(split.topic.as_str(), split.partition)
353                .and_then(|position| match position.offset() {
354                    Offset::Offset(next_offset) => Self::offset_before(next_offset),
355                    _ => None,
356                });
357            let Some(inclusive_offset) = Self::max_known_offset(
358                position_offset,
359                self.known_eof_offsets.get(&split_id).copied(),
360            ) else {
361                continue;
362            };
363            eof_offsets.insert(split_id, inclusive_offset);
364        }
365
366        Ok(eof_offsets)
367    }
368
369    fn apply_split_progress(
370        stop_offsets: &mut HashMap<SplitId, i64>,
371        split_progress: Option<HashMap<SplitId, String>>,
372        is_bounded: bool,
373    ) -> Option<HashMap<SplitId, String>> {
374        let progress = split_progress?;
375        if is_bounded {
376            Self::drain_stop_offsets_with_progress(stop_offsets, &progress);
377        }
378        Some(progress)
379    }
380
381    #[try_stream(ok = SourceMessageEvent, error = crate::error::ConnectorError)]
382    async fn into_data_event_stream(self) {
383        if self.offsets.values().all(|(start_offset, stop_offset)| {
384            match (start_offset, stop_offset) {
385                (Some(start), Some(stop)) if (*start + 1) >= *stop => true,
386                (_, Some(stop)) if *stop == 0 => true,
387                _ => false,
388            }
389        }) {
390            yield SourceMessageEvent::Data(Vec::new());
391            return Ok(());
392        };
393
394        let mut stop_offsets: HashMap<_, _> = self
395            .offsets
396            .iter()
397            .flat_map(|(split_id, (_, stop_offset))| {
398                stop_offset
399                    .filter(|offset| *offset > 0)
400                    .map(|offset| (split_id.clone() as SplitId, offset))
401            })
402            .collect();
403        let is_bounded = !stop_offsets.is_empty();
404
405        let mut interval = tokio::time::interval(Duration::from_secs(1));
406        interval.tick().await;
407        let mut bytes_current_second = 0;
408        let mut num_messages = 0;
409        let max_chunk_size = self.source_ctx.source_ctrl_opts.chunk_size;
410        let mut res = Vec::with_capacity(max_chunk_size);
411        // ingest kafka message header can be expensive, do it only when required
412        let require_message_header = self.parser_config.common.rw_columns.iter().any(|col_desc| {
413            matches!(
414                col_desc.additional_column.column_type,
415                Some(AdditionalColumnType::Headers(_) | AdditionalColumnType::HeaderInner(_))
416            )
417        });
418
419        let mut latest_message_id_metrics: HashMap<String, RelabeledAggregatedIntGauge> =
420            HashMap::new();
421        let mut partition_eof_count_metrics: HashMap<String, LabelGuardedIntCounter> =
422            HashMap::new();
423        let mut partition_eof_offset_metrics: HashMap<String, LabelGuardedIntGauge> =
424            HashMap::new();
425        let mut latest_progress_offsets = self
426            .splits
427            .iter()
428            .filter_map(|split| split.start_offset.map(|offset| (split.id(), offset)))
429            .collect::<HashMap<_, _>>();
430
431        #[for_await]
432        'for_outer_loop: for raw_msgs in self.consumer.stream().ready_chunks(max_chunk_size) {
433            let mut eof_partitions = HashSet::new();
434            let mut msgs = Vec::with_capacity(max_chunk_size);
435            for msg in raw_msgs {
436                match msg {
437                    Ok(msg) => msgs.push(msg),
438                    Err(KafkaError::PartitionEOF(partition)) => {
439                        let split_id = partition.to_string();
440                        partition_eof_count_metrics
441                            .entry(split_id.clone())
442                            .or_insert_with(|| {
443                                self.source_ctx
444                                    .metrics
445                                    .partition_eof_count
446                                    .with_guarded_label_values(&[
447                                        &self.source_ctx.source_id.to_string(),
448                                        &split_id,
449                                        &self.source_ctx.source_name,
450                                        &self.source_ctx.fragment_id.to_string(),
451                                    ])
452                            })
453                            .inc();
454                        eof_partitions.insert(partition);
455                    }
456                    Err(err) => return Err(err.into()),
457                }
458            }
459
460            let eof_offsets = if eof_partitions.is_empty() {
461                HashMap::new()
462            } else {
463                self.resolve_eof_offsets(&eof_partitions)?
464            };
465            for (split_id, eof_offset) in &eof_offsets {
466                let split_id_label = split_id.as_ref().to_owned();
467                partition_eof_offset_metrics
468                    .entry(split_id_label.clone())
469                    .or_insert_with(|| {
470                        self.source_ctx
471                            .metrics
472                            .partition_eof_offset
473                            .with_guarded_label_values(&[
474                                &self.source_ctx.source_id.to_string(),
475                                &split_id_label,
476                                &self.source_ctx.source_name,
477                                &self.source_ctx.fragment_id.to_string(),
478                            ])
479                    })
480                    .set(*eof_offset);
481                tracing::info!(
482                    actor_id = %self.source_ctx.actor_id,
483                    source_id = %self.source_ctx.source_id,
484                    source_name = self.source_ctx.source_name,
485                    fragment_id = %self.source_ctx.fragment_id,
486                    split_id = %split_id,
487                    eof_offset,
488                    "received kafka partition EOF"
489                );
490            }
491            let split_progress =
492                Self::snapshot_split_progress(&eof_offsets, &mut latest_progress_offsets);
493            let split_progress =
494                Self::apply_split_progress(&mut stop_offsets, split_progress, is_bounded);
495
496            if msgs.is_empty() {
497                if let Some(progress) = split_progress {
498                    yield SourceMessageEvent::SplitProgress(progress);
499                }
500                if is_bounded && stop_offsets.is_empty() {
501                    if !res.is_empty() {
502                        yield SourceMessageEvent::Data(res);
503                    }
504                    break 'for_outer_loop;
505                }
506                continue;
507            }
508
509            let mut split_msg_offsets = HashMap::new();
510
511            for msg in &msgs {
512                split_msg_offsets.insert(msg.partition(), msg.offset());
513            }
514
515            for (partition, offset) in split_msg_offsets {
516                let split_id = partition.to_string();
517                latest_message_id_metrics
518                    .entry(split_id.clone())
519                    .or_insert_with(|| {
520                        self.source_ctx
521                            .metrics
522                            .latest_message_id
523                            .with_guarded_label_values(&[
524                                // source name is not available here
525                                &self.source_ctx.actor_id.to_string(),
526                                &self.source_ctx.source_id.to_string(),
527                                &split_id,
528                            ])
529                    })
530                    .set(offset);
531            }
532
533            for msg in msgs {
534                let cur_offset = msg.offset();
535                bytes_current_second += match &msg.payload() {
536                    None => 0,
537                    Some(payload) => payload.len(),
538                };
539                num_messages += 1;
540                let source_message =
541                    SourceMessage::from_kafka_message(&msg, require_message_header);
542                let split_id = source_message.split_id.clone();
543                res.push(source_message);
544
545                if let Entry::Occupied(o) = stop_offsets.entry(split_id) {
546                    let stop_offset = *o.get();
547
548                    if cur_offset == stop_offset - 1 {
549                        tracing::debug!(
550                            "stop offset reached for split {}, stop reading, offset: {}, stop offset: {}",
551                            o.key(),
552                            cur_offset,
553                            stop_offset
554                        );
555
556                        o.remove();
557
558                        if stop_offsets.is_empty() {
559                            yield SourceMessageEvent::Data(res);
560                            break 'for_outer_loop;
561                        }
562
563                        continue;
564                    }
565                }
566
567                if bytes_current_second > self.bytes_per_second {
568                    // swap to make compiler happy
569                    let mut cur = Vec::with_capacity(res.capacity());
570                    swap(&mut cur, &mut res);
571                    yield SourceMessageEvent::Data(cur);
572                    interval.tick().await;
573                    bytes_current_second = 0;
574                    res.clear();
575                }
576                if num_messages >= self.max_num_messages {
577                    yield SourceMessageEvent::Data(res);
578                    break 'for_outer_loop;
579                }
580            }
581            let mut cur = Vec::with_capacity(res.capacity());
582            swap(&mut cur, &mut res);
583            yield SourceMessageEvent::Data(cur);
584            if let Some(progress) = split_progress {
585                yield SourceMessageEvent::SplitProgress(progress);
586            }
587            if is_bounded && stop_offsets.is_empty() {
588                break 'for_outer_loop;
589            }
590            // don't clear `bytes_current_second` here as it is only related to `.tick()`.
591            // yield in the outer loop so that we can always guarantee that some messages are read
592            // every `MAX_CHUNK_SIZE`.
593        }
594        tracing::info!("kafka reader finished");
595    }
596}
597
598#[cfg(test)]
599mod tests {
600    use std::collections::HashMap;
601
602    use super::KafkaSplitReader;
603    use crate::source::SplitId;
604
605    #[test]
606    fn test_drain_stop_offsets_with_progress() {
607        let mut stop_offsets = HashMap::from([
608            (SplitId::from("0"), 10_i64),
609            (SplitId::from("1"), 15_i64),
610            (SplitId::from("2"), 20_i64),
611        ]);
612        let progress = HashMap::from([
613            (SplitId::from("0"), "9".to_owned()),
614            (SplitId::from("1"), "13".to_owned()),
615            (SplitId::from("2"), "invalid".to_owned()),
616        ]);
617
618        KafkaSplitReader::drain_stop_offsets_with_progress(&mut stop_offsets, &progress);
619
620        assert!(!stop_offsets.contains_key("0"));
621        assert_eq!(stop_offsets.get("1"), Some(&15));
622        assert_eq!(stop_offsets.get("2"), Some(&20));
623    }
624
625    #[test]
626    fn test_drain_stop_offsets_with_progress_for_zero_stop_offset() {
627        let mut stop_offsets = HashMap::from([(SplitId::from("0"), 0_i64)]);
628        let progress = HashMap::from([(SplitId::from("0"), "0".to_owned())]);
629
630        KafkaSplitReader::drain_stop_offsets_with_progress(&mut stop_offsets, &progress);
631
632        assert!(stop_offsets.is_empty());
633    }
634
635    #[test]
636    fn test_max_known_offset_prefers_larger_offset() {
637        assert_eq!(
638            KafkaSplitReader::max_known_offset(Some(3), Some(4)),
639            Some(4)
640        );
641        assert_eq!(
642            KafkaSplitReader::max_known_offset(Some(7), Some(4)),
643            Some(7)
644        );
645        assert_eq!(KafkaSplitReader::max_known_offset(None, Some(4)), Some(4));
646        assert_eq!(KafkaSplitReader::max_known_offset(Some(4), None), Some(4));
647    }
648
649    #[test]
650    fn test_apply_split_progress_for_empty_message_batch() {
651        let mut stop_offsets = HashMap::from([(SplitId::from("0"), 5_i64)]);
652        let split_progress = Some(HashMap::from([(SplitId::from("0"), "4".to_owned())]));
653
654        let applied =
655            KafkaSplitReader::apply_split_progress(&mut stop_offsets, split_progress, true);
656
657        assert_eq!(
658            applied,
659            Some(HashMap::from([(SplitId::from("0"), "4".to_owned())]))
660        );
661        assert!(stop_offsets.is_empty());
662    }
663}
risingwave_connector/source/kafka/source/reader.rs

risingwave_connector/source/kafka/source/
reader.rs
