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)) => Some(RegionStatistic {
52 num_rows: metadata_stat.num_rows + data_stat.num_rows,
53 memtable_size: metadata_stat.memtable_size + data_stat.memtable_size,
54 wal_size: metadata_stat.wal_size + data_stat.wal_size,
55 manifest_size: metadata_stat.manifest_size + data_stat.manifest_size,
56 sst_size: metadata_stat.sst_size + data_stat.sst_size,
57 index_size: metadata_stat.index_size + data_stat.index_size,
58 manifest: RegionManifestInfo::Metric {
59 data_flushed_entry_id: data_stat.manifest.data_flushed_entry_id(),
60 data_manifest_version: data_stat.manifest.data_manifest_version(),
61 metadata_flushed_entry_id: metadata_stat.manifest.data_flushed_entry_id(),
62 metadata_manifest_version: metadata_stat.manifest.data_manifest_version(),
63 },
64 data_topic_latest_entry_id: data_stat.data_topic_latest_entry_id,
65 metadata_topic_latest_entry_id: metadata_stat.metadata_topic_latest_entry_id,
66 }),
67 _ => {
68 warn!(
69 "Failed to get region statistic for region {}, metadata_stat: {:?}, data_stat: {:?}",
70 region_id, metadata_stat, data_stat
71 );
72 None
73 }
74 }
75}
76
77pub(crate) fn append_manifest_info(
79 mito: &MitoEngine,
80 region_id: RegionId,
81 manifest_infos: &mut Vec<(RegionId, RegionManifestInfo)>,
82) {
83 if let Some(statistic) = get_region_statistic(mito, region_id) {
84 manifest_infos.push((region_id, statistic.manifest));
85 }
86}
87
88pub(crate) fn encode_manifest_info_to_extensions(
90 manifest_infos: &[(RegionId, RegionManifestInfo)],
91 extensions: &mut HashMap<String, Vec<u8>>,
92) -> Result<()> {
93 extensions.insert(
94 MANIFEST_INFO_EXTENSION_KEY.to_string(),
95 RegionManifestInfo::encode_list(manifest_infos)
96 .context(SerializeRegionManifestInfoSnafu)?,
97 );
98 for (region_id, manifest_info) in manifest_infos {
99 info!(
100 "Added manifest info: {:?} to extensions, region_id: {:?}",
101 manifest_info, region_id
102 );
103 }
104 Ok(())
105}
106
107#[cfg(test)]
108mod tests {
109
110 use super::*;
111
112 #[test]
113 fn test_to_metadata_region_id() {
114 let region_id = RegionId::new(1, 2);
115 let expected_region_id = RegionId::with_group_and_seq(1, METRIC_METADATA_REGION_GROUP, 2);
116 assert_eq!(to_metadata_region_id(region_id), expected_region_id);
117
118 let region_id = RegionId::with_group_and_seq(1, 243, 2);
119 let expected_region_id = RegionId::with_group_and_seq(1, METRIC_METADATA_REGION_GROUP, 2);
120 assert_eq!(to_metadata_region_id(region_id), expected_region_id);
121 }
122
123 #[test]
124 fn test_to_data_region_id() {
125 let region_id = RegionId::new(1, 2);
126 let expected_region_id = RegionId::with_group_and_seq(1, METRIC_DATA_REGION_GROUP, 2);
127 assert_eq!(to_data_region_id(region_id), expected_region_id);
128
129 let region_id = RegionId::with_group_and_seq(1, 243, 2);
130 let expected_region_id = RegionId::with_group_and_seq(1, METRIC_DATA_REGION_GROUP, 2);
131 assert_eq!(to_data_region_id(region_id), expected_region_id);
132 }
133}