Skip to main content

metric_engine/engine/
state.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
15//! Internal states of metric engine
16
17use std::collections::{HashMap, HashSet};
18
19use api::v1::SemanticType;
20use common_time::timestamp::TimeUnit;
21use snafu::OptionExt;
22use store_api::codec::PrimaryKeyEncoding;
23use store_api::metadata::ColumnMetadata;
24use store_api::storage::RegionId;
25
26use crate::engine::options::PhysicalRegionOptions;
27use crate::error::{PhysicalRegionNotFoundSnafu, Result};
28use crate::metrics::LOGICAL_REGION_COUNT;
29use crate::utils::to_data_region_id;
30
31pub struct PhysicalRegionState {
32    logical_regions: HashSet<RegionId>,
33    physical_columns: HashMap<String, ColumnMetadata>,
34    /// Name of the time index column, cached at region load so that the write
35    /// path doesn't have to scan `physical_columns` for the timestamp on every
36    /// row batch. The time index is fixed at region creation and never
37    /// changes, so this stays in sync with `physical_columns`.
38    time_index_column_name: String,
39    primary_key_encoding: PrimaryKeyEncoding,
40    options: PhysicalRegionOptions,
41    time_index_unit: TimeUnit,
42}
43
44impl PhysicalRegionState {
45    pub fn new(
46        physical_columns: HashMap<String, ColumnMetadata>,
47        primary_key_encoding: PrimaryKeyEncoding,
48        options: PhysicalRegionOptions,
49        time_index_unit: TimeUnit,
50    ) -> Self {
51        // Safety: a valid physical region always has exactly one time index
52        // column; callers validate this before reaching here (see
53        // `create_data_region_request` and the open path).
54        let time_index_column_name = physical_columns
55            .iter()
56            .find(|(_, meta)| meta.semantic_type == SemanticType::Timestamp)
57            .map(|(name, _)| name.clone())
58            .unwrap_or_default();
59        Self {
60            logical_regions: HashSet::new(),
61            physical_columns,
62            time_index_column_name,
63            primary_key_encoding,
64            options,
65            time_index_unit,
66        }
67    }
68
69    /// Returns a reference to the logical region ids.
70    pub fn logical_regions(&self) -> &HashSet<RegionId> {
71        &self.logical_regions
72    }
73
74    /// Returns a reference to the physical columns.
75    pub fn physical_columns(&self) -> &HashMap<String, ColumnMetadata> {
76        &self.physical_columns
77    }
78
79    /// Returns the cached name of the time index column.
80    pub fn time_index_column_name(&self) -> &str {
81        &self.time_index_column_name
82    }
83
84    /// Returns a reference to the physical region options.
85    pub fn options(&self) -> &PhysicalRegionOptions {
86        &self.options
87    }
88
89    /// Removes a logical region id from the physical region state.
90    /// Returns true if the logical region id was present.
91    pub fn remove_logical_region(&mut self, logical_region_id: RegionId) -> bool {
92        self.logical_regions.remove(&logical_region_id)
93    }
94}
95
96/// Internal states of metric engine
97#[derive(Default)]
98pub(crate) struct MetricEngineState {
99    /// Physical regions states.
100    physical_regions: HashMap<RegionId, PhysicalRegionState>,
101    /// Mapping from logical region id to physical region id.
102    logical_regions: HashMap<RegionId, RegionId>,
103    /// Cache for the column metadata of logical regions.
104    /// The column order is the same with the order in the metadata, which is
105    /// alphabetically ordered on column name.
106    logical_columns: HashMap<RegionId, Vec<ColumnMetadata>>,
107}
108
109impl MetricEngineState {
110    pub fn add_physical_region(
111        &mut self,
112        physical_region_id: RegionId,
113        physical_columns: HashMap<String, ColumnMetadata>,
114        primary_key_encoding: PrimaryKeyEncoding,
115        options: PhysicalRegionOptions,
116        time_index_unit: TimeUnit,
117    ) {
118        let physical_region_id = to_data_region_id(physical_region_id);
119        self.physical_regions.insert(
120            physical_region_id,
121            PhysicalRegionState::new(
122                physical_columns,
123                primary_key_encoding,
124                options,
125                time_index_unit,
126            ),
127        );
128    }
129
130    /// # Panic
131    /// if the physical region does not exist
132    pub fn add_physical_columns(
133        &mut self,
134        physical_region_id: RegionId,
135        physical_columns: impl IntoIterator<Item = (String, ColumnMetadata)>,
136    ) {
137        let physical_region_id = to_data_region_id(physical_region_id);
138        let state = self.physical_regions.get_mut(&physical_region_id).unwrap();
139        for (col, meta) in physical_columns {
140            // The time index is fixed at region creation and alter cannot add
141            // a new one; keep the cached name in sync defensively.
142            debug_assert_ne!(
143                meta.semantic_type,
144                SemanticType::Timestamp,
145                "unexpected time index column {col} added to an existing physical region"
146            );
147            state.physical_columns.insert(col, meta);
148        }
149    }
150
151    /// # Panic
152    /// if the physical region does not exist
153    pub fn add_logical_regions(
154        &mut self,
155        physical_region_id: RegionId,
156        logical_region_ids: impl IntoIterator<Item = RegionId>,
157    ) {
158        let physical_region_id = to_data_region_id(physical_region_id);
159        let state = self.physical_regions.get_mut(&physical_region_id).unwrap();
160        for logical_region_id in logical_region_ids {
161            state.logical_regions.insert(logical_region_id);
162            self.logical_regions
163                .insert(logical_region_id, physical_region_id);
164        }
165    }
166
167    pub fn invalid_logical_regions_cache(
168        &mut self,
169        logical_region_ids: impl IntoIterator<Item = RegionId>,
170    ) {
171        for logical_region_id in logical_region_ids {
172            self.logical_columns.remove(&logical_region_id);
173        }
174    }
175
176    /// # Panic
177    /// if the physical region does not exist
178    pub fn add_logical_region(
179        &mut self,
180        physical_region_id: RegionId,
181        logical_region_id: RegionId,
182    ) {
183        let physical_region_id = to_data_region_id(physical_region_id);
184        self.physical_regions
185            .get_mut(&physical_region_id)
186            .unwrap()
187            .logical_regions
188            .insert(logical_region_id);
189        self.logical_regions
190            .insert(logical_region_id, physical_region_id);
191    }
192
193    /// Replace the logical columns of the logical region with given columns.
194    pub fn set_logical_columns(
195        &mut self,
196        logical_region_id: RegionId,
197        columns: Vec<ColumnMetadata>,
198    ) {
199        self.logical_columns.insert(logical_region_id, columns);
200    }
201
202    pub fn get_physical_region_id(&self, logical_region_id: RegionId) -> Option<RegionId> {
203        self.logical_regions.get(&logical_region_id).copied()
204    }
205
206    pub fn logical_columns(&self) -> &HashMap<RegionId, Vec<ColumnMetadata>> {
207        &self.logical_columns
208    }
209
210    pub fn physical_region_states(&self) -> &HashMap<RegionId, PhysicalRegionState> {
211        &self.physical_regions
212    }
213
214    pub fn exist_physical_region(&self, physical_region_id: RegionId) -> bool {
215        self.physical_regions.contains_key(&physical_region_id)
216    }
217
218    pub fn physical_region_time_index_unit(
219        &self,
220        physical_region_id: RegionId,
221    ) -> Option<TimeUnit> {
222        self.physical_regions
223            .get(&physical_region_id)
224            .map(|state| state.time_index_unit)
225    }
226
227    pub fn get_primary_key_encoding(
228        &self,
229        physical_region_id: RegionId,
230    ) -> Option<PrimaryKeyEncoding> {
231        self.physical_regions
232            .get(&physical_region_id)
233            .map(|state| state.primary_key_encoding)
234    }
235
236    pub fn logical_regions(&self) -> &HashMap<RegionId, RegionId> {
237        &self.logical_regions
238    }
239
240    /// Remove all data that are related to the physical region id.
241    pub fn remove_physical_region(&mut self, physical_region_id: RegionId) -> Result<()> {
242        let physical_region_id = to_data_region_id(physical_region_id);
243
244        let logical_regions = &self
245            .physical_regions
246            .get(&physical_region_id)
247            .context(PhysicalRegionNotFoundSnafu {
248                region_id: physical_region_id,
249            })?
250            .logical_regions;
251
252        LOGICAL_REGION_COUNT.sub(logical_regions.len() as i64);
253
254        for logical_region in logical_regions {
255            self.logical_regions.remove(logical_region);
256        }
257        self.physical_regions.remove(&physical_region_id);
258        Ok(())
259    }
260
261    /// Remove all data that are related to the logical region id.
262    pub fn remove_logical_region(&mut self, logical_region_id: RegionId) -> Result<()> {
263        let physical_region_id = self.logical_regions.remove(&logical_region_id).context(
264            PhysicalRegionNotFoundSnafu {
265                region_id: logical_region_id,
266            },
267        )?;
268
269        self.physical_regions
270            .get_mut(&physical_region_id)
271            .unwrap() // Safety: physical_region_id is got from physical_regions
272            .remove_logical_region(logical_region_id);
273
274        self.logical_columns.remove(&logical_region_id);
275
276        Ok(())
277    }
278
279    pub fn is_logical_region_exist(&self, logical_region_id: RegionId) -> bool {
280        self.logical_regions().contains_key(&logical_region_id)
281    }
282}