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

common_meta/
region_keeper.rs
