130 lines
3.5 KiB
Rust
130 lines
3.5 KiB
Rust
//! Random load balancing policy
|
|
|
|
use std::sync::Arc;
|
|
|
|
use rand::Rng;
|
|
|
|
use super::{get_healthy_worker_indices, LoadBalancingPolicy};
|
|
use crate::{core::Worker, metrics::RouterMetrics};
|
|
|
|
/// Random selection policy
|
|
///
|
|
/// Selects workers randomly with uniform distribution among healthy workers.
|
|
#[derive(Debug, Default)]
|
|
pub struct RandomPolicy;
|
|
|
|
impl RandomPolicy {
|
|
pub fn new() -> Self {
|
|
Self
|
|
}
|
|
}
|
|
|
|
impl LoadBalancingPolicy for RandomPolicy {
|
|
fn select_worker(
|
|
&self,
|
|
workers: &[Arc<dyn Worker>],
|
|
_request_text: Option<&str>,
|
|
) -> Option<usize> {
|
|
let healthy_indices = get_healthy_worker_indices(workers);
|
|
|
|
if healthy_indices.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
let mut rng = rand::rng();
|
|
let random_idx = rng.random_range(0..healthy_indices.len());
|
|
let worker = workers[healthy_indices[random_idx]].url();
|
|
|
|
RouterMetrics::record_processed_request(worker);
|
|
RouterMetrics::record_policy_decision(self.name(), worker);
|
|
Some(healthy_indices[random_idx])
|
|
}
|
|
|
|
fn name(&self) -> &'static str {
|
|
"random"
|
|
}
|
|
|
|
fn as_any(&self) -> &dyn std::any::Any {
|
|
self
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::collections::HashMap;
|
|
|
|
use super::*;
|
|
use crate::core::{BasicWorkerBuilder, WorkerType};
|
|
|
|
#[test]
|
|
fn test_random_selection() {
|
|
let policy = RandomPolicy::new();
|
|
let workers: Vec<Arc<dyn Worker>> = vec![
|
|
Arc::new(
|
|
BasicWorkerBuilder::new("http://w1:8000")
|
|
.worker_type(WorkerType::Regular)
|
|
.build(),
|
|
),
|
|
Arc::new(
|
|
BasicWorkerBuilder::new("http://w2:8000")
|
|
.worker_type(WorkerType::Regular)
|
|
.build(),
|
|
),
|
|
Arc::new(
|
|
BasicWorkerBuilder::new("http://w3:8000")
|
|
.worker_type(WorkerType::Regular)
|
|
.build(),
|
|
),
|
|
];
|
|
|
|
let mut counts = HashMap::new();
|
|
for _ in 0..100 {
|
|
if let Some(idx) = policy.select_worker(&workers, None) {
|
|
*counts.entry(idx).or_insert(0) += 1;
|
|
}
|
|
}
|
|
|
|
// All workers should be selected at least once
|
|
assert_eq!(counts.len(), 3);
|
|
assert!(counts.values().all(|&count| count > 0));
|
|
}
|
|
|
|
#[test]
|
|
fn test_random_with_unhealthy_workers() {
|
|
let policy = RandomPolicy::new();
|
|
let workers: Vec<Arc<dyn Worker>> = vec![
|
|
Arc::new(
|
|
BasicWorkerBuilder::new("http://w1:8000")
|
|
.worker_type(WorkerType::Regular)
|
|
.build(),
|
|
),
|
|
Arc::new(
|
|
BasicWorkerBuilder::new("http://w2:8000")
|
|
.worker_type(WorkerType::Regular)
|
|
.build(),
|
|
),
|
|
];
|
|
|
|
// Mark first worker as unhealthy
|
|
workers[0].set_healthy(false);
|
|
|
|
// Should always select the healthy worker (index 1)
|
|
for _ in 0..10 {
|
|
assert_eq!(policy.select_worker(&workers, None), Some(1));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_random_no_healthy_workers() {
|
|
let policy = RandomPolicy::new();
|
|
let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
|
|
BasicWorkerBuilder::new("http://w1:8000")
|
|
.worker_type(WorkerType::Regular)
|
|
.build(),
|
|
)];
|
|
|
|
workers[0].set_healthy(false);
|
|
assert_eq!(policy.select_worker(&workers, None), None);
|
|
}
|
|
}
|