satex_load_balancer/
load_balancer.rs

1use crate::health_check::HealthCheck;
2use crate::selector::{BackendIter, BoxSelector, Selector};
3use crate::{Backend, Backends};
4use satex_core::Error;
5use std::time::Duration;
6
7/// 负载均衡器
8pub struct LoadBalancer {
9    selector: BoxSelector,
10    pub(crate) backends: Backends,
11    pub(crate) update_frequency: Option<Duration>,
12    pub(crate) health_check_frequency: Option<Duration>,
13    pub(crate) health_check_parallel: bool,
14}
15
16impl LoadBalancer {
17    /// 创建负载均衡器实例
18    pub fn new<S>(backends: Backends, selector: S) -> Self
19    where
20        S: Selector + Send + Sync + 'static,
21        S::Iter: Send + Sync + 'static,
22    {
23        let selector = BoxSelector::new(selector);
24        Self {
25            backends,
26            selector,
27            update_frequency: None,
28            health_check_frequency: None,
29            health_check_parallel: false,
30        }
31    }
32
33    /// 设置健康检查
34    pub fn with_health_check(
35        mut self,
36        health_check: impl HealthCheck + Send + Sync + 'static,
37    ) -> Self {
38        self.backends = self.backends.with_health_check(health_check);
39        self
40    }
41
42    /// 设置更新频率
43    pub fn with_update_frequency(mut self, update_frequency: Duration) -> Self {
44        self.update_frequency = Some(update_frequency);
45        self
46    }
47
48    /// 设置健康检查的频率
49    pub fn with_health_check_frequency(mut self, health_check_frequency: Duration) -> Self {
50        self.health_check_frequency = Some(health_check_frequency);
51        self
52    }
53
54    pub fn with_health_check_parallel(mut self, health_check_parallel: bool) -> Self {
55        self.health_check_parallel = health_check_parallel;
56        self
57    }
58
59    /// 运行服务发现并更新选择算法。
60    ///
61    /// 如果这个 [LoadBalancer] 实例作为后台服务运行，此函数将每隔 `update_frequency` 被调用一次。
62    pub async fn update(&self) -> Result<(), Error> {
63        self.backends
64            .update(|backends| self.selector.update(&backends))
65            .await
66    }
67
68    /// 根据选择算法和健康检查结果返回第一个健康的 [Backend]。
69    ///
70    /// `key` 用于基于哈希的选择，如果选择是随机或轮询，则忽略此参数。
71    ///
72    /// `max_iterations` 用于限制搜索下一个 Backend 的时间。在某些算法中，
73    /// 如 Ketama 哈希，搜索下一个后端是线性的，可能需要很多步骤。
74    pub fn select(&self, key: &[u8]) -> Option<Backend> {
75        self.select_with(key, |_, health| health)
76    }
77
78    /// 类似于 [Self::select]，根据选择算法和用户定义的 `accept` 函数返回第一个健康的 [Backend]。
79    ///
80    /// `accept` 函数接受两个输入：正在选择的后端和该后端的内部健康状态。该函数可以执行一些操作，
81    /// 比如忽略内部健康检查或因为之前失败而跳过这个后端。`accept` 函数会被多次调用，遍历后端，
82    /// 直到返回 `true`。
83    pub fn select_with<F>(&self, key: &[u8], accept: F) -> Option<Backend>
84    where
85        F: Fn(&Backend, bool) -> bool,
86    {
87        let mut iter = self.selector.iter(key);
88        while let Some(b) = iter.next() {
89            if accept(b, self.backends.ready(b)) {
90                return Some(b.clone());
91            }
92        }
93        None
94    }
95}
satex_load_balancer/load_balancer.rs

satex_load_balancer/
load_balancer.rs
