diff --git a/src/common/meta/src/rpc/router.rs b/src/common/meta/src/rpc/router.rs index 0e700cc6da..03a1e0bfa0 100644 --- a/src/common/meta/src/rpc/router.rs +++ b/src/common/meta/src/rpc/router.rs @@ -32,6 +32,10 @@ use crate::key::RegionDistribution; use crate::peer::Peer; use crate::DatanodeId; +/// Returns the distribution of regions to datanodes. +/// +/// The distribution is a map of datanode id to a list of region ids. +/// The list of region ids is sorted in ascending order. pub fn region_distribution(region_routes: &[RegionRoute]) -> RegionDistribution { let mut regions_id_map = RegionDistribution::new(); for route in region_routes.iter() { @@ -39,6 +43,10 @@ pub fn region_distribution(region_routes: &[RegionRoute]) -> RegionDistribution let region_id = route.region.id.region_number(); regions_id_map.entry(peer.id).or_default().push(region_id); } + for peer in route.follower_peers.iter() { + let region_id = route.region.id.region_number(); + regions_id_map.entry(peer.id).or_default().push(region_id); + } } for (_, regions) in regions_id_map.iter_mut() { // id asc @@ -550,4 +558,40 @@ mod tests { assert_eq!(got, p); } + + #[test] + fn test_region_distribution() { + let region_routes = vec![ + RegionRoute { + region: Region { + id: RegionId::new(1, 1), + name: "r1".to_string(), + partition: None, + attrs: BTreeMap::new(), + }, + leader_peer: Some(Peer::new(1, "a1")), + follower_peers: vec![Peer::new(2, "a2"), Peer::new(3, "a3")], + leader_state: None, + leader_down_since: None, + }, + RegionRoute { + region: Region { + id: RegionId::new(1, 2), + name: "r2".to_string(), + partition: None, + attrs: BTreeMap::new(), + }, + leader_peer: Some(Peer::new(2, "a2")), + follower_peers: vec![Peer::new(1, "a1"), Peer::new(3, "a3")], + leader_state: None, + leader_down_since: None, + }, + ]; + + let distribution = region_distribution(®ion_routes); + assert_eq!(distribution.len(), 3); + assert_eq!(distribution[&1], vec![1, 2]); + assert_eq!(distribution[&2], vec![1, 2]); + assert_eq!(distribution[&3], vec![1, 2]); + } }