log_store/kafka/

consumer.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod range;

pub(crate) use range::RegionWalRange;

use std::collections::VecDeque;
use std::ops::Range;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use common_telemetry::debug;
use derive_builder::Builder;
use futures::future::{BoxFuture, Fuse, FusedFuture};
use futures::{FutureExt, Stream};
use pin_project::pin_project;
use rskafka::client::partition::PartitionClient;
use rskafka::record::RecordAndOffset;

use crate::kafka::consumer::range::NextBatchHint;

pub struct FetchResult {
    /// The offsets of the fetched records.
    pub records: Vec<RecordAndOffset>,

    /// The high watermark of the partition.
    pub high_watermark: i64,

    /// The size of the response encoded in bytes.
    pub encoded_response_size: usize,
}

#[async_trait::async_trait]
pub trait FetchClient: std::fmt::Debug + Send + Sync {
    /// Fetch records.
    ///
    /// Arguments are identical to [`PartitionClient::fetch_records`].
    async fn fetch_records(
        &self,
        offset: i64,
        bytes: Range<i32>,
        max_wait_ms: i32,
    ) -> rskafka::client::error::Result<FetchResult>;

    fn topic(&self) -> &str;
}

#[async_trait::async_trait]
impl FetchClient for PartitionClient {
    async fn fetch_records(
        &self,
        offset: i64,
        bytes: Range<i32>,
        max_wait_ms: i32,
    ) -> rskafka::client::error::Result<FetchResult> {
        self.fetch_records(offset, bytes, max_wait_ms)
            .await
            .map(|r| FetchResult {
                records: r.records,
                high_watermark: r.high_watermark,
                encoded_response_size: r.encoded_response_size,
            })
    }

    fn topic(&self) -> &str {
        self.topic()
    }
}

struct FetchResultInner {
    records_and_offsets: Vec<RecordAndOffset>,
    batch_size: usize,
    fetch_bytes: usize,
    watermark: i64,
    used_offset: i64,
}

const MAX_BATCH_SIZE: usize = 52428800;
const AVG_RECORD_SIZE: usize = 256 * 1024;

/// The [`Consumer`] struct represents a Kafka consumer that fetches messages from
/// a Kafka cluster. Yielding records respecting the [`RegionWalRange`].
#[pin_project]
#[derive(Builder)]
#[builder(pattern = "owned")]
pub struct Consumer {
    #[builder(default = "-1")]
    last_high_watermark: i64,

    /// The client is used to fetch records from kafka topic.
    client: Arc<dyn FetchClient>,

    /// The max batch size in a single fetch request.
    #[builder(default = "MAX_BATCH_SIZE")]
    max_batch_size: usize,

    /// The max wait milliseconds.
    #[builder(default = "500")]
    max_wait_ms: u32,

    /// The avg record size
    #[builder(default = "AVG_RECORD_SIZE")]
    avg_record_size: usize,

    /// Termination flag
    #[builder(default = "false")]
    terminated: bool,

    /// The buffer of records.
    buffer: RecordsBuffer,

    /// The fetch future.
    #[builder(default = "Fuse::terminated()")]
    fetch_fut: Fuse<BoxFuture<'static, rskafka::client::error::Result<FetchResultInner>>>,

    /// Total fetched bytes.
    #[builder(default = "0")]
    total_fetched_bytes: u64,
}

impl Consumer {
    /// Returns the total fetched bytes.
    pub fn total_fetched_bytes(&self) -> u64 {
        self.total_fetched_bytes
    }

    /// Returns the topic name.
    pub fn topic(&self) -> &str {
        self.client.topic()
    }
}

pub(crate) struct RecordsBuffer {
    buffer: VecDeque<RecordAndOffset>,

    index: RegionWalRange,
}

impl RecordsBuffer {
    /// Creates an empty [`RecordsBuffer`]
    pub fn new(index: RegionWalRange) -> Self {
        RecordsBuffer {
            buffer: VecDeque::new(),
            index,
        }
    }
}

impl RecordsBuffer {
    fn pop_front(&mut self) -> Option<RecordAndOffset> {
        while let Some(index) = self.index.peek() {
            let record_and_offset = self.buffer.pop_front()?;
            if index == record_and_offset.offset as u64 {
                self.index.next();
                return Some(record_and_offset);
            }
        }

        self.buffer.clear();
        None
    }

    fn extend(&mut self, records: Vec<RecordAndOffset>) {
        if let (Some(first), Some(index)) = (records.first(), self.index.peek()) {
            // TODO(weny): throw an error?
            assert!(
                index <= first.offset as u64,
                "index: {index}, first offset: {}",
                first.offset
            );
        }
        self.buffer.extend(records);
    }
}

impl Stream for Consumer {
    type Item = rskafka::client::error::Result<(RecordAndOffset, i64)>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.project();

        loop {
            if *this.terminated {
                return Poll::Ready(None);
            }

            if this.buffer.index.peek().is_none() {
                return Poll::Ready(None);
            }

            if let Some(x) = this.buffer.pop_front() {
                debug!("Yielding record with offset: {}", x.offset);
                return Poll::Ready(Some(Ok((x, *this.last_high_watermark))));
            }

            if this.fetch_fut.is_terminated() {
                match this.buffer.index.peek() {
                    Some(next_offset) => {
                        let client = Arc::clone(this.client);
                        let max_wait_ms = *this.max_wait_ms as i32;
                        let offset = next_offset as i64;
                        let NextBatchHint { bytes, len } = this
                            .buffer
                            .index
                            .next_batch_hint(*this.avg_record_size)
                            .unwrap_or(NextBatchHint {
                                bytes: *this.avg_record_size,
                                len: 1,
                            });

                        let fetch_range =
                            1i32..(bytes.saturating_add(1).min(*this.max_batch_size) as i32);
                        *this.fetch_fut = FutureExt::fuse(Box::pin(async move {
                            let FetchResult {
                                records: records_and_offsets,
                                high_watermark: watermark,
                                encoded_response_size,
                                ..
                            } = client
                                .fetch_records(offset, fetch_range, max_wait_ms)
                                .await?;

                            Ok(FetchResultInner {
                                records_and_offsets,
                                watermark,
                                used_offset: offset,
                                fetch_bytes: encoded_response_size,
                                batch_size: len,
                            })
                        }));
                    }
                    None => {
                        return Poll::Ready(None);
                    }
                }
            }

            let data = futures::ready!(this.fetch_fut.poll_unpin(cx));

            match data {
                Ok(FetchResultInner {
                    mut records_and_offsets,
                    watermark,
                    used_offset,
                    fetch_bytes,
                    batch_size,
                }) => {
                    // Sort records by offset in case they aren't in order
                    records_and_offsets.sort_unstable_by_key(|x| x.offset);
                    *this.last_high_watermark = watermark;
                    if !records_and_offsets.is_empty() {
                        *this.avg_record_size = fetch_bytes / records_and_offsets.len();
                        debug!("set avg_record_size: {}", *this.avg_record_size);
                    }
                    *this.total_fetched_bytes += fetch_bytes as u64;

                    debug!(
                        "Fetch result: {:?}, used_offset: {used_offset}, max_batch_size: {fetch_bytes}, expected batch_num: {batch_size}, actual batch_num: {}",
                        records_and_offsets
                            .iter()
                            .map(|record| record.offset)
                            .collect::<Vec<_>>(),
                        records_and_offsets.len()
                    );
                    this.buffer.extend(records_and_offsets);
                    continue;
                }
                Err(e) => {
                    *this.terminated = true;

                    return Poll::Ready(Some(Err(e)));
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::VecDeque;
    use std::ops::Range;
    use std::sync::Arc;

    use chrono::{TimeZone, Utc};
    use futures::TryStreamExt;
    use futures::future::Fuse;
    use rskafka::record::{Record, RecordAndOffset};

    use super::*;
    use crate::kafka::consumer::{Consumer, RecordsBuffer};

    #[derive(Debug)]
    struct MockFetchClient {
        record: Record,
    }

    #[async_trait::async_trait]
    impl FetchClient for MockFetchClient {
        async fn fetch_records(
            &self,
            offset: i64,
            bytes: Range<i32>,
            _max_wait_ms: i32,
        ) -> rskafka::client::error::Result<FetchResult> {
            let record_size = self.record.approximate_size();
            let num = (bytes.end.unsigned_abs() as usize / record_size).max(1);

            let records = (0..num)
                .map(|idx| RecordAndOffset {
                    record: self.record.clone(),
                    offset: offset + idx as i64,
                })
                .collect::<Vec<_>>();

            let max_offset = offset + records.len() as i64;
            let encoded_response_size = records.iter().map(|r| r.record.approximate_size()).sum();
            Ok(FetchResult {
                records,
                high_watermark: max_offset,
                encoded_response_size,
            })
        }

        fn topic(&self) -> &str {
            "test"
        }
    }

    fn test_record() -> Record {
        Record {
            key: Some(vec![0; 4]),
            value: Some(vec![0; 6]),
            headers: Default::default(),
            timestamp: Utc.timestamp_millis_opt(1337).unwrap(),
        }
    }

    #[tokio::test]
    async fn test_consumer_without_index() {
        common_telemetry::init_default_ut_logging();
        let record = test_record();
        let mock_client = MockFetchClient {
            record: record.clone(),
        };
        let index = RegionWalRange::new(0..30, 1024);
        let consumer = Consumer {
            last_high_watermark: -1,
            client: Arc::new(mock_client),
            max_batch_size: usize::MAX,
            max_wait_ms: 500,
            avg_record_size: record.approximate_size(),
            terminated: false,
            buffer: RecordsBuffer {
                buffer: VecDeque::new(),
                index,
            },
            fetch_fut: Fuse::terminated(),
            total_fetched_bytes: 0,
        };

        let records = consumer.try_collect::<Vec<_>>().await.unwrap();
        assert_eq!(
            records
                .into_iter()
                .map(|(x, _)| x.offset)
                .collect::<Vec<_>>(),
            (0..30).collect::<Vec<_>>()
        )
    }

    #[tokio::test]
    async fn test_consumer_range_boundaries() {
        for (start, end, max_batch_size) in [
            (0, 0, 1024),
            (5, 5, 1024),
            (6, 5, 1024),
            (5, 6, 1024),
            (5, 35, 128),
        ] {
            let record = test_record();
            let consumer = ConsumerBuilder::default()
                .client(Arc::new(MockFetchClient {
                    record: record.clone(),
                }))
                .buffer(RecordsBuffer::new(RegionWalRange::new(
                    start..end,
                    max_batch_size,
                )))
                .avg_record_size(record.approximate_size())
                .build()
                .unwrap();
            let offsets = consumer
                .try_collect::<Vec<_>>()
                .await
                .unwrap()
                .into_iter()
                .map(|(record, _)| record.offset as u64)
                .collect::<Vec<_>>();
            assert_eq!(offsets, (start..end).collect::<Vec<_>>());
        }
    }
}