log_store/kafka/

log_store.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.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::{Duration, Instant};

use common_base::readable_size::ReadableSize;
use common_meta::datanode::TopicStatsReporter;
use common_meta::distributed_time_constants::TOPIC_STATS_REPORT_INTERVAL_SECS;
use common_telemetry::{debug, info, warn};
use common_time::util::current_time_millis;
use common_wal::config::kafka::DatanodeKafkaConfig;
use dashmap::DashMap;
use futures::future::try_join_all;
use futures_util::StreamExt;
use rskafka::client::partition::OffsetAt;
use snafu::{OptionExt, ResultExt};
use store_api::logstore::entry::{
    Entry, Id as EntryId, MultiplePartEntry, MultiplePartHeader, NaiveEntry,
};
use store_api::logstore::provider::{KafkaProvider, Provider};
use store_api::logstore::{AppendBatchResponse, LogStore, SendableEntryStream};
use store_api::storage::RegionId;

use crate::error::{self, ConsumeRecordSnafu, Error, GetOffsetSnafu, InvalidProviderSnafu, Result};
use crate::kafka::client_manager::{ClientManager, ClientManagerRef};
use crate::kafka::consumer::{ConsumerBuilder, RecordsBuffer, RegionWalRange};
use crate::kafka::periodic_offset_fetcher::PeriodicOffsetFetcher;
use crate::kafka::producer::OrderedBatchProducerRef;
use crate::kafka::util::record::{
    ESTIMATED_META_SIZE, Record, convert_to_kafka_records, is_record_meta_key, maybe_emit_entry,
    remaining_entries,
};
use crate::metrics;

/// Statistics for a topic.
#[derive(Debug, Clone, Copy, Default)]
pub struct TopicStat {
    /// Latest offset for the topic.
    ///
    /// The latest offset is updated in two ways:
    /// - Automatically when the producer successfully commits data to Kafka
    /// - Periodically by the [PeriodicOffsetFetcher](crate::kafka::periodic_offset_fetcher::PeriodicOffsetFetcher).
    ///
    pub latest_offset: u64,
    /// Total size in bytes of records appended to the topic.
    pub record_size: u64,
    /// Total number of records appended to the topic.
    pub record_num: u64,
}

/// A log store backed by Kafka.
#[derive(Debug)]
pub struct KafkaLogStore {
    /// The manager of topic clients.
    client_manager: ClientManagerRef,
    /// The max size of a batch.
    max_batch_bytes: usize,
    /// The consumer wait timeout.
    consumer_wait_timeout: Duration,
    /// Ignore missing entries during read WAL.
    overwrite_entry_start_id: bool,
    /// The stats of each topic.
    ///
    /// This shared map allows multiple components to access the latest stats
    /// information without needing to query Kafka directly.
    topic_stats: Arc<DashMap<Arc<KafkaProvider>, TopicStat>>,
}

struct PeriodicTopicStatsReporter {
    topic_stats: Arc<DashMap<Arc<KafkaProvider>, TopicStat>>,
    last_reported_timestamp_millis: i64,
    report_interval_millis: i64,
}

impl PeriodicTopicStatsReporter {
    fn align_ts(ts: i64, report_interval_millis: i64) -> i64 {
        (ts / report_interval_millis) * report_interval_millis
    }

    /// Creates a new [PeriodicTopicStatsReporter].
    ///
    /// # Panics
    ///
    /// Panics if `report_interval` is zero.
    fn new(
        topic_stats: Arc<DashMap<Arc<KafkaProvider>, TopicStat>>,
        report_interval: Duration,
    ) -> Self {
        assert!(!report_interval.is_zero());
        let report_interval_millis = report_interval.as_millis() as i64;
        let last_reported_timestamp_millis =
            Self::align_ts(current_time_millis(), report_interval_millis);

        Self {
            topic_stats,
            last_reported_timestamp_millis,
            report_interval_millis,
        }
    }
}

impl TopicStatsReporter for PeriodicTopicStatsReporter {
    fn reportable_topics(&mut self) -> Vec<common_meta::datanode::TopicStat> {
        let now = Self::align_ts(current_time_millis(), self.report_interval_millis);
        if now < self.last_reported_timestamp_millis + self.report_interval_millis {
            debug!("Skip reporting topic stats because the interval is not reached");
            return vec![];
        }

        self.last_reported_timestamp_millis = now;
        let mut reportable_topics = Vec::with_capacity(self.topic_stats.len());
        for e in self.topic_stats.iter() {
            let topic_stat = e.value();
            let topic_stat = common_meta::datanode::TopicStat {
                topic: e.key().topic.clone(),
                latest_entry_id: topic_stat.latest_offset,
                record_size: topic_stat.record_size,
                record_num: topic_stat.record_num,
            };
            debug!("Reportable topic: {:?}", topic_stat);
            reportable_topics.push(topic_stat);
        }
        debug!("Reportable {} topics at {}", reportable_topics.len(), now);
        reportable_topics
    }
}

impl KafkaLogStore {
    /// Tries to create a Kafka log store.
    pub async fn try_new(config: &DatanodeKafkaConfig) -> Result<Self> {
        let topic_stats = Arc::new(DashMap::new());
        let client_manager = Arc::new(ClientManager::try_new(config, topic_stats.clone()).await?);
        let fetcher = PeriodicOffsetFetcher::new(
            config.topic_latest_offset_fetch_interval,
            client_manager.clone(),
        );
        fetcher.run().await;

        Ok(Self {
            client_manager,
            max_batch_bytes: config.max_batch_bytes.as_bytes() as usize,
            consumer_wait_timeout: config.consumer_wait_timeout,
            overwrite_entry_start_id: config.overwrite_entry_start_id,
            topic_stats,
        })
    }

    /// Returns the topic stats.
    pub fn topic_stats_reporter(&self) -> Box<dyn TopicStatsReporter> {
        Box::new(PeriodicTopicStatsReporter::new(
            self.topic_stats.clone(),
            Duration::from_secs(TOPIC_STATS_REPORT_INTERVAL_SECS),
        ))
    }
}

fn build_entry(
    data: Vec<u8>,
    entry_id: EntryId,
    region_id: RegionId,
    provider: &Provider,
    max_data_size: usize,
) -> Entry {
    if data.len() <= max_data_size {
        Entry::Naive(NaiveEntry {
            provider: provider.clone(),
            region_id,
            entry_id,
            data,
        })
    } else {
        let parts = data
            .chunks(max_data_size)
            .map(|s| s.into())
            .collect::<Vec<_>>();
        let num_parts = parts.len();

        let mut headers = Vec::with_capacity(num_parts);
        headers.push(MultiplePartHeader::First);
        headers.extend((1..num_parts - 1).map(MultiplePartHeader::Middle));
        headers.push(MultiplePartHeader::Last);

        Entry::MultiplePart(MultiplePartEntry {
            provider: provider.clone(),
            region_id,
            entry_id,
            headers,
            parts,
        })
    }
}

#[async_trait::async_trait]
impl LogStore for KafkaLogStore {
    type Error = Error;

    /// Creates an [Entry].
    fn entry(
        &self,
        data: Vec<u8>,
        entry_id: EntryId,
        region_id: RegionId,
        provider: &Provider,
    ) -> Result<Entry> {
        provider
            .as_kafka_provider()
            .with_context(|| InvalidProviderSnafu {
                expected: KafkaProvider::type_name(),
                actual: provider.type_name(),
            })?;

        let max_data_size = self.max_batch_bytes - ESTIMATED_META_SIZE;
        Ok(build_entry(
            data,
            entry_id,
            region_id,
            provider,
            max_data_size,
        ))
    }

    /// Appends a batch of entries and returns a response containing a map where the key is a region id
    /// while the value is the id of the last successfully written entry of the region.
    async fn append_batch(&self, entries: Vec<Entry>) -> Result<AppendBatchResponse> {
        metrics::METRIC_KAFKA_APPEND_BATCH_BYTES_TOTAL.inc_by(
            entries
                .iter()
                .map(|entry| entry.estimated_size())
                .sum::<usize>() as u64,
        );
        let _timer = metrics::METRIC_KAFKA_APPEND_BATCH_ELAPSED.start_timer();

        if entries.is_empty() {
            return Ok(AppendBatchResponse::default());
        }

        let region_ids = entries
            .iter()
            .map(|entry| entry.region_id())
            .collect::<HashSet<_>>();
        let mut region_grouped_records: HashMap<RegionId, (OrderedBatchProducerRef, Vec<_>)> =
            HashMap::with_capacity(region_ids.len());
        let mut region_to_provider = HashMap::with_capacity(region_ids.len());
        for entry in entries {
            let provider = entry.provider().as_kafka_provider().with_context(|| {
                error::InvalidProviderSnafu {
                    expected: KafkaProvider::type_name(),
                    actual: entry.provider().type_name(),
                }
            })?;
            region_to_provider.insert(entry.region_id(), provider.clone());
            let region_id = entry.region_id();
            match region_grouped_records.entry(region_id) {
                std::collections::hash_map::Entry::Occupied(mut slot) => {
                    slot.get_mut().1.extend(convert_to_kafka_records(entry)?);
                }
                std::collections::hash_map::Entry::Vacant(slot) => {
                    let producer = self
                        .client_manager
                        .get_or_insert(provider)
                        .await?
                        .producer()
                        .clone();

                    slot.insert((producer, convert_to_kafka_records(entry)?));
                }
            }
        }

        let mut region_grouped_result_receivers = Vec::with_capacity(region_ids.len());
        for (region_id, (producer, records)) in region_grouped_records {
            // Safety: `KafkaLogStore::entry` will ensure that the
            // `Record`'s `approximate_size` must be less or equal to `max_batch_bytes`.
            region_grouped_result_receivers.push((region_id, producer.produce(records).await?))
        }

        let region_grouped_max_offset =
            try_join_all(region_grouped_result_receivers.into_iter().map(
                |(region_id, receiver)| async move {
                    receiver.wait().await.map(|offset| (region_id, offset))
                },
            ))
            .await?;
        debug!(
            "Appended batch to Kafka, region_grouped_max_offset: {:?}",
            region_grouped_max_offset
        );

        Ok(AppendBatchResponse {
            last_entry_ids: region_grouped_max_offset.into_iter().collect(),
        })
    }

    /// Creates a new `EntryStream` to asynchronously generates `Entry` with entry ids.
    /// Returns entries belonging to `provider`, starting from `entry_id`.
    async fn read(
        &self,
        provider: &Provider,
        mut entry_id: EntryId,
    ) -> Result<SendableEntryStream<'static, Entry, Self::Error>> {
        let provider = provider
            .as_kafka_provider()
            .with_context(|| InvalidProviderSnafu {
                expected: KafkaProvider::type_name(),
                actual: provider.type_name(),
            })?;

        let _timer = metrics::METRIC_KAFKA_READ_ELAPSED.start_timer();

        // Gets the client associated with the topic.
        let client = self
            .client_manager
            .get_or_insert(provider)
            .await?
            .client()
            .clone();

        if self.overwrite_entry_start_id {
            let start_offset =
                client
                    .get_offset(OffsetAt::Earliest)
                    .await
                    .context(GetOffsetSnafu {
                        topic: &provider.topic,
                    })?;

            if entry_id as i64 <= start_offset {
                warn!(
                    "The entry_id: {} is less than start_offset: {}, topic: {}. Overwriting entry_id with start_offset",
                    entry_id, start_offset, &provider.topic
                );

                entry_id = start_offset as u64;
            }
        }

        // Gets the offset of the latest record in the topic. Actually, it's the latest record of the single partition in the topic.
        // The read operation terminates when this record is consumed.
        // Warning: the `get_offset` returns the end offset of the latest record. For our usage, it should be decremented.
        // See: https://kafka.apache.org/36/javadoc/org/apache/kafka/clients/consumer/KafkaConsumer.html#endOffsets(java.util.Collection)
        let end_offset = client
            .get_offset(OffsetAt::Latest)
            .await
            .context(GetOffsetSnafu {
                topic: &provider.topic,
            })?;
        let latest_offset = (end_offset as u64).saturating_sub(1);
        self.topic_stats
            .entry(provider.clone())
            .and_modify(|stat| {
                stat.latest_offset = stat.latest_offset.max(latest_offset);
            })
            .or_insert_with(|| TopicStat {
                latest_offset,
                record_size: 0,
                record_num: 0,
            });

        let range = entry_id..end_offset as u64;
        if range.is_empty() {
            warn!("No new entries in range {:?} of ns {}", range, provider);
            return Ok(futures_util::stream::empty().boxed());
        }

        let iterator = RegionWalRange::new(range, self.max_batch_bytes);
        debug!("Reading entries with {:?} of ns {}", iterator, provider);

        // Safety: Must be ok.
        let mut stream_consumer = ConsumerBuilder::default()
            .client(client)
            // Safety: checked before.
            .buffer(RecordsBuffer::new(iterator))
            .max_batch_size(self.max_batch_bytes)
            .max_wait_ms(self.consumer_wait_timeout.as_millis() as u32)
            .build()
            .unwrap();

        // A buffer is used to collect records to construct a complete entry.
        let mut entry_records: HashMap<RegionId, Vec<Record>> = HashMap::new();
        let provider = provider.clone();
        let stream = async_stream::stream!({
            let now = Instant::now();
            while let Some(consume_result) = stream_consumer.next().await {
                // Each next on the stream consumer produces a `RecordAndOffset` and a high watermark offset.
                // The `RecordAndOffset` contains the record data and its start offset.
                // The high watermark offset is the offset of the last record plus one.
                let (record_and_offset, high_watermark) =
                    consume_result.context(ConsumeRecordSnafu {
                        topic: &provider.topic,
                    })?;
                let (kafka_record, offset) = (record_and_offset.record, record_and_offset.offset);

                debug!(
                    "Read a record at offset {} for topic {}, high watermark: {}",
                    offset, provider.topic, high_watermark
                );

                // Ignores no-op records.
                if kafka_record.value.is_some() {
                    let key = kafka_record.key;
                    match Record::try_from_parts(key.as_deref(), kafka_record.value) {
                        Ok(record) => {
                            // Tries to construct an entry from records consumed so far.
                            if let Some(mut entry) =
                                maybe_emit_entry(&provider, record, &mut entry_records)?
                            {
                                // We don't rely on the EntryId generated by mito2.
                                // Instead, we use the offset return from Kafka as EntryId.
                                // Therefore, we MUST overwrite the EntryId with RecordOffset.
                                entry.set_entry_id(offset as u64);
                                yield Ok(vec![entry]);
                            }
                        }
                        Err(err) => {
                            // A key that cannot be decoded is skipped only when the log store
                            // never wrote it. A key that is a JSON object with a `version`
                            // field may be a `RecordMeta` written by an incompatible version,
                            // for example one with a newer `RecordType` variant read after a
                            // rollback or during a rolling upgrade; skipping it would lose the
                            // entry once the WAL is pruned past its offset, so it stays fatal.
                            if key.as_deref().is_some_and(is_record_meta_key) {
                                yield Err(err);
                                return;
                            }
                            warn!(
                                "Skipping unparsable record at offset {} of topic {}, reason: {}",
                                offset, provider.topic, err
                            );
                        }
                    }
                }

                // The termination check runs for every consumed record, including skipped ones.
                if check_termination(offset, end_offset) {
                    if let Some(entries) = remaining_entries(&provider, &mut entry_records) {
                        yield Ok(entries);
                    }
                    break;
                }
            }

            metrics::METRIC_KAFKA_READ_BYTES_TOTAL.inc_by(stream_consumer.total_fetched_bytes());

            info!(
                "Fetched {} bytes from topic: {}, start_entry_id: {}, end_offset: {}, elapsed: {:?}",
                ReadableSize(stream_consumer.total_fetched_bytes()),
                stream_consumer.topic(),
                entry_id,
                end_offset,
                now.elapsed()
            );
        });
        Ok(Box::pin(stream))
    }

    /// Creates a new `Namespace` from the given ref.
    async fn create_namespace(&self, _provider: &Provider) -> Result<()> {
        Ok(())
    }

    /// Deletes an existing `Namespace` specified by the given ref.
    async fn delete_namespace(&self, _provider: &Provider) -> Result<()> {
        Ok(())
    }

    /// Lists all existing namespaces.
    async fn list_namespaces(&self) -> Result<Vec<Provider>> {
        Ok(vec![])
    }

    /// Marks all entries with ids `<=entry_id` of the given `namespace` as obsolete,
    /// so that the log store can safely delete those entries. This method does not guarantee
    /// that the obsolete entries are deleted immediately.
    async fn obsolete(
        &self,
        _provider: &Provider,
        _region_id: RegionId,
        _entry_id: EntryId,
    ) -> Result<()> {
        // Kafka WAL pruning is coordinated by metasrv.
        Ok(())
    }

    async fn obsolete_all(&self, _provider: &Provider, _region_id: RegionId) -> Result<()> {
        Ok(())
    }

    /// Returns the highest entry id of the specified topic in remote WAL.
    fn latest_entry_id(&self, provider: &Provider) -> Result<EntryId> {
        let provider = provider
            .as_kafka_provider()
            .with_context(|| InvalidProviderSnafu {
                expected: KafkaProvider::type_name(),
                actual: provider.type_name(),
            })?;

        let stat = self
            .topic_stats
            .get(provider)
            .as_deref()
            .copied()
            .unwrap_or_default();

        Ok(stat.latest_offset)
    }

    /// Stops components of the logstore.
    async fn stop(&self) -> Result<()> {
        Ok(())
    }
}

fn check_termination(offset: i64, end_offset: i64) -> bool {
    // Terminates the stream if the entry with the end offset was read.
    if offset >= end_offset {
        debug!("Stream consumer terminates at offset {}", offset);
        // There must have no records when the stream terminates.
        true
    } else {
        false
    }
}

#[cfg(test)]
mod tests {

    use std::assert_matches;
    use std::collections::HashMap;
    use std::sync::Arc;
    use std::time::Duration;

    use common_base::readable_size::ReadableSize;
    use common_meta::datanode::TopicStatsReporter;
    use common_telemetry::info;
    use common_wal::config::kafka::DatanodeKafkaConfig;
    use common_wal::config::kafka::common::KafkaConnectionConfig;
    use common_wal::maybe_skip_kafka_integration_test;
    use common_wal::test_util::get_kafka_endpoints;
    use dashmap::DashMap;
    use futures::TryStreamExt;
    use rand::Rng;
    use rand::prelude::SliceRandom;
    use rskafka::client::partition::{Compression, OffsetAt};
    use rskafka::record::Record as KafkaRecord;
    use store_api::logstore::LogStore;
    use store_api::logstore::entry::{Entry, MultiplePartEntry, MultiplePartHeader, NaiveEntry};
    use store_api::logstore::provider::Provider;
    use store_api::storage::RegionId;

    use super::build_entry;
    use crate::error;
    use crate::kafka::log_store::{KafkaLogStore, PeriodicTopicStatsReporter, TopicStat};
    use crate::kafka::util::record::{ESTIMATED_META_SIZE, convert_to_kafka_records};

    #[test]
    fn test_build_naive_entry() {
        let provider = Provider::kafka_provider("my_topic".to_string());
        let region_id = RegionId::new(1, 1);
        let entry = build_entry(vec![1; 100], 1, region_id, &provider, 120);

        assert_eq!(
            entry.into_naive_entry().unwrap(),
            NaiveEntry {
                provider,
                region_id,
                entry_id: 1,
                data: vec![1; 100]
            }
        )
    }

    #[test]
    fn test_build_into_multiple_part_entry() {
        let provider = Provider::kafka_provider("my_topic".to_string());
        let region_id = RegionId::new(1, 1);
        let entry = build_entry(vec![1; 100], 1, region_id, &provider, 50);

        assert_eq!(
            entry.into_multiple_part_entry().unwrap(),
            MultiplePartEntry {
                provider: provider.clone(),
                region_id,
                entry_id: 1,
                headers: vec![MultiplePartHeader::First, MultiplePartHeader::Last],
                parts: vec![vec![1; 50], vec![1; 50]],
            }
        );

        let region_id = RegionId::new(1, 1);
        let entry = build_entry(vec![1; 100], 1, region_id, &provider, 21);

        assert_eq!(
            entry.into_multiple_part_entry().unwrap(),
            MultiplePartEntry {
                provider,
                region_id,
                entry_id: 1,
                headers: vec![
                    MultiplePartHeader::First,
                    MultiplePartHeader::Middle(1),
                    MultiplePartHeader::Middle(2),
                    MultiplePartHeader::Middle(3),
                    MultiplePartHeader::Last
                ],
                parts: vec![
                    vec![1; 21],
                    vec![1; 21],
                    vec![1; 21],
                    vec![1; 21],
                    vec![1; 16]
                ],
            }
        )
    }

    fn generate_entries(
        logstore: &KafkaLogStore,
        provider: &Provider,
        num_entries: usize,
        region_id: RegionId,
        data_len: usize,
    ) -> Vec<Entry> {
        (0..num_entries)
            .map(|_| {
                let data: Vec<u8> = (0..data_len).map(|_| rand::random::<u8>()).collect();
                // Always set `entry_id` to 0, the real entry_id will be set during the read.
                logstore.entry(data, 0, region_id, provider).unwrap()
            })
            .collect()
    }

    async fn prepare_topic(logstore: &KafkaLogStore, topic_name: &str) {
        let controller_client = logstore.client_manager.controller_client();
        controller_client
            .create_topic(topic_name.to_string(), 1, 1, 5000)
            .await
            .unwrap();
    }

    async fn new_test_logstore(max_batch_bytes: ReadableSize) -> KafkaLogStore {
        let config = DatanodeKafkaConfig {
            connection: KafkaConnectionConfig {
                broker_endpoints: get_kafka_endpoints(),
                ..Default::default()
            },
            max_batch_bytes,
            ..Default::default()
        };
        KafkaLogStore::try_new(&config).await.unwrap()
    }

    async fn produce_records(
        logstore: &KafkaLogStore,
        provider: &Provider,
        records: Vec<KafkaRecord>,
    ) {
        // A freshly created topic can be missing from the broker's metadata for a short
        // while; `try_create_client` already retries that case.
        let client = logstore
            .client_manager
            .get_or_insert(provider.as_kafka_provider().unwrap())
            .await
            .unwrap()
            .client()
            .clone();
        client
            .produce(records, Compression::NoCompression)
            .await
            .unwrap();
    }

    /// Builds a record that was not written by this log store.
    fn foreign_record(key: Option<&[u8]>) -> KafkaRecord {
        KafkaRecord {
            key: key.map(|key| key.to_vec()),
            value: Some(b"foreign payload".to_vec()),
            timestamp: chrono::Utc::now(),
            headers: Default::default(),
        }
    }

    /// Returns the log end offset of the topic, i.e. the exclusive end of the range that
    /// `read` consumes.
    async fn log_end_offset(logstore: &KafkaLogStore, provider: &Provider) -> i64 {
        let client = logstore
            .client_manager
            .get_or_insert(provider.as_kafka_provider().unwrap())
            .await
            .unwrap();
        client.client().get_offset(OffsetAt::Latest).await.unwrap()
    }

    /// Returns one record per kind of key that the log store never writes: a missing key, a
    /// key that is not JSON, a JSON value that is not an object, and a JSON object without a
    /// `version` field.
    fn foreign_records() -> Vec<KafkaRecord> {
        vec![
            foreign_record(None),
            foreign_record(Some(b"not-a-json-key")),
            foreign_record(Some(b"[1, 2, 3]")),
            foreign_record(Some(br#"{"hello": "world"}"#)),
        ]
    }

    #[tokio::test]
    async fn test_append_batch_basic() {
        common_telemetry::init_default_ut_logging();
        maybe_skip_kafka_integration_test!();
        let logstore = new_test_logstore(ReadableSize::kb(32)).await;
        let topic_name = uuid::Uuid::new_v4().to_string();
        prepare_topic(&logstore, &topic_name).await;
        let provider = Provider::kafka_provider(topic_name);

        let region_entries = (0..5)
            .map(|i| {
                let region_id = RegionId::new(1, i);
                (
                    region_id,
                    generate_entries(&logstore, &provider, 20, region_id, 1024),
                )
            })
            .collect::<HashMap<RegionId, Vec<_>>>();

        let mut all_entries = region_entries
            .values()
            .flatten()
            .cloned()
            .collect::<Vec<_>>();
        all_entries.shuffle(&mut rand::rng());

        let response = logstore.append_batch(all_entries.clone()).await.unwrap();
        // 5 region
        assert_eq!(response.last_entry_ids.len(), 5);
        let got_entries = logstore
            .read(&provider, 0)
            .await
            .unwrap()
            .try_collect::<Vec<_>>()
            .await
            .unwrap()
            .into_iter()
            .flatten()
            .collect::<Vec<_>>();
        for (region_id, _) in region_entries {
            let expected_entries = all_entries
                .iter()
                .filter(|entry| entry.region_id() == region_id)
                .cloned()
                .collect::<Vec<_>>();
            let mut actual_entries = got_entries
                .iter()
                .filter(|entry| entry.region_id() == region_id)
                .cloned()
                .collect::<Vec<_>>();
            actual_entries
                .iter_mut()
                .for_each(|entry| entry.set_entry_id(0));
            assert_eq!(expected_entries, actual_entries);
        }
        let latest_entry_id = logstore.latest_entry_id(&provider).unwrap();
        let client = logstore
            .client_manager
            .get_or_insert(provider.as_kafka_provider().unwrap())
            .await
            .unwrap();
        assert_eq!(latest_entry_id, 99);
        // The latest offset is the offset of the last record plus one.
        let latest = client.client().get_offset(OffsetAt::Latest).await.unwrap();
        assert_eq!(latest, 100);
    }

    #[tokio::test]
    async fn test_append_batch_basic_large() {
        common_telemetry::init_default_ut_logging();
        maybe_skip_kafka_integration_test!();
        let data_size_kb = rand::rng().random_range(9..31usize);
        info!("Entry size: {}Ki", data_size_kb);
        let logstore = new_test_logstore(ReadableSize::kb(8)).await;
        let topic_name = uuid::Uuid::new_v4().to_string();
        prepare_topic(&logstore, &topic_name).await;
        let provider = Provider::kafka_provider(topic_name);
        let region_entries = (0..5)
            .map(|i| {
                let region_id = RegionId::new(1, i);
                (
                    region_id,
                    generate_entries(&logstore, &provider, 20, region_id, data_size_kb * 1024),
                )
            })
            .collect::<HashMap<RegionId, Vec<_>>>();

        let mut all_entries = region_entries
            .values()
            .flatten()
            .cloned()
            .collect::<Vec<_>>();
        assert_matches!(all_entries[0], Entry::MultiplePart(_));
        all_entries.shuffle(&mut rand::rng());

        let response = logstore.append_batch(all_entries.clone()).await.unwrap();
        // 5 region
        assert_eq!(response.last_entry_ids.len(), 5);
        let got_entries = logstore
            .read(&provider, 0)
            .await
            .unwrap()
            .try_collect::<Vec<_>>()
            .await
            .unwrap()
            .into_iter()
            .flatten()
            .collect::<Vec<_>>();
        for (region_id, _) in region_entries {
            let expected_entries = all_entries
                .iter()
                .filter(|entry| entry.region_id() == region_id)
                .cloned()
                .collect::<Vec<_>>();
            let mut actual_entries = got_entries
                .iter()
                .filter(|entry| entry.region_id() == region_id)
                .cloned()
                .collect::<Vec<_>>();
            actual_entries
                .iter_mut()
                .for_each(|entry| entry.set_entry_id(0));
            assert_eq!(expected_entries, actual_entries);
        }
        let high_wathermark = logstore.latest_entry_id(&provider).unwrap();
        assert_eq!(high_wathermark, (data_size_kb as u64 / 8 + 1) * 20 * 5 - 1);
    }

    #[tokio::test]
    async fn test_read_skips_unparsable_record_inside_multipart_entry() {
        maybe_skip_kafka_integration_test!();
        let logstore = new_test_logstore(ReadableSize::kb(8)).await;
        let topic_name = uuid::Uuid::new_v4().to_string();
        prepare_topic(&logstore, &topic_name).await;
        let provider = Provider::kafka_provider(topic_name);

        let region_id = RegionId::new(1, 1);
        // Splits the entry into First, Middle(1) and Last.
        let max_data_size = 8 * 1024 - ESTIMATED_META_SIZE;
        let data = vec![7u8; max_data_size * 2 + 10];
        let entry = logstore
            .entry(data.clone(), 0, region_id, &provider)
            .unwrap();
        assert_matches!(entry, Entry::MultiplePart(_));
        let mut records = convert_to_kafka_records(entry).unwrap();
        assert_eq!(records.len(), 3);

        let last = records.pop().unwrap();
        // Places an unparsable record before the entry and between every pair of parts.
        produce_records(
            &logstore,
            &provider,
            vec![foreign_record(Some(b"not-a-json-key"))],
        )
        .await;
        produce_records(&logstore, &provider, vec![records[0].clone()]).await;
        produce_records(
            &logstore,
            &provider,
            vec![foreign_record(Some(br#"{"hello": "world"}"#))],
        )
        .await;
        produce_records(&logstore, &provider, vec![records[1].clone()]).await;
        produce_records(
            &logstore,
            &provider,
            vec![foreign_record(Some(b"not-a-json-key"))],
        )
        .await;
        produce_records(&logstore, &provider, vec![last]).await;

        let entries = logstore
            .read(&provider, 0)
            .await
            .unwrap()
            .try_collect::<Vec<_>>()
            .await
            .unwrap()
            .into_iter()
            .flatten()
            .collect::<Vec<_>>();
        assert_eq!(entries.len(), 1);
        let entry = entries
            .into_iter()
            .next()
            .unwrap()
            .into_multiple_part_entry()
            .unwrap();
        assert_eq!(
            entry.headers,
            vec![
                MultiplePartHeader::First,
                MultiplePartHeader::Middle(1),
                MultiplePartHeader::Last
            ]
        );
        assert_eq!(entry.parts.concat(), data);
    }

    #[tokio::test]
    async fn test_read_skips_unparsable_record_across_regions() {
        maybe_skip_kafka_integration_test!();
        let logstore = new_test_logstore(ReadableSize::kb(8)).await;
        let topic_name = uuid::Uuid::new_v4().to_string();
        prepare_topic(&logstore, &topic_name).await;
        let provider = Provider::kafka_provider(topic_name);

        let region_a = RegionId::new(1, 1);
        let region_b = RegionId::new(1, 2);
        let max_data_size = 8 * 1024 - ESTIMATED_META_SIZE;
        let data_a = vec![1u8; max_data_size * 2 + 10];
        let data_b = vec![2u8; max_data_size * 2 + 10];
        let records_a = convert_to_kafka_records(
            logstore
                .entry(data_a.clone(), 0, region_a, &provider)
                .unwrap(),
        )
        .unwrap();
        let records_b = convert_to_kafka_records(
            logstore
                .entry(data_b.clone(), 0, region_b, &provider)
                .unwrap(),
        )
        .unwrap();
        assert_eq!(records_a.len(), 3);
        assert_eq!(records_b.len(), 3);

        // Interleaves the parts of both entries with unparsable records.
        for i in 0..3 {
            produce_records(&logstore, &provider, vec![records_a[i].clone()]).await;
            produce_records(&logstore, &provider, vec![records_b[i].clone()]).await;
            produce_records(
                &logstore,
                &provider,
                vec![foreign_record(Some(b"not-a-json-key"))],
            )
            .await;
        }

        let entries = logstore
            .read(&provider, 0)
            .await
            .unwrap()
            .try_collect::<Vec<_>>()
            .await
            .unwrap()
            .into_iter()
            .flatten()
            .collect::<Vec<_>>();
        assert_eq!(entries.len(), 2);

        let mut entries = entries
            .into_iter()
            .map(|entry| (entry.region_id(), entry.into_multiple_part_entry().unwrap()))
            .collect::<HashMap<_, _>>();
        assert_eq!(entries.remove(&region_a).unwrap().parts.concat(), data_a);
        assert_eq!(entries.remove(&region_b).unwrap().parts.concat(), data_b);
        assert!(entries.is_empty());
    }

    #[tokio::test]
    async fn test_read_finishes_when_last_record_is_unparsable() {
        maybe_skip_kafka_integration_test!();
        let logstore = new_test_logstore(ReadableSize::kb(32)).await;
        let topic_name = uuid::Uuid::new_v4().to_string();
        prepare_topic(&logstore, &topic_name).await;
        let provider = Provider::kafka_provider(topic_name);

        let region_id = RegionId::new(1, 1);
        let mut expected = generate_entries(&logstore, &provider, 3, region_id, 128);
        logstore.append_batch(expected.clone()).await.unwrap();
        // Appends an unparsable record as the last record of the read range.
        produce_records(
            &logstore,
            &provider,
            vec![foreign_record(Some(b"not-a-json-key"))],
        )
        .await;

        let entries = tokio::time::timeout(
            Duration::from_secs(30),
            logstore
                .read(&provider, 0)
                .await
                .unwrap()
                .try_collect::<Vec<_>>(),
        )
        .await
        .expect("read should finish after skipping the unparsable record")
        .unwrap()
        .into_iter()
        .flatten()
        .collect::<Vec<_>>();

        expected.iter_mut().for_each(|entry| entry.set_entry_id(0));
        let mut entries = entries;
        entries.iter_mut().for_each(|entry| entry.set_entry_id(0));
        assert_eq!(entries, expected);
    }

    #[tokio::test]
    async fn test_read_skips_foreign_records_of_every_kind() {
        maybe_skip_kafka_integration_test!();
        let logstore = new_test_logstore(ReadableSize::kb(32)).await;
        let topic_name = uuid::Uuid::new_v4().to_string();
        prepare_topic(&logstore, &topic_name).await;
        let provider = Provider::kafka_provider(topic_name);

        // Places every kind of unparsable record before the first entry.
        for record in foreign_records() {
            produce_records(&logstore, &provider, vec![record]).await;
        }

        let region_id = RegionId::new(1, 1);
        let mut expected = generate_entries(&logstore, &provider, 1, region_id, 128);
        logstore.append_batch(expected.clone()).await.unwrap();

        let entries = logstore
            .read(&provider, 0)
            .await
            .unwrap()
            .try_collect::<Vec<_>>()
            .await
            .unwrap()
            .into_iter()
            .flatten()
            .collect::<Vec<_>>();

        expected.iter_mut().for_each(|entry| entry.set_entry_id(0));
        let mut entries = entries;
        entries.iter_mut().for_each(|entry| entry.set_entry_id(0));
        assert_eq!(entries, expected);
    }

    #[tokio::test]
    async fn test_read_finishes_when_topic_only_has_foreign_records() {
        maybe_skip_kafka_integration_test!();
        let logstore = new_test_logstore(ReadableSize::kb(32)).await;
        let topic_name = uuid::Uuid::new_v4().to_string();
        prepare_topic(&logstore, &topic_name).await;
        let provider = Provider::kafka_provider(topic_name);

        for record in foreign_records() {
            produce_records(&logstore, &provider, vec![record]).await;
        }

        // Asserts the read range covers the produced records; an empty range would let `read`
        // return without exercising the skip path.
        assert_eq!(log_end_offset(&logstore, &provider).await, 4);

        let entries = tokio::time::timeout(
            Duration::from_secs(30),
            logstore
                .read(&provider, 0)
                .await
                .unwrap()
                .try_collect::<Vec<_>>(),
        )
        .await
        .expect("read should finish after skipping all unparsable records")
        .unwrap();

        assert!(entries.into_iter().flatten().next().is_none());
    }

    #[tokio::test]
    async fn test_read_fails_on_incompatible_record_meta() {
        maybe_skip_kafka_integration_test!();
        let logstore = new_test_logstore(ReadableSize::kb(32)).await;

        // Keys that look like a `RecordMeta` but cannot be decoded by this version. They may
        // come from an incompatible version of the log store, so the read must not skip them.
        let keys: [&[u8]; 2] = [
            // A `RecordType` variant this version does not know.
            br#"{"version": 0, "tp": "Compressed", "entry_id": 0, "ns": {"region_id": 1, "topic": ""}}"#,
            // A field with the wrong type.
            br#"{"version": "0", "tp": "Full", "entry_id": 0, "ns": {"region_id": 1, "topic": ""}}"#,
        ];
        for key in keys {
            let topic_name = uuid::Uuid::new_v4().to_string();
            prepare_topic(&logstore, &topic_name).await;
            let provider = Provider::kafka_provider(topic_name);
            produce_records(&logstore, &provider, vec![foreign_record(Some(key))]).await;

            let err = logstore
                .read(&provider, 0)
                .await
                .unwrap()
                .try_collect::<Vec<Vec<Entry>>>()
                .await
                .unwrap_err();
            assert_matches!(err, error::Error::DecodeJson { .. });
        }
    }

    #[tokio::test]
    async fn test_read_fails_on_missing_middle_part() {
        maybe_skip_kafka_integration_test!();
        let logstore = new_test_logstore(ReadableSize::kb(8)).await;
        let topic_name = uuid::Uuid::new_v4().to_string();
        prepare_topic(&logstore, &topic_name).await;
        let provider = Provider::kafka_provider(topic_name);

        let region_id = RegionId::new(1, 1);
        // Splits the entry into First, Middle(1), Middle(2) and Last.
        let max_data_size = 8 * 1024 - ESTIMATED_META_SIZE;
        let data = vec![7u8; max_data_size * 3 + 10];
        let entry = logstore.entry(data, 0, region_id, &provider).unwrap();
        let records = convert_to_kafka_records(entry).unwrap();
        assert_eq!(records.len(), 4);

        // Only a missing part of an entry we wrote is fatal; a foreign record is skipped
        // because it cannot be one of our parts.
        produce_records(
            &logstore,
            &provider,
            vec![foreign_record(Some(b"not-a-json-key"))],
        )
        .await;
        produce_records(&logstore, &provider, vec![records[0].clone()]).await;
        produce_records(&logstore, &provider, vec![records[2].clone()]).await;

        let err = logstore
            .read(&provider, 0)
            .await
            .unwrap()
            .try_collect::<Vec<Vec<Entry>>>()
            .await
            .unwrap_err();
        assert_matches!(err, error::Error::IllegalSequence { .. });
    }

    #[tokio::test]
    async fn test_read_succeeds_on_two_part_entry() {
        maybe_skip_kafka_integration_test!();
        let logstore = new_test_logstore(ReadableSize::kb(8)).await;
        let topic_name = uuid::Uuid::new_v4().to_string();
        prepare_topic(&logstore, &topic_name).await;
        let provider = Provider::kafka_provider(topic_name);

        let region_id = RegionId::new(1, 1);
        // Splits the entry into First and Last, without any middle part.
        let max_data_size = 8 * 1024 - ESTIMATED_META_SIZE;
        let data = vec![7u8; max_data_size + 10];
        let entry = logstore
            .entry(data.clone(), 0, region_id, &provider)
            .unwrap();
        let records = convert_to_kafka_records(entry).unwrap();
        assert_eq!(records.len(), 2);

        for record in records {
            produce_records(&logstore, &provider, vec![record]).await;
        }

        let entries = logstore
            .read(&provider, 0)
            .await
            .unwrap()
            .try_collect::<Vec<_>>()
            .await
            .unwrap()
            .into_iter()
            .flatten()
            .collect::<Vec<_>>();
        assert_eq!(entries.len(), 1);
        let entry = entries
            .into_iter()
            .next()
            .unwrap()
            .into_multiple_part_entry()
            .unwrap();
        assert_eq!(
            entry.headers,
            vec![MultiplePartHeader::First, MultiplePartHeader::Last]
        );
        assert_eq!(entry.parts.concat(), data);
    }

    #[tokio::test]
    async fn test_topic_stats_reporter() {
        common_telemetry::init_default_ut_logging();
        let topic_stats = Arc::new(DashMap::new());
        let provider = Provider::kafka_provider("my_topic".to_string());
        topic_stats.insert(
            provider.as_kafka_provider().unwrap().clone(),
            TopicStat {
                latest_offset: 0,
                record_size: 0,
                record_num: 0,
            },
        );
        let mut reporter = PeriodicTopicStatsReporter::new(topic_stats, Duration::from_secs(1));
        // The first reportable topics should be empty.
        let reportable_topics = reporter.reportable_topics();
        assert_eq!(reportable_topics.len(), 0);

        // After 1 second, the reportable topics should be the topic in the topic_stats.
        tokio::time::sleep(Duration::from_secs(1)).await;
        let reportable_topics = reporter.reportable_topics();
        assert_eq!(reportable_topics.len(), 1);

        // Call it immediately, should be empty.
        let reportable_topics = reporter.reportable_topics();
        assert_eq!(reportable_topics.len(), 0);
    }
}