common_meta/
region_keeper.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::hash_map::Entry;
16use std::collections::{HashMap, HashSet};
17use std::sync::{Arc, RwLock};
18
19use store_api::region_engine::RegionRole;
20use store_api::storage::RegionId;
21
22use crate::DatanodeId;
23
24/// Tracks the operating(i.e., creating, opening, dropping) regions.
25#[derive(Debug, Clone)]
26pub struct OperatingRegionGuard {
27    datanode_id: DatanodeId,
28    region_id: RegionId,
29    inner: Arc<RwLock<HashMap<(DatanodeId, RegionId), RegionRole>>>,
30}
31
32impl Drop for OperatingRegionGuard {
33    fn drop(&mut self) {
34        let mut inner = self.inner.write().unwrap();
35        inner.remove(&(self.datanode_id, self.region_id));
36    }
37}
38
39impl OperatingRegionGuard {
40    /// Returns opening region info.
41    pub fn info(&self) -> (DatanodeId, RegionId) {
42        (self.datanode_id, self.region_id)
43    }
44}
45
46pub type MemoryRegionKeeperRef = Arc<MemoryRegionKeeper>;
47
48/// Tracks regions in memory.
49///
50/// It used in following cases:
51/// - Tracks the opening regions before the corresponding metadata is created.
52/// - Tracks the deleting regions after the corresponding metadata is deleted.
53#[derive(Debug, Clone, Default)]
54pub struct MemoryRegionKeeper {
55    inner: Arc<RwLock<HashMap<(DatanodeId, RegionId), RegionRole>>>,
56}
57
58impl MemoryRegionKeeper {
59    pub fn new() -> Self {
60        Default::default()
61    }
62
63    /// Returns [OperatingRegionGuard] if Region(`region_id`) on Peer(`datanode_id`) does not exist.
64    pub fn register_with_role(
65        &self,
66        datanode_id: DatanodeId,
67        region_id: RegionId,
68        role: RegionRole,
69    ) -> Option<OperatingRegionGuard> {
70        let mut inner = self.inner.write().unwrap();
71
72        match inner.entry((datanode_id, region_id)) {
73            Entry::Occupied(_) => None,
74            Entry::Vacant(vacant_entry) => {
75                vacant_entry.insert(role);
76                Some(OperatingRegionGuard {
77                    datanode_id,
78                    region_id,
79                    inner: self.inner.clone(),
80                })
81            }
82        }
83    }
84
85    /// Returns true if the keeper contains a (`datanoe_id`, `region_id`) tuple.
86    pub fn contains(&self, datanode_id: DatanodeId, region_id: RegionId) -> bool {
87        let inner = self.inner.read().unwrap();
88        inner.contains_key(&(datanode_id, region_id))
89    }
90
91    /// Extracts all operating regions with roles from `region_ids`.
92    pub fn extract_operating_region_roles(
93        &self,
94        datanode_id: DatanodeId,
95        region_ids: &HashSet<RegionId>,
96    ) -> HashMap<RegionId, RegionRole> {
97        let inner = self.inner.read().unwrap();
98        region_ids
99            .iter()
100            .filter_map(|region_id| {
101                inner
102                    .get(&(datanode_id, *region_id))
103                    .map(|role| (*region_id, *role))
104            })
105            .collect()
106    }
107
108    /// Returns number of element in tracking set.
109    pub fn len(&self) -> usize {
110        let inner = self.inner.read().unwrap();
111        inner.len()
112    }
113
114    /// Returns true if it's empty.
115    pub fn is_empty(&self) -> bool {
116        self.len() == 0
117    }
118
119    #[cfg(test)]
120    pub fn clear(&self) {
121        let mut inner = self.inner.write().unwrap();
122        inner.clear();
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use std::collections::{HashMap, HashSet};
129
130    use store_api::region_engine::RegionRole;
131    use store_api::storage::RegionId;
132
133    use crate::region_keeper::MemoryRegionKeeper;
134
135    #[test]
136    fn test_opening_region_keeper() {
137        let keeper = MemoryRegionKeeper::new();
138
139        let guard = keeper
140            .register_with_role(1, RegionId::from_u64(1), RegionRole::Leader)
141            .unwrap();
142        assert!(
143            keeper
144                .register_with_role(1, RegionId::from_u64(1), RegionRole::Leader)
145                .is_none()
146        );
147        let guard2 = keeper
148            .register_with_role(1, RegionId::from_u64(2), RegionRole::Follower)
149            .unwrap();
150
151        let regions = HashSet::from([
152            RegionId::from_u64(1),
153            RegionId::from_u64(2),
154            RegionId::from_u64(3),
155        ]);
156        let output = keeper.extract_operating_region_roles(1, &regions);
157        assert_eq!(output.len(), 2);
158
159        assert!(output.contains_key(&RegionId::from_u64(1)));
160        assert!(output.contains_key(&RegionId::from_u64(2)));
161        assert_eq!(keeper.len(), 2);
162
163        let regions = HashSet::from([
164            RegionId::from_u64(1),
165            RegionId::from_u64(2),
166            RegionId::from_u64(3),
167        ]);
168        let output = keeper.extract_operating_region_roles(1, &regions);
169        assert_eq!(
170            output,
171            HashMap::from([
172                (RegionId::from_u64(1), RegionRole::Leader),
173                (RegionId::from_u64(2), RegionRole::Follower),
174            ])
175        );
176        assert_eq!(keeper.len(), 2);
177
178        drop(guard);
179        assert_eq!(keeper.len(), 1);
180        assert!(keeper.contains(1, RegionId::from_u64(2)));
181
182        drop(guard2);
183        assert!(keeper.is_empty());
184    }
185}
common_meta/region_keeper.rs

common_meta/
region_keeper.rs
