meta_srv/procedure/wal_prune/
utils.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::collections::HashMap;
16use std::sync::Arc;
17use std::time::Duration;
18
19use common_meta::key::TableMetadataManagerRef;
20use common_meta::region_registry::LeaderRegionRegistryRef;
21use common_telemetry::warn;
22use itertools::{Itertools, MinMaxResult};
23use rskafka::client::Client;
24use rskafka::client::partition::{OffsetAt, PartitionClient, UnknownTopicHandling};
25use snafu::ResultExt;
26use store_api::logstore::EntryId;
27use store_api::storage::RegionId;
28
29use crate::error::{
30    BuildPartitionClientSnafu, DeleteRecordsSnafu, GetOffsetSnafu, Result,
31    TableMetadataManagerSnafu, UpdateTopicNameValueSnafu,
32};
33
34/// The default timeout for deleting records.
35const DELETE_RECORDS_TIMEOUT: Duration = Duration::from_secs(5);
36/// The default partition.
37const DEFAULT_PARTITION: i32 = 0;
38
39fn missing_region_ids(
40    all_region_ids: &[RegionId],
41    result_set: &HashMap<RegionId, u64>,
42) -> Vec<RegionId> {
43    let mut missing_region_ids = Vec::new();
44    for region_id in all_region_ids {
45        if !result_set.contains_key(region_id) {
46            missing_region_ids.push(*region_id);
47        }
48    }
49    missing_region_ids
50}
51
52/// Finds the prunable entry id for the topic.
53///
54/// Returns `None` if:
55/// - The topic has no region.
56/// - Some region info is missing from heartbeat.
57pub(crate) async fn find_pruneable_entry_id_for_topic(
58    table_metadata_manager: &TableMetadataManagerRef,
59    leader_region_registry: &LeaderRegionRegistryRef,
60    topic: &str,
61) -> Result<Option<u64>> {
62    let region_ids = table_metadata_manager
63        .topic_region_manager()
64        .regions(topic)
65        .await
66        .context(TableMetadataManagerSnafu)?
67        .into_keys()
68        .collect::<Vec<_>>();
69    if region_ids.is_empty() {
70        return Ok(None);
71    }
72
73    // Get the prunable entry id for each region.
74    let prunable_entry_ids_map = leader_region_registry
75        .batch_get(region_ids.iter().cloned())
76        .into_iter()
77        .map(|(region_id, region)| {
78            let prunable_entry_id = region.manifest.prunable_entry_id();
79            (region_id, prunable_entry_id)
80        })
81        .collect();
82    let missing_region_ids = missing_region_ids(&region_ids, &prunable_entry_ids_map);
83    if !missing_region_ids.is_empty() {
84        warn!(
85            "Cannot determine prunable entry id: missing region info from heartbeat. Topic: {}, missing region ids: {:?}",
86            topic, missing_region_ids
87        );
88        return Ok(None);
89    }
90
91    let min_max_result = prunable_entry_ids_map.values().minmax();
92    match min_max_result {
93        MinMaxResult::NoElements => Ok(None),
94        MinMaxResult::OneElement(prunable_entry_id) => Ok(Some(*prunable_entry_id)),
95        MinMaxResult::MinMax(min_prunable_entry_id, _) => Ok(Some(*min_prunable_entry_id)),
96    }
97}
98
99/// Determines whether pruning should be triggered based on the current pruned entry id and the prunable entry id.
100/// Returns true if:
101/// - There is no current pruned entry id (i.e., pruning has never occurred).
102/// - The current pruned entry id is greater than the prunable entry id (i.e., there is something to prune).
103pub(crate) fn should_trigger_prune(current: Option<u64>, prunable_entry_id: u64) -> bool {
104    match current {
105        None => true, // No pruning has occurred yet, should trigger immediately.
106        Some(current) => prunable_entry_id > current,
107    }
108}
109
110/// Returns a partition client for the given topic.
111pub(crate) async fn get_partition_client(
112    client: &Arc<Client>,
113    topic: &str,
114) -> Result<PartitionClient> {
115    client
116        .partition_client(topic, DEFAULT_PARTITION, UnknownTopicHandling::Retry)
117        .await
118        .context(BuildPartitionClientSnafu {
119            topic,
120            partition: DEFAULT_PARTITION,
121        })
122}
123
124/// Returns the earliest and latest offsets for the given topic.
125pub(crate) async fn get_offsets_for_topic(
126    partition_client: &PartitionClient,
127    topic: &str,
128) -> Result<(u64, u64)> {
129    let earliest_offset = partition_client
130        .get_offset(OffsetAt::Earliest)
131        .await
132        .context(GetOffsetSnafu { topic })?;
133    let latest_offset = partition_client
134        .get_offset(OffsetAt::Latest)
135        .await
136        .context(GetOffsetSnafu { topic })?;
137
138    Ok((earliest_offset as u64, latest_offset as u64))
139}
140
141/// Updates the pruned entry id for the given topic.
142pub(crate) async fn update_pruned_entry_id(
143    table_metadata_manager: &TableMetadataManagerRef,
144    topic: &str,
145    pruned_entry_id: EntryId,
146) -> Result<()> {
147    let prev = table_metadata_manager
148        .topic_name_manager()
149        .get(topic)
150        .await
151        .context(TableMetadataManagerSnafu)?;
152
153    table_metadata_manager
154        .topic_name_manager()
155        .update(topic, pruned_entry_id, prev)
156        .await
157        .context(UpdateTopicNameValueSnafu { topic })?;
158
159    Ok(())
160}
161
162/// Deletes the records for the given topic.
163pub(crate) async fn delete_records(
164    partition_client: &PartitionClient,
165    topic: &str,
166    pruned_entry_id: u64,
167) -> Result<()> {
168    partition_client
169        .delete_records(
170            // Note: here no "+1" is needed because the offset arg is exclusive,
171            // and it's defensive programming just in case somewhere else have a off by one error,
172            // see https://kafka.apache.org/36/javadoc/org/apache/kafka/clients/consumer/KafkaConsumer.html#endOffsets(java.util.Collection)
173            // which we use to get the end offset from high watermark
174            pruned_entry_id as i64,
175            DELETE_RECORDS_TIMEOUT.as_millis() as i32,
176        )
177        .await
178        .context(DeleteRecordsSnafu {
179            topic,
180            partition: DEFAULT_PARTITION,
181            offset: pruned_entry_id,
182        })
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn test_should_trigger_prune_none_current() {
191        // No pruning has occurred yet, should trigger
192        assert!(should_trigger_prune(None, 10));
193        assert!(should_trigger_prune(None, 0));
194    }
195
196    #[test]
197    fn test_should_trigger_prune_prunable_greater_than_current() {
198        // Prunable entry id is greater than current, should trigger
199        assert!(should_trigger_prune(Some(5), 6));
200        assert!(should_trigger_prune(Some(0), 1));
201        assert!(should_trigger_prune(Some(99), 100));
202    }
203
204    #[test]
205    fn test_should_not_trigger_prune_prunable_equal_to_current() {
206        // Prunable entry id is equal to current, should not trigger
207        assert!(!should_trigger_prune(Some(10), 10));
208        assert!(!should_trigger_prune(Some(0), 0));
209    }
210
211    #[test]
212    fn test_should_not_trigger_prune_prunable_less_than_current() {
213        // Prunable entry id is less than current, should not trigger
214        assert!(!should_trigger_prune(Some(10), 9));
215        assert!(!should_trigger_prune(Some(100), 99));
216    }
217
218    #[test]
219    fn test_missing_region_ids_none_missing() {
220        let all_region_ids = vec![RegionId::new(1, 1), RegionId::new(2, 2)];
221        let mut result_set = HashMap::new();
222        result_set.insert(RegionId::new(1, 1), 10);
223        result_set.insert(RegionId::new(2, 2), 20);
224        let missing = missing_region_ids(&all_region_ids, &result_set);
225        assert!(missing.is_empty());
226    }
227
228    #[test]
229    fn test_missing_region_ids_some_missing() {
230        let all_region_ids = vec![
231            RegionId::new(1, 1),
232            RegionId::new(2, 2),
233            RegionId::new(3, 3),
234        ];
235        let mut result_set = HashMap::new();
236        result_set.insert(RegionId::new(1, 1), 10);
237        let missing = missing_region_ids(&all_region_ids, &result_set);
238        assert_eq!(missing, vec![RegionId::new(2, 2), RegionId::new(3, 3)]);
239    }
240
241    #[test]
242    fn test_missing_region_ids_all_missing() {
243        let all_region_ids = vec![RegionId::new(1, 1), RegionId::new(2, 2)];
244        let result_set = HashMap::new();
245        let missing = missing_region_ids(&all_region_ids, &result_set);
246        assert_eq!(missing, all_region_ids);
247    }
248
249    #[test]
250    fn test_missing_region_ids_empty_all() {
251        let all_region_ids: Vec<RegionId> = vec![];
252        let mut result_set = HashMap::new();
253        result_set.insert(RegionId::new(1, 1), 10);
254        let missing = missing_region_ids(&all_region_ids, &result_set);
255        assert!(missing.is_empty());
256    }
257}
meta_srv/procedure/wal_prune/utils.rs

meta_srv/procedure/wal_prune/
utils.rs
