feat: add previous parameter to BackendSelection::build to keep state…

… between updates

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(())
     }

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<Backend>) -> Self {
+    fn build(backends: &BTreeSet<Backend>, _previous: Option<Arc<Self>>) -> 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");

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<Backend>) -> Self;
+    /// `previous` is previous instance of self, when it's created during backend update 
+    fn build(backends: &BTreeSet<Backend>, previous: Option<Arc<Self>>) -> Self;
     /// Select backends for a given key.
     ///
     /// An [BackendIter] should be returned. The first item in the iter is the first

pingora-load-balancing/src/selection/weighted.rs

@@ -32,7 +32,7 @@ pub struct Weighted<H = FnvHasher> {
 impl<H: SelectionAlgorithm> BackendSelection for Weighted<H> {
     type Iter = WeightedIterator<H>;
 
-    fn build(backends: &BTreeSet<Backend>) -> Self {
+    fn build(backends: &BTreeSet<Backend>, _previous: Option<Arc<Self>>) -> 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<Weighted> = Arc::new(Weighted::build(&backends));
+        let hash: Arc<Weighted> = 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<Weighted<RoundRobin>> = Arc::new(Weighted::build(&backends));
+        let hash: Arc<Weighted<RoundRobin>> = 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<Weighted<Random>> = Arc::new(Weighted::build(&backends));
+        let hash: Arc<Weighted<Random>> = Arc::new(Weighted::build(&backends, None));
 
         let mut count = HashMap::new();
         count.insert(b1.clone(), 0);