diff --git a/pingora-load-balancing/src/lib.rs b/pingora-load-balancing/src/lib.rs index d2c1e34ba..63106ada8 100644 --- a/pingora-load-balancing/src/lib.rs +++ b/pingora-load-balancing/src/lib.rs @@ -305,7 +305,7 @@ where /// Build a [LoadBalancer] with the given [Backends]. pub fn from_backends(backends: Backends) -> Self { - let selector = ArcSwap::new(Arc::new(S::build(&backends.get_backend()))); + let selector = ArcSwap::new(Arc::new(S::build(&backends.get_backend(), None))); LoadBalancer { backends, selector, @@ -321,8 +321,10 @@ where /// is running as a background service. pub async fn update(&self) -> Result<()> { if self.backends.update().await? { - self.selector - .store(Arc::new(S::build(&self.backends.get_backend()))) + self.selector.store(Arc::new(S::build( + &self.backends.get_backend(), + Some(self.selector.load_full()), + ))) } Ok(()) } diff --git a/pingora-load-balancing/src/selection/consistent.rs b/pingora-load-balancing/src/selection/consistent.rs index 9c627260c..e8bd8d8d0 100644 --- a/pingora-load-balancing/src/selection/consistent.rs +++ b/pingora-load-balancing/src/selection/consistent.rs @@ -29,7 +29,7 @@ pub struct KetamaHashing { impl BackendSelection for KetamaHashing { type Iter = OwnedNodeIterator; - fn build(backends: &BTreeSet) -> Self { + fn build(backends: &BTreeSet, _previous: Option>) -> Self { let buckets: Vec<_> = backends .iter() .filter_map(|b| { @@ -84,7 +84,7 @@ mod test { let b2 = Backend::new("1.0.0.1:80").unwrap(); let b3 = Backend::new("1.0.0.255:80").unwrap(); let backends = BTreeSet::from_iter([b1.clone(), b2.clone(), b3.clone()]); - let hash = Arc::new(KetamaHashing::build(&backends)); + let hash = Arc::new(KetamaHashing::build(&backends, None)); let mut iter = hash.iter(b"test0"); assert_eq!(iter.next(), Some(&b2)); @@ -109,7 +109,7 @@ mod test { // remove b3 let backends = BTreeSet::from_iter([b1.clone(), b2.clone()]); - let hash = Arc::new(KetamaHashing::build(&backends)); + let hash = Arc::new(KetamaHashing::build(&backends, None)); let mut iter = hash.iter(b"test0"); assert_eq!(iter.next(), Some(&b2)); let mut iter = hash.iter(b"test1"); diff --git a/pingora-load-balancing/src/selection/mod.rs b/pingora-load-balancing/src/selection/mod.rs index 15e29adaf..ec36e4a44 100644 --- a/pingora-load-balancing/src/selection/mod.rs +++ b/pingora-load-balancing/src/selection/mod.rs @@ -28,7 +28,8 @@ pub trait BackendSelection { /// The [BackendIter] returned from iter() below. type Iter: BackendIter; /// The function to create a [BackendSelection] implementation. - fn build(backends: &BTreeSet) -> Self; + /// `previous` is previous instance of self, when it's created during backend update + fn build(backends: &BTreeSet, previous: Option>) -> Self; /// Select backends for a given key. /// /// An [BackendIter] should be returned. The first item in the iter is the first diff --git a/pingora-load-balancing/src/selection/weighted.rs b/pingora-load-balancing/src/selection/weighted.rs index 3f37de605..ea1ca45b6 100644 --- a/pingora-load-balancing/src/selection/weighted.rs +++ b/pingora-load-balancing/src/selection/weighted.rs @@ -32,7 +32,7 @@ pub struct Weighted { impl BackendSelection for Weighted { type Iter = WeightedIterator; - fn build(backends: &BTreeSet) -> Self { + fn build(backends: &BTreeSet, _previous: Option>) -> Self { assert!( backends.len() <= u16::MAX as usize, "support up to 2^16 backends" @@ -113,7 +113,7 @@ mod test { b2.weight = 10; // 10x than the rest let b3 = Backend::new("1.0.0.255:80").unwrap(); let backends = BTreeSet::from_iter([b1.clone(), b2.clone(), b3.clone()]); - let hash: Arc = Arc::new(Weighted::build(&backends)); + let hash: Arc = Arc::new(Weighted::build(&backends, None)); // same hash iter over let mut iter = hash.iter(b"test"); @@ -155,7 +155,7 @@ mod test { b2.weight = 8; // 8x than the rest let b3 = Backend::new("1.0.0.255:80").unwrap(); let backends = BTreeSet::from_iter([b1.clone(), b2.clone(), b3.clone()]); - let hash: Arc> = Arc::new(Weighted::build(&backends)); + let hash: Arc> = Arc::new(Weighted::build(&backends, None)); // same hash iter over let mut iter = hash.iter(b"test"); @@ -191,7 +191,7 @@ mod test { b2.weight = 8; // 8x than the rest let b3 = Backend::new("1.0.0.255:80").unwrap(); let backends = BTreeSet::from_iter([b1.clone(), b2.clone(), b3.clone()]); - let hash: Arc> = Arc::new(Weighted::build(&backends)); + let hash: Arc> = Arc::new(Weighted::build(&backends, None)); let mut count = HashMap::new(); count.insert(b1.clone(), 0);