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