139 lines
3.5 KiB
Rust
139 lines
3.5 KiB
Rust
//! Random load balancing policy
|
|
|
|
use std::sync::Arc;
|
|
|
|
use async_trait::async_trait;
|
|
use rand::Rng;
|
|
|
|
use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
|
|
use crate::core::Worker;
|
|
|
|
/// 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
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl LoadBalancingPolicy for RandomPolicy {
|
|
async fn select_worker(
|
|
&self,
|
|
workers: &[Arc<dyn Worker>],
|
|
_info: &SelectWorkerInfo<'_>,
|
|
) -> 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());
|
|
|
|
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};
|
|
|
|
#[tokio::test]
|
|
async 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, &SelectWorkerInfo::default())
|
|
.await
|
|
{
|
|
*counts.entry(idx).or_insert(0) += 1;
|
|
}
|
|
}
|
|
|
|
assert_eq!(counts.len(), 3);
|
|
assert!(counts.values().all(|&count| count > 0));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async 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(),
|
|
),
|
|
];
|
|
|
|
workers[0].set_healthy(false);
|
|
|
|
for _ in 0..10 {
|
|
assert_eq!(
|
|
policy
|
|
.select_worker(&workers, &SelectWorkerInfo::default())
|
|
.await,
|
|
Some(1)
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async 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, &SelectWorkerInfo::default())
|
|
.await,
|
|
None
|
|
);
|
|
}
|
|
}
|