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        region_ids
94            .extract_if(|region_id| inner.contains(&(datanode_id, *region_id)))
95            .collect::<HashSet<_>>()
96    }
97
98    /// Returns number of element in tracking set.
99    pub fn len(&self) -> usize {
100        let inner = self.inner.read().unwrap();
101        inner.len()
102    }
103
104    /// Returns true if it's empty.
105    pub fn is_empty(&self) -> bool {
106        self.len() == 0
107    }
108
109    #[cfg(test)]
110    pub fn clear(&self) {
111        let mut inner = self.inner.write().unwrap();
112        inner.clear();
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use std::collections::HashSet;
119
120    use store_api::storage::RegionId;
121
122    use crate::region_keeper::MemoryRegionKeeper;
123
124    #[test]
125    fn test_opening_region_keeper() {
126        let keeper = MemoryRegionKeeper::new();
127
128        let guard = keeper.register(1, RegionId::from_u64(1)).unwrap();
129        assert!(keeper.register(1, RegionId::from_u64(1)).is_none());
130        let guard2 = keeper.register(1, RegionId::from_u64(2)).unwrap();
131
132        let mut regions = HashSet::from([
133            RegionId::from_u64(1),
134            RegionId::from_u64(2),
135            RegionId::from_u64(3),
136        ]);
137        let output = keeper.extract_operating_regions(1, &mut regions);
138        assert_eq!(output.len(), 2);
139
140        assert!(output.contains(&RegionId::from_u64(1)));
141        assert!(output.contains(&RegionId::from_u64(2)));
142        assert_eq!(keeper.len(), 2);
143
144        drop(guard);
145        assert_eq!(keeper.len(), 1);
146        assert!(keeper.contains(1, RegionId::from_u64(2)));
147
148        drop(guard2);
149        assert!(keeper.is_empty());
150    }
151}
common_meta/region_keeper.rs

common_meta/
region_keeper.rs
