log_store/kafka/
periodic_offset_fetcher.rs

1// Copyright 2023 Greptime Team
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::time::Duration;
16
17use common_telemetry::{debug, error, info};
18use tokio::time::{MissedTickBehavior, interval};
19
20use crate::error::Result;
21use crate::kafka::client_manager::ClientManagerRef;
22
23/// PeriodicOffsetFetcher is responsible for periodically updating the offset
24/// for each Kafka topic.
25pub(crate) struct PeriodicOffsetFetcher {
26    /// Interval to fetch the latest offset.
27    interval: Duration,
28    /// Client manager to send requests.
29    client_manager: ClientManagerRef,
30}
31
32impl PeriodicOffsetFetcher {
33    pub(crate) fn new(interval: Duration, client_manager: ClientManagerRef) -> Self {
34        Self {
35            interval,
36            client_manager,
37        }
38    }
39
40    /// Starts the offset fetcher as a background task
41    ///
42    /// This spawns a task that periodically queries Kafka for the latest
43    /// offset values for all registered topics and updates the shared map.
44    pub(crate) async fn run(self) {
45        common_runtime::spawn_global(async move {
46            info!("PeriodicOffsetFetcher started");
47            let mut interval = interval(self.interval);
48            interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
49            loop {
50                interval.tick().await;
51                if let Err(e) = self.try_update().await {
52                    error!(e; "Failed to update latest offset");
53                }
54            }
55        });
56    }
57
58    /// Tries to refresh the latest offset for every registered topic.
59    ///
60    /// For each topic in the stats map, retrieves its producer and requests
61    /// an update of the latest offset from Kafka.
62    pub(crate) async fn try_update(&self) -> Result<()> {
63        let topics = self.client_manager.list_topics().await;
64        for topic in topics.iter() {
65            debug!("Fetching latest offset for topic: {}", topic.topic);
66            let producer = self
67                .client_manager
68                .get_or_insert(topic)
69                .await?
70                .producer()
71                .clone();
72            producer.fetch_latest_offset().await?;
73        }
74        Ok(())
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use std::sync::Arc;
81
82    use common_wal::maybe_skip_kafka_integration_test;
83    use common_wal::test_util::get_kafka_endpoints;
84    use store_api::logstore::provider::KafkaProvider;
85
86    use super::*;
87    use crate::kafka::test_util::{prepare, record};
88
89    #[tokio::test]
90    async fn test_try_update_latest_offset() {
91        common_telemetry::init_default_ut_logging();
92        maybe_skip_kafka_integration_test!();
93        let broker_endpoints = get_kafka_endpoints();
94
95        let (manager, topics) = prepare("test_try_update_latest_offset", 1, broker_endpoints).await;
96        let manager = Arc::new(manager);
97        let fetcher = PeriodicOffsetFetcher::new(Duration::from_millis(100), manager.clone());
98        let topic_stats = manager.topic_stats().clone();
99        fetcher.run().await;
100
101        let topic = topics[0].clone();
102        let provider = Arc::new(KafkaProvider::new(topic.clone()));
103        let producer = manager
104            .get_or_insert(&provider)
105            .await
106            .unwrap()
107            .producer()
108            .clone();
109
110        tokio::time::sleep(Duration::from_millis(150)).await;
111        let current_latest_offset = topic_stats.get(&provider).unwrap().latest_offset;
112        assert_eq!(current_latest_offset, 0);
113
114        let record = vec![record(), record()];
115        producer.produce(record.clone()).await.unwrap();
116        tokio::time::sleep(Duration::from_millis(150)).await;
117        let current_latest_offset = topic_stats.get(&provider).unwrap().latest_offset;
118        assert_eq!(current_latest_offset, record.len() as u64 - 1);
119    }
120}
log_store/kafka/periodic_offset_fetcher.rs

log_store/kafka/
periodic_offset_fetcher.rs
