Skip to content

Commit

Permalink
opt: use gcd to reduce memory overhead
Browse files Browse the repository at this point in the history
  • Loading branch information
jizhuozhi committed Mar 1, 2024
1 parent 8797329 commit 88bf007
Showing 1 changed file with 16 additions and 1 deletion.
17 changes: 16 additions & 1 deletion pingora-load-balancing/src/selection/weighted.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>;

Expand All @@ -39,8 +49,13 @@ impl<H: SelectionAlgorithm> BackendSelection for Weighted<H> {
);
let backends = Vec::from_iter(backends.iter().cloned()).into_boxed_slice();
let mut weighted = Vec::with_capacity(backends.len());
let mut g = 0;
// use gcd to reduce the memory overhead
for (_, b) in backends.iter().enumerate() {
g = gcd(g, b.weight)
}
for (index, b) in backends.iter().enumerate() {
for _ in 0..b.weight {
for _ in 0..(b.weight / g) {
weighted.push(index as u16);
}
}
Expand Down

0 comments on commit 88bf007

Please sign in to comment.