From 1ba9ff25bc1f59d5ce1c6a30f162e0aad4bba76b Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 31 Oct 2024 15:47:36 +0800 Subject: [PATCH 1/2] chore: minor refactor for weighted choose --- 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 +- src/meta-srv/src/selector/weighted_choose.rs | 5 +++ 4 files changed, 39 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: {:?}.", diff --git a/src/meta-srv/src/selector/weighted_choose.rs b/src/meta-srv/src/selector/weighted_choose.rs index 9fe8d7b28d14..f8be7f299331 100644 --- a/src/meta-srv/src/selector/weighted_choose.rs +++ b/src/meta-srv/src/selector/weighted_choose.rs @@ -92,6 +92,11 @@ where } fn choose_multiple(&mut self, amount: usize) -> Result> { + if amount >= self.items.len() { + // fast path + return Ok(self.items.iter().map(|item| item.item.clone()).collect()); + } + Ok(self .items .choose_multiple_weighted(&mut thread_rng(), amount, |item| item.weight as f64) From 2c1903ff3134a7296ccf75cb5e0fc081220090c3 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 1 Nov 2024 15:14:52 +0800 Subject: [PATCH 2/2] chore: by comment, remove the fast path of choose_multiple --- src/meta-srv/src/selector/weighted_choose.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/meta-srv/src/selector/weighted_choose.rs b/src/meta-srv/src/selector/weighted_choose.rs index f8be7f299331..9fe8d7b28d14 100644 --- a/src/meta-srv/src/selector/weighted_choose.rs +++ b/src/meta-srv/src/selector/weighted_choose.rs @@ -92,11 +92,6 @@ where } fn choose_multiple(&mut self, amount: usize) -> Result> { - if amount >= self.items.len() { - // fast path - return Ok(self.items.iter().map(|item| item.item.clone()).collect()); - } - Ok(self .items .choose_multiple_weighted(&mut thread_rng(), amount, |item| item.weight as f64)