From f3509fa312e04eb113cf816f9c0b2e32688e40c7 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 5 Nov 2024 10:55:11 +0800 Subject: [PATCH] chore: minor refactor for weighted choose (#4917) * chore: minor refactor for weighted choose * chore: by comment, remove the fast path of choose_multiple --- src/meta-srv/src/selector/common.rs | 44 ++++++++++++++++-------- src/meta-srv/src/selector/lease_based.rs | 4 +-- src/meta-srv/src/selector/load_based.rs | 4 +-- 3 files changed, 34 insertions(+), 18 deletions(-) diff --git a/src/meta-srv/src/selector/common.rs b/src/meta-srv/src/selector/common.rs index 11e5b3741a68..85abcdfca24a 100644 --- a/src/meta-srv/src/selector/common.rs +++ b/src/meta-srv/src/selector/common.rs @@ -22,7 +22,7 @@ use crate::metasrv::SelectTarget; use crate::selector::SelectorOptions; /// According to the `opts`, choose peers from the `weight_array` through `weighted_choose`. -pub fn choose_peers(opts: &SelectorOptions, weighted_choose: &mut W) -> Result> +pub fn choose_items(opts: &SelectorOptions, weighted_choose: &mut W) -> Result> where W: WeightedChoose, { @@ -36,20 +36,36 @@ where } ); + if min_required_items == 1 { + // fast path + return Ok(vec![weighted_choose.choose_one()?]); + } + + let available_count = weighted_choose.len(); + if opts.allow_duplication { - (0..min_required_items) - .map(|_| weighted_choose.choose_one()) - .collect::>() - } else { - let weight_array_len = weighted_choose.len(); + // Calculate how many complete rounds of `available_count` items to select, + // plus any additional items needed after complete rounds. + let complete_batches = min_required_items / available_count; + let leftover_items = min_required_items % available_count; + if complete_batches == 0 { + return weighted_choose.choose_multiple(leftover_items); + } + + let mut result = Vec::with_capacity(min_required_items); + for _ in 0..complete_batches { + result.extend(weighted_choose.choose_multiple(available_count)?); + } + result.extend(weighted_choose.choose_multiple(leftover_items)?); - // When opts.allow_duplication is false, we need to check that the length of the weighted array is greater than - // or equal to min_required_items, otherwise it may cause an infinite loop. + Ok(result) + } else { + // Ensure the available items are sufficient when duplication is not allowed. ensure!( - weight_array_len >= min_required_items, + available_count >= min_required_items, error::NoEnoughAvailableNodeSnafu { required: min_required_items, - available: weight_array_len, + available: available_count, select_target: SelectTarget::Datanode } ); @@ -64,7 +80,7 @@ mod tests { use common_meta::peer::Peer; - use crate::selector::common::choose_peers; + use crate::selector::common::choose_items; use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem}; use crate::selector::SelectorOptions; @@ -115,7 +131,7 @@ mod tests { }; let selected_peers: HashSet<_> = - choose_peers(&opts, &mut RandomWeightedChoose::new(weight_array.clone())) + choose_items(&opts, &mut RandomWeightedChoose::new(weight_array.clone())) .unwrap() .into_iter() .collect(); @@ -129,7 +145,7 @@ mod tests { }; let selected_result = - choose_peers(&opts, &mut RandomWeightedChoose::new(weight_array.clone())); + choose_items(&opts, &mut RandomWeightedChoose::new(weight_array.clone())); assert!(selected_result.is_err()); for i in 1..=50 { @@ -139,7 +155,7 @@ mod tests { }; let selected_peers = - choose_peers(&opts, &mut RandomWeightedChoose::new(weight_array.clone())).unwrap(); + choose_items(&opts, &mut RandomWeightedChoose::new(weight_array.clone())).unwrap(); assert_eq!(i, selected_peers.len()); } diff --git a/src/meta-srv/src/selector/lease_based.rs b/src/meta-srv/src/selector/lease_based.rs index 3ab99eb31e6b..d9af63da6555 100644 --- a/src/meta-srv/src/selector/lease_based.rs +++ b/src/meta-srv/src/selector/lease_based.rs @@ -17,7 +17,7 @@ use common_meta::peer::Peer; use crate::error::Result; use crate::lease; use crate::metasrv::SelectorContext; -use crate::selector::common::choose_peers; +use crate::selector::common::choose_items; use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem}; use crate::selector::{Namespace, Selector, SelectorOptions}; @@ -53,7 +53,7 @@ impl Selector for LeaseBasedSelector { // 3. choose peers by weight_array. let mut weighted_choose = RandomWeightedChoose::new(weight_array); - let selected = choose_peers(&opts, &mut weighted_choose)?; + let selected = choose_items(&opts, &mut weighted_choose)?; Ok(selected) } diff --git a/src/meta-srv/src/selector/load_based.rs b/src/meta-srv/src/selector/load_based.rs index f52d6f9fc38e..8a00c7fdb7bd 100644 --- a/src/meta-srv/src/selector/load_based.rs +++ b/src/meta-srv/src/selector/load_based.rs @@ -26,7 +26,7 @@ use crate::error::{self, Result}; use crate::key::{DatanodeLeaseKey, LeaseValue}; use crate::lease; use crate::metasrv::SelectorContext; -use crate::selector::common::choose_peers; +use crate::selector::common::choose_items; use crate::selector::weight_compute::{RegionNumsBasedWeightCompute, WeightCompute}; use crate::selector::weighted_choose::RandomWeightedChoose; use crate::selector::{Namespace, Selector, SelectorOptions}; @@ -94,7 +94,7 @@ where // 5. choose peers by weight_array. let mut weighted_choose = RandomWeightedChoose::new(weight_array); - let selected = choose_peers(&opts, &mut weighted_choose)?; + let selected = choose_items(&opts, &mut weighted_choose)?; debug!( "LoadBasedSelector select peers: {:?}, namespace: {}, opts: {:?}.",