fix: lb random range bug

src/backend/load_balance.rs

@@ -1,4 +1,5 @@
 use derive_builder::Builder;
+use rand::Rng;
 use std::sync::{
   atomic::{AtomicUsize, Ordering},
   Arc,
@@ -13,50 +14,85 @@ pub(super) mod load_balance_options {
 }
 
 #[derive(Debug, Clone, Builder)]
-/// Counter object as a pointer to the current serving upstream destination
-pub struct LbRoundRobinCount {
+/// Round Robin LB object as a pointer to the current serving upstream destination
+pub struct LbRoundRobin {
   #[builder(default)]
-  cnt: Arc<AtomicUsize>,
+  /// Pointer to the index of the last served upstream destination
+  ptr: Arc<AtomicUsize>,
   #[builder(setter(custom), default)]
-  max_val: usize,
+  /// Number of upstream destinations
+  num_upstreams: usize,
 }
-impl LbRoundRobinCountBuilder {
-  pub fn max_val(&mut self, v: &usize) -> &mut Self {
-    self.max_val = Some(*v);
+impl LbRoundRobinBuilder {
+  pub fn num_upstreams(&mut self, v: &usize) -> &mut Self {
+    self.num_upstreams = Some(*v);
     self
   }
 }
-impl LbRoundRobinCount {
+impl LbRoundRobin {
   /// Get a current count of upstream served
-  fn current_cnt(&self) -> usize {
-    self.cnt.load(Ordering::Relaxed)
+  fn current_ptr(&self) -> usize {
+    self.ptr.load(Ordering::Relaxed)
   }
 
   /// Increment the count of upstream served up to the max value
-  pub fn increment_cnt(&self) -> usize {
-    if self.current_cnt() < self.max_val - 1 {
-      self.cnt.fetch_add(1, Ordering::Relaxed)
+  pub fn increment_ptr(&self) -> usize {
+    if self.current_ptr() < self.num_upstreams - 1 {
+      self.ptr.fetch_add(1, Ordering::Relaxed)
     } else {
       // Clear the counter
-      self.cnt.fetch_and(0, Ordering::Relaxed)
+      self.ptr.fetch_and(0, Ordering::Relaxed)
     }
   }
 }
 
+#[derive(Debug, Clone, Builder)]
+/// Random LB object to keep the object of random pools
+pub struct LbRandom {
+  #[builder(setter(custom), default)]
+  /// Number of upstream destinations
+  num_upstreams: usize,
+}
+impl LbRandomBuilder {
+  pub fn num_upstreams(&mut self, v: &usize) -> &mut Self {
+    self.num_upstreams = Some(*v);
+    self
+  }
+}
+impl LbRandom {
+  /// Returns the random index within the range
+  pub fn get_ptr(&self) -> usize {
+    let mut rng = rand::thread_rng();
+    rng.gen_range(0..self.num_upstreams)
+  }
+}
+
 #[derive(Debug, Clone)]
 /// Load Balancing Option
 pub enum LoadBalance {
   /// Fix to the first upstream. Use if only one upstream destination is specified
   FixToFirst,
   /// Randomly chose one upstream server
-  Random,
+  Random(LbRandom),
   /// Simple round robin without session persistance
-  RoundRobin(LbRoundRobinCount),
+  RoundRobin(LbRoundRobin),
   /// Round robin with session persistance using cookie
-  StickyRoundRobin(LbRoundRobinCount),
+  StickyRoundRobin(LbRoundRobin),
 }
 impl Default for LoadBalance {
   fn default() -> Self {
     Self::FixToFirst
   }
 }
+
+impl LoadBalance {
+  /// Get the index of the upstream serving the incoming request
+  pub(super) fn get_idx(&self) -> usize {
+    match self {
+      LoadBalance::FixToFirst => 0usize,
+      LoadBalance::RoundRobin(ptr) => ptr.increment_ptr(),
+      LoadBalance::Random(v) => v.get_ptr(),
+      LoadBalance::StickyRoundRobin(_ptr) => 0usize, // todo!(), // TODO: TODO: TODO: TODO: tentative value
+    }
+  }
+}