Skip to main content

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)) => {
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        ingested_parquet_bytes: metadata_stat.ingested_parquet_bytes
84            + data_stat.ingested_parquet_bytes,
85        query_cpu_time: data_stat.query_cpu_time,
86        query_scanned_bytes: data_stat.query_scanned_bytes,
87        data_topic_latest_entry_id: data_stat.data_topic_latest_entry_id,
88        metadata_topic_latest_entry_id: metadata_stat.metadata_topic_latest_entry_id,
89        // Metadata rows are written at timestamp 0, so merging their bounds would
90        // pin every metric region's minimum to the epoch.
91        min_timestamp: data_stat.min_timestamp,
92        max_timestamp: data_stat.max_timestamp,
93    }
94}
95
96/// Appends the given [RegionId]'s manifest info to the given list.
97pub(crate) fn append_manifest_info(
98    mito: &MitoEngine,
99    region_id: RegionId,
100    manifest_infos: &mut Vec<(RegionId, RegionManifestInfo)>,
101) {
102    if let Some(statistic) = get_region_statistic(mito, region_id) {
103        manifest_infos.push((region_id, statistic.manifest));
104    }
105}
106
107/// Encodes the given list of ([RegionId], [RegionManifestInfo]) to extensions(key: MANIFEST_INFO_EXTENSION_KEY).
108pub(crate) fn encode_manifest_info_to_extensions(
109    manifest_infos: &[(RegionId, RegionManifestInfo)],
110    extensions: &mut HashMap<String, Vec<u8>>,
111) -> Result<()> {
112    extensions.insert(
113        MANIFEST_INFO_EXTENSION_KEY.to_string(),
114        RegionManifestInfo::encode_list(manifest_infos)
115            .context(SerializeRegionManifestInfoSnafu)?,
116    );
117    for (region_id, manifest_info) in manifest_infos {
118        info!(
119            "Added manifest info: {:?} to extensions, region_id: {:?}",
120            manifest_info, region_id
121        );
122    }
123    Ok(())
124}
125
126#[cfg(test)]
127mod tests {
128
129    use common_time::Timestamp;
130
131    use super::*;
132
133    #[test]
134    fn test_to_metadata_region_id() {
135        let region_id = RegionId::new(1, 2);
136        let expected_region_id = RegionId::with_group_and_seq(1, METRIC_METADATA_REGION_GROUP, 2);
137        assert_eq!(to_metadata_region_id(region_id), expected_region_id);
138
139        let region_id = RegionId::with_group_and_seq(1, 243, 2);
140        let expected_region_id = RegionId::with_group_and_seq(1, METRIC_METADATA_REGION_GROUP, 2);
141        assert_eq!(to_metadata_region_id(region_id), expected_region_id);
142    }
143
144    #[test]
145    fn test_to_data_region_id() {
146        let region_id = RegionId::new(1, 2);
147        let expected_region_id = RegionId::with_group_and_seq(1, METRIC_DATA_REGION_GROUP, 2);
148        assert_eq!(to_data_region_id(region_id), expected_region_id);
149
150        let region_id = RegionId::with_group_and_seq(1, 243, 2);
151        let expected_region_id = RegionId::with_group_and_seq(1, METRIC_DATA_REGION_GROUP, 2);
152        assert_eq!(to_data_region_id(region_id), expected_region_id);
153    }
154
155    #[test]
156    fn merge_region_statistic_sums_ingested_parquet_bytes() {
157        let metadata_stat = RegionStatistic {
158            ingested_parquet_bytes: 123,
159            ..Default::default()
160        };
161        let data_stat = RegionStatistic {
162            ingested_parquet_bytes: 456,
163            ..Default::default()
164        };
165        assert_eq!(
166            merge_region_statistic(&metadata_stat, &data_stat).ingested_parquet_bytes,
167            579
168        );
169    }
170
171    #[test]
172    fn merge_region_statistic_ignores_metadata_timestamps() {
173        // Metadata rows are written at timestamp 0; folding them in would report
174        // 1970 as the minimum for a region holding only modern samples.
175        let metadata_stat = RegionStatistic {
176            min_timestamp: Some(Timestamp::new_millisecond(0)),
177            max_timestamp: Some(Timestamp::new_millisecond(0)),
178            ..Default::default()
179        };
180        let data_stat = RegionStatistic {
181            min_timestamp: Some(Timestamp::new_millisecond(1_700_000_000_000)),
182            max_timestamp: Some(Timestamp::new_millisecond(1_700_000_001_000)),
183            ..Default::default()
184        };
185
186        let statistic = merge_region_statistic(&metadata_stat, &data_stat);
187
188        assert_eq!(
189            statistic.min_timestamp,
190            Some(Timestamp::new_millisecond(1_700_000_000_000))
191        );
192        assert_eq!(
193            statistic.max_timestamp,
194            Some(Timestamp::new_millisecond(1_700_000_001_000))
195        );
196    }
197
198    #[test]
199    fn merge_region_statistic_reports_no_range_for_an_empty_data_region() {
200        // Metadata present but no samples: the region has no data to bound, so
201        // both ends must stay NULL rather than collapsing onto the epoch.
202        let metadata_stat = RegionStatistic {
203            min_timestamp: Some(Timestamp::new_millisecond(0)),
204            max_timestamp: Some(Timestamp::new_millisecond(0)),
205            ..Default::default()
206        };
207
208        let statistic = merge_region_statistic(&metadata_stat, &RegionStatistic::default());
209
210        assert_eq!(statistic.min_timestamp, None);
211        assert_eq!(statistic.max_timestamp, None);
212    }
213
214    #[test]
215    fn merge_region_statistic_keeps_pre_epoch_samples() {
216        // Samples entirely before 1970 must not have their maximum pulled up to
217        // the metadata timestamp.
218        let metadata_stat = RegionStatistic {
219            min_timestamp: Some(Timestamp::new_millisecond(0)),
220            max_timestamp: Some(Timestamp::new_millisecond(0)),
221            ..Default::default()
222        };
223        let data_stat = RegionStatistic {
224            min_timestamp: Some(Timestamp::new_millisecond(-2_000)),
225            max_timestamp: Some(Timestamp::new_millisecond(-1_000)),
226            ..Default::default()
227        };
228
229        let statistic = merge_region_statistic(&metadata_stat, &data_stat);
230
231        assert_eq!(
232            statistic.min_timestamp,
233            Some(Timestamp::new_millisecond(-2_000))
234        );
235        assert_eq!(
236            statistic.max_timestamp,
237            Some(Timestamp::new_millisecond(-1_000))
238        );
239    }
240
241    #[test]
242    fn merge_region_statistic_uses_data_region_query_stats() {
243        let metadata_stat = RegionStatistic {
244            query_cpu_time: 100,
245            query_scanned_bytes: 200,
246            manifest: RegionManifestInfo::mito(1, 2, 0),
247            ..Default::default()
248        };
249        let data_stat = RegionStatistic {
250            query_cpu_time: 10,
251            query_scanned_bytes: 20,
252            manifest: RegionManifestInfo::mito(3, 4, 0),
253            ..Default::default()
254        };
255
256        let statistic = merge_region_statistic(&metadata_stat, &data_stat);
257
258        assert_eq!(statistic.query_cpu_time, 10);
259        assert_eq!(statistic.query_scanned_bytes, 20);
260    }
261}