random.rs 3.33 KB
Newer Older
1
2
3
4
//! Random load balancing policy

use super::{get_healthy_worker_indices, LoadBalancingPolicy};
use crate::core::Worker;
5
use crate::metrics::RouterMetrics;
6
use rand::Rng;
7
use std::sync::Arc;
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

/// 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,
24
        workers: &[Arc<dyn Worker>],
25
26
27
28
29
30
31
32
        _request_text: Option<&str>,
    ) -> Option<usize> {
        let healthy_indices = get_healthy_worker_indices(workers);

        if healthy_indices.is_empty() {
            return None;
        }

33
34
        let mut rng = rand::rng();
        let random_idx = rng.random_range(0..healthy_indices.len());
35
36
37
38
        let worker = workers[healthy_indices[random_idx]].url();

        RouterMetrics::record_processed_request(worker);
        RouterMetrics::record_policy_decision(self.name(), worker);
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
        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 super::*;
    use crate::core::{BasicWorker, WorkerType};
    use std::collections::HashMap;

    #[test]
    fn test_random_selection() {
        let policy = RandomPolicy::new();
60
61
        let workers: Vec<Arc<dyn Worker>> = vec![
            Arc::new(BasicWorker::new(
62
63
64
                "http://w1:8000".to_string(),
                WorkerType::Regular,
            )),
65
            Arc::new(BasicWorker::new(
66
67
68
                "http://w2:8000".to_string(),
                WorkerType::Regular,
            )),
69
            Arc::new(BasicWorker::new(
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
                "http://w3:8000".to_string(),
                WorkerType::Regular,
            )),
        ];

        // Test multiple selections to ensure randomness
        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();
91
92
        let workers: Vec<Arc<dyn Worker>> = vec![
            Arc::new(BasicWorker::new(
93
94
95
                "http://w1:8000".to_string(),
                WorkerType::Regular,
            )),
96
            Arc::new(BasicWorker::new(
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
                "http://w2:8000".to_string(),
                WorkerType::Regular,
            )),
        ];

        // 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();
114
        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(BasicWorker::new(
115
116
117
118
119
120
121
122
            "http://w1:8000".to_string(),
            WorkerType::Regular,
        ))];

        workers[0].set_healthy(false);
        assert_eq!(policy.select_worker(&workers, None), None);
    }
}