From 3d9373f58f64ecb90cdd760d4f66e157d1c4e750 Mon Sep 17 00:00:00 2001
From: jizhuozhi <jizhuozhi.george@gmail.com>
Date: Fri, 1 Mar 2024 11:37:47 +0800
Subject: [PATCH] opt: use gcd to reduce memory overhead

---
 .../src/selection/weighted.rs                 | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/pingora-load-balancing/src/selection/weighted.rs b/pingora-load-balancing/src/selection/weighted.rs
index 3f37de60..d962bad9 100644
--- a/pingora-load-balancing/src/selection/weighted.rs
+++ b/pingora-load-balancing/src/selection/weighted.rs
@@ -29,6 +29,16 @@ pub struct Weighted<H = FnvHasher> {
     algorithm: H,
 }
 
+fn gcd(mut a: usize, mut b: usize) -> usize {
+    let mut r;
+    while b != 0 {
+        r = a % b;
+        a = b;
+        b = r;
+    }
+    a
+}
+
 impl<H: SelectionAlgorithm> BackendSelection for Weighted<H> {
     type Iter = WeightedIterator<H>;
 
@@ -38,9 +48,16 @@ impl<H: SelectionAlgorithm> BackendSelection for Weighted<H> {
             "support up to 2^16 backends"
         );
         let backends = Vec::from_iter(backends.iter().cloned()).into_boxed_slice();
-        let mut weighted = Vec::with_capacity(backends.len());
+        let mut g = 0;
+        let mut total = 0;
+        // use gcd to reduce the memory overhead
+        for (_, b) in backends.iter().enumerate() {
+            g = gcd(g, b.weight);
+            total += b.weight;
+        }
+        let mut weighted = Vec::with_capacity(total / g);
         for (index, b) in backends.iter().enumerate() {
-            for _ in 0..b.weight {
+            for _ in 0..(b.weight / g) {
                 weighted.push(index as u16);
             }
         }