refactor(meta): store operating region roles in memory keeper

Signed-off-by: WenyXu <wenymedia@gmail.com>

src/common/meta/src/region_keeper.rs

@@ -12,9 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::sync::{Arc, RwLock};
 
+use store_api::region_engine::RegionRole;
 use store_api::storage::RegionId;
 
 use crate::DatanodeId;
@@ -24,7 +25,7 @@ use crate::DatanodeId;
 pub struct OperatingRegionGuard {
     datanode_id: DatanodeId,
     region_id: RegionId,
-    inner: Arc<RwLock<HashSet<(DatanodeId, RegionId)>>>,
+    inner: Arc<RwLock<HashMap<(DatanodeId, RegionId), RegionRole>>>,
 }
 
 impl Drop for OperatingRegionGuard {
@@ -50,7 +51,7 @@ pub type MemoryRegionKeeperRef = Arc<MemoryRegionKeeper>;
 /// - Tracks the deleting regions after the corresponding metadata is deleted.
 #[derive(Debug, Clone, Default)]
 pub struct MemoryRegionKeeper {
-    inner: Arc<RwLock<HashSet<(DatanodeId, RegionId)>>>,
+    inner: Arc<RwLock<HashMap<(DatanodeId, RegionId), RegionRole>>>,
 }
 
 impl MemoryRegionKeeper {
@@ -63,10 +64,20 @@ impl MemoryRegionKeeper {
         &self,
         datanode_id: DatanodeId,
         region_id: RegionId,
+    ) -> Option<OperatingRegionGuard> {
+        self.register_with_role(datanode_id, region_id, RegionRole::Leader)
+    }
+
+    /// Returns [OperatingRegionGuard] if Region(`region_id`) on Peer(`datanode_id`) does not exist.
+    pub fn register_with_role(
+        &self,
+        datanode_id: DatanodeId,
+        region_id: RegionId,
+        role: RegionRole,
     ) -> Option<OperatingRegionGuard> {
         let mut inner = self.inner.write().unwrap();
 
-        if inner.insert((datanode_id, region_id)) {
+        if inner.insert((datanode_id, region_id), role).is_none() {
             Some(OperatingRegionGuard {
                 datanode_id,
                 region_id,
@@ -80,7 +91,7 @@ impl MemoryRegionKeeper {
     /// Returns true if the keeper contains a (`datanoe_id`, `region_id`) tuple.
     pub fn contains(&self, datanode_id: DatanodeId, region_id: RegionId) -> bool {
         let inner = self.inner.read().unwrap();
-        inner.contains(&(datanode_id, region_id))
+        inner.contains_key(&(datanode_id, region_id))
     }
 
     /// Extracts all operating regions from `region_ids` and returns operating regions.
@@ -91,8 +102,30 @@ impl MemoryRegionKeeper {
     ) -> HashSet<RegionId> {
         let inner = self.inner.read().unwrap();
         region_ids
-            .extract_if(|region_id| inner.contains(&(datanode_id, *region_id)))
-            .collect::<HashSet<_>>()
+            .extract_if(|region_id| inner.contains_key(&(datanode_id, *region_id)))
+            .collect()
+    }
+
+    /// Extracts all operating regions with roles from `region_ids`.
+    pub fn extract_operating_region_roles(
+        &self,
+        datanode_id: DatanodeId,
+        region_ids: &mut HashSet<RegionId>,
+    ) -> HashMap<RegionId, RegionRole> {
+        let inner = self.inner.read().unwrap();
+        let operating_regions = region_ids
+            .extract_if(|region_id| inner.contains_key(&(datanode_id, *region_id)))
+            .collect::<Vec<_>>();
+
+        operating_regions
+            .into_iter()
+            .map(|region_id| {
+                let role = *inner
+                    .get(&(datanode_id, region_id))
+                    .expect("operating region role must exist");
+                (region_id, role)
+            })
+            .collect()
     }
 
     /// Returns number of element in tracking set.
@@ -115,8 +148,9 @@ impl MemoryRegionKeeper {
 
 #[cfg(test)]
 mod tests {
-    use std::collections::HashSet;
+    use std::collections::{HashMap, HashSet};
 
+    use store_api::region_engine::RegionRole;
     use store_api::storage::RegionId;
 
     use crate::region_keeper::MemoryRegionKeeper;
@@ -127,7 +161,9 @@ mod tests {
 
         let guard = keeper.register(1, RegionId::from_u64(1)).unwrap();
         assert!(keeper.register(1, RegionId::from_u64(1)).is_none());
-        let guard2 = keeper.register(1, RegionId::from_u64(2)).unwrap();
+        let guard2 = keeper
+            .register_with_role(1, RegionId::from_u64(2), RegionRole::Follower)
+            .unwrap();
 
         let mut regions = HashSet::from([
             RegionId::from_u64(1),
@@ -139,6 +175,23 @@ mod tests {
 
         assert!(output.contains(&RegionId::from_u64(1)));
         assert!(output.contains(&RegionId::from_u64(2)));
+        assert_eq!(regions, HashSet::from([RegionId::from_u64(3)]));
+        assert_eq!(keeper.len(), 2);
+
+        let mut regions = HashSet::from([
+            RegionId::from_u64(1),
+            RegionId::from_u64(2),
+            RegionId::from_u64(3),
+        ]);
+        let output = keeper.extract_operating_region_roles(1, &mut regions);
+        assert_eq!(
+            output,
+            HashMap::from([
+                (RegionId::from_u64(1), RegionRole::Leader),
+                (RegionId::from_u64(2), RegionRole::Follower),
+            ])
+        );
+        assert_eq!(regions, HashSet::from([RegionId::from_u64(3)]));
         assert_eq!(keeper.len(), 2);
 
         drop(guard);