1use std::collections::HashMap;
16
17use common_telemetry::{info, warn};
18use mito2::engine::MitoEngine;
19use snafu::ResultExt;
20use store_api::metric_engine_consts::{
21 MANIFEST_INFO_EXTENSION_KEY, METRIC_DATA_REGION_GROUP, METRIC_METADATA_REGION_GROUP,
22};
23use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionStatistic};
24use store_api::storage::RegionId;
25
26use crate::error::{Result, SerializeRegionManifestInfoSnafu};
27
28pub fn to_metadata_region_id(region_id: RegionId) -> RegionId {
30 let table_id = region_id.table_id();
31 let region_sequence = region_id.region_sequence();
32 RegionId::with_group_and_seq(table_id, METRIC_METADATA_REGION_GROUP, region_sequence)
33}
34
35pub fn to_data_region_id(region_id: RegionId) -> RegionId {
37 let table_id = region_id.table_id();
38 let region_sequence = region_id.region_sequence();
39 RegionId::with_group_and_seq(table_id, METRIC_DATA_REGION_GROUP, region_sequence)
40}
41
42pub fn get_region_statistic(mito: &MitoEngine, region_id: RegionId) -> Option<RegionStatistic> {
44 let metadata_region_id = to_metadata_region_id(region_id);
45 let data_region_id = to_data_region_id(region_id);
46
47 let metadata_stat = mito.region_statistic(metadata_region_id);
48 let data_stat = mito.region_statistic(data_region_id);
49
50 match (&metadata_stat, &data_stat) {
51 (Some(metadata_stat), Some(data_stat)) => {
52 Some(merge_region_statistic(metadata_stat, data_stat))
53 }
54 _ => {
55 warn!(
56 "Failed to get region statistic for region {}, metadata_stat: {:?}, data_stat: {:?}",
57 region_id, metadata_stat, data_stat
58 );
59 None
60 }
61 }
62}
63
64fn merge_region_statistic(
65 metadata_stat: &RegionStatistic,
66 data_stat: &RegionStatistic,
67) -> RegionStatistic {
68 RegionStatistic {
69 num_rows: metadata_stat.num_rows + data_stat.num_rows,
70 memtable_size: metadata_stat.memtable_size + data_stat.memtable_size,
71 wal_size: metadata_stat.wal_size + data_stat.wal_size,
72 manifest_size: metadata_stat.manifest_size + data_stat.manifest_size,
73 sst_size: metadata_stat.sst_size + data_stat.sst_size,
74 sst_num: metadata_stat.sst_num + data_stat.sst_num,
75 index_size: metadata_stat.index_size + data_stat.index_size,
76 manifest: RegionManifestInfo::Metric {
77 data_flushed_entry_id: data_stat.manifest.data_flushed_entry_id(),
78 data_manifest_version: data_stat.manifest.data_manifest_version(),
79 metadata_flushed_entry_id: metadata_stat.manifest.data_flushed_entry_id(),
80 metadata_manifest_version: metadata_stat.manifest.data_manifest_version(),
81 },
82 written_bytes: metadata_stat.written_bytes + data_stat.written_bytes,
83 query_cpu_time: data_stat.query_cpu_time,
84 query_scanned_bytes: data_stat.query_scanned_bytes,
85 data_topic_latest_entry_id: data_stat.data_topic_latest_entry_id,
86 metadata_topic_latest_entry_id: metadata_stat.metadata_topic_latest_entry_id,
87 min_timestamp: data_stat.min_timestamp,
90 max_timestamp: data_stat.max_timestamp,
91 }
92}
93
94pub(crate) fn append_manifest_info(
96 mito: &MitoEngine,
97 region_id: RegionId,
98 manifest_infos: &mut Vec<(RegionId, RegionManifestInfo)>,
99) {
100 if let Some(statistic) = get_region_statistic(mito, region_id) {
101 manifest_infos.push((region_id, statistic.manifest));
102 }
103}
104
105pub(crate) fn encode_manifest_info_to_extensions(
107 manifest_infos: &[(RegionId, RegionManifestInfo)],
108 extensions: &mut HashMap<String, Vec<u8>>,
109) -> Result<()> {
110 extensions.insert(
111 MANIFEST_INFO_EXTENSION_KEY.to_string(),
112 RegionManifestInfo::encode_list(manifest_infos)
113 .context(SerializeRegionManifestInfoSnafu)?,
114 );
115 for (region_id, manifest_info) in manifest_infos {
116 info!(
117 "Added manifest info: {:?} to extensions, region_id: {:?}",
118 manifest_info, region_id
119 );
120 }
121 Ok(())
122}
123
124#[cfg(test)]
125mod tests {
126
127 use common_time::Timestamp;
128
129 use super::*;
130
131 #[test]
132 fn test_to_metadata_region_id() {
133 let region_id = RegionId::new(1, 2);
134 let expected_region_id = RegionId::with_group_and_seq(1, METRIC_METADATA_REGION_GROUP, 2);
135 assert_eq!(to_metadata_region_id(region_id), expected_region_id);
136
137 let region_id = RegionId::with_group_and_seq(1, 243, 2);
138 let expected_region_id = RegionId::with_group_and_seq(1, METRIC_METADATA_REGION_GROUP, 2);
139 assert_eq!(to_metadata_region_id(region_id), expected_region_id);
140 }
141
142 #[test]
143 fn test_to_data_region_id() {
144 let region_id = RegionId::new(1, 2);
145 let expected_region_id = RegionId::with_group_and_seq(1, METRIC_DATA_REGION_GROUP, 2);
146 assert_eq!(to_data_region_id(region_id), expected_region_id);
147
148 let region_id = RegionId::with_group_and_seq(1, 243, 2);
149 let expected_region_id = RegionId::with_group_and_seq(1, METRIC_DATA_REGION_GROUP, 2);
150 assert_eq!(to_data_region_id(region_id), expected_region_id);
151 }
152
153 #[test]
154 fn merge_region_statistic_ignores_metadata_timestamps() {
155 let metadata_stat = RegionStatistic {
158 min_timestamp: Some(Timestamp::new_millisecond(0)),
159 max_timestamp: Some(Timestamp::new_millisecond(0)),
160 ..Default::default()
161 };
162 let data_stat = RegionStatistic {
163 min_timestamp: Some(Timestamp::new_millisecond(1_700_000_000_000)),
164 max_timestamp: Some(Timestamp::new_millisecond(1_700_000_001_000)),
165 ..Default::default()
166 };
167
168 let statistic = merge_region_statistic(&metadata_stat, &data_stat);
169
170 assert_eq!(
171 statistic.min_timestamp,
172 Some(Timestamp::new_millisecond(1_700_000_000_000))
173 );
174 assert_eq!(
175 statistic.max_timestamp,
176 Some(Timestamp::new_millisecond(1_700_000_001_000))
177 );
178 }
179
180 #[test]
181 fn merge_region_statistic_reports_no_range_for_an_empty_data_region() {
182 let metadata_stat = RegionStatistic {
185 min_timestamp: Some(Timestamp::new_millisecond(0)),
186 max_timestamp: Some(Timestamp::new_millisecond(0)),
187 ..Default::default()
188 };
189
190 let statistic = merge_region_statistic(&metadata_stat, &RegionStatistic::default());
191
192 assert_eq!(statistic.min_timestamp, None);
193 assert_eq!(statistic.max_timestamp, None);
194 }
195
196 #[test]
197 fn merge_region_statistic_keeps_pre_epoch_samples() {
198 let metadata_stat = RegionStatistic {
201 min_timestamp: Some(Timestamp::new_millisecond(0)),
202 max_timestamp: Some(Timestamp::new_millisecond(0)),
203 ..Default::default()
204 };
205 let data_stat = RegionStatistic {
206 min_timestamp: Some(Timestamp::new_millisecond(-2_000)),
207 max_timestamp: Some(Timestamp::new_millisecond(-1_000)),
208 ..Default::default()
209 };
210
211 let statistic = merge_region_statistic(&metadata_stat, &data_stat);
212
213 assert_eq!(
214 statistic.min_timestamp,
215 Some(Timestamp::new_millisecond(-2_000))
216 );
217 assert_eq!(
218 statistic.max_timestamp,
219 Some(Timestamp::new_millisecond(-1_000))
220 );
221 }
222
223 #[test]
224 fn merge_region_statistic_uses_data_region_query_stats() {
225 let metadata_stat = RegionStatistic {
226 query_cpu_time: 100,
227 query_scanned_bytes: 200,
228 manifest: RegionManifestInfo::mito(1, 2, 0),
229 ..Default::default()
230 };
231 let data_stat = RegionStatistic {
232 query_cpu_time: 10,
233 query_scanned_bytes: 20,
234 manifest: RegionManifestInfo::mito(3, 4, 0),
235 ..Default::default()
236 };
237
238 let statistic = merge_region_statistic(&metadata_stat, &data_stat);
239
240 assert_eq!(statistic.query_cpu_time, 10);
241 assert_eq!(statistic.query_scanned_bytes, 20);
242 }
243}