satex_load_balancer/
lib.rs

1//!
2//! Satex 服务发现和负载均衡的库
3//!
4//! 参考实现: [`pingora-load-balancing`](https://github.com/cloudflare/pingora/tree/main/pingora-load-balancing)
5//!
6mod background;
7mod load_balancer;
8
9pub mod discovery;
10pub mod health_check;
11pub mod resolver;
12pub mod selector;
13
14pub use load_balancer::LoadBalancer;
15
16use crate::discovery::Discovery;
17use crate::health_check::HealthCheck;
18use crate::health_check::health::Health;
19use arc_swap::ArcSwap;
20use derivative::Derivative;
21use http::Extensions;
22use satex_core::Error;
23use std::collections::{BTreeSet, HashMap};
24use std::hash::{DefaultHasher, Hash, Hasher};
25use std::net::{AddrParseError, SocketAddr};
26use std::str::FromStr;
27use std::sync::Arc;
28use tokio::spawn;
29use tracing::{info, warn};
30
31/// 后端服务
32#[derive(Derivative)]
33#[derivative(Clone, Hash, PartialEq, PartialOrd, Eq, Ord, Debug)]
34pub struct Backend {
35    ///
36    /// 后端服务地址
37    ///
38    pub addr: SocketAddr,
39
40    ///
41    /// 后端服务权重,负载均衡算法会使用到该权重值.
42    ///
43    pub weight: usize,
44
45    ///
46    /// 拓展信息
47    ///
48    #[derivative(PartialEq = "ignore")]
49    #[derivative(PartialOrd = "ignore")]
50    #[derivative(Hash = "ignore")]
51    #[derivative(Ord = "ignore")]
52    pub extension: Extensions,
53}
54
55impl Backend {
56    /// 创建后端服务实例
57    pub fn new(addr: impl Into<SocketAddr>) -> Self {
58        Self::new_with_weight(addr.into(), 1)
59    }
60
61    /// 创建后端服务实例
62    pub fn new_with_weight(addr: impl Into<SocketAddr>, weight: usize) -> Self {
63        Self {
64            addr: addr.into(),
65            weight,
66            extension: Extensions::new(),
67        }
68    }
69
70    /// 计算后端服务的哈希值
71    pub(crate) fn key(&self) -> u64 {
72        let mut hasher = DefaultHasher::new();
73        self.hash(&mut hasher);
74        hasher.finish()
75    }
76}
77
78impl FromStr for Backend {
79    type Err = AddrParseError;
80
81    fn from_str(addr: &str) -> Result<Self, Self::Err> {
82        SocketAddr::from_str(addr).map(Backend::new)
83    }
84}
85
86impl<'a> TryFrom<(&'a str, usize)> for Backend {
87    type Error = AddrParseError;
88
89    fn try_from((addr, weight): (&'a str, usize)) -> Result<Self, Self::Error> {
90        let addr = SocketAddr::from_str(addr)?;
91        Ok(Backend::new_with_weight(addr, weight))
92    }
93}
94
95
96pub struct Backends {
97    discovery: Box<dyn Discovery + Send + Sync>,
98    health_check: Option<Arc<dyn HealthCheck + Send + Sync>>,
99    backends: ArcSwap<BTreeSet<Backend>>,
100    health: ArcSwap<HashMap<u64, Health>>,
101}
102
103impl Backends {
104    pub fn new(discovery: impl Discovery + Send + Sync + 'static) -> Self {
105        Self {
106            discovery: Box::new(discovery),
107            health_check: None,
108            backends: Default::default(),
109            health: Default::default(),
110        }
111    }
112
113    pub(crate) fn with_health_check(
114        self,
115        health_check: impl HealthCheck + Send + Sync + 'static,
116    ) -> Self {
117        Self {
118            health_check: Some(Arc::new(health_check)),
119            ..self
120        }
121    }
122
123    /// 当新的后端集合与当前集合不同时更新后端，
124    /// 当新的后端集合与当前集合不同时，回调函数将被调用，
125    /// 以便调用者可以相应地更新选择器。
126    fn do_update<F>(
127        &self,
128        new_backends: BTreeSet<Backend>,
129        enablement: HashMap<u64, bool>,
130        callback: F,
131    ) where
132        F: Fn(Arc<BTreeSet<Backend>>),
133    {
134        if (**self.backends.load()) != new_backends {
135            let old_health = self.health.load();
136            let mut new_health = HashMap::with_capacity(new_backends.len());
137            for backend in new_backends.iter() {
138                let key = backend.key();
139                // use the default health if the backend is new
140                let health = old_health.get(&key).cloned().unwrap_or_default();
141
142                // override enablement
143                if let Some(enabled) = enablement.get(&key) {
144                    health.enable(*enabled);
145                }
146                new_health.insert(key, health);
147            }
148
149            // 确保 `callback()` 首先执行是很重要的，因为计算选择器后端可能会很耗时。
150            // 例如，如果调用者检查 `backends` 以查看是否有可用的后端，
151            // 如果选择器尚未准备好，他们可能会遇到误报。
152            let new_backends = Arc::new(new_backends);
153            callback(new_backends.clone());
154            self.backends.store(new_backends);
155            self.health.store(Arc::new(new_health));
156        } else {
157            // no backend change, just check enablement
158            for (key, backend_enabled) in enablement.iter() {
159                // override enablement if set
160                // this get should always be Some(_) because we already populate `health`` for all known backends
161                if let Some(health) = self.health.load().get(key) {
162                    health.enable(*backend_enabled);
163                }
164            }
165        }
166    }
167
168    /// 检查指定的后端服务是否可以接收流量
169    ///
170    /// 在以下情况下返回 true:
171    /// - 后端服务既健康又启用时
172    /// - 未配置健康检查但后端服务已启用时
173    ///
174    /// 当配置了健康检查时，对于任何未知的后端服务都将返回 false
175    pub fn ready(&self, backend: &Backend) -> bool {
176        self.health
177            .load()
178            .get(&backend.key())
179            // Racing: return `None` when this function is called between the
180            // backend store and the health store
181            .map_or(self.health_check.is_none(), |h| h.ready())
182    }
183
184    /// 手动设置一个 [Backend] 是否可以接收流量。
185    ///
186    /// 此方法不会覆盖后端的健康状态。它的目的是在后端仍然健康时，停止其接受流量。
187    ///
188    /// 如果给定的后端不存在于服务发现中，此方法将无操作。
189    pub fn set_enable(&self, backend: &Backend, enabled: bool) {
190        // this should always be Some(_) because health is always populated during update
191        if let Some(health) = self.health.load().get(&backend.key()) {
192            health.enable(enabled)
193        };
194    }
195
196    /// Return the collection of the backends.
197    pub fn items(&self) -> Arc<BTreeSet<Backend>> {
198        self.backends.load_full()
199    }
200
201    /// Call the service discovery method to update the collection of backends.
202    ///
203    /// The callback will be invoked when the new set of backend is different
204    /// from the current one so that the caller can update the selector accordingly.
205    pub async fn update<F>(&self, callback: F) -> Result<(), Error>
206    where
207        F: Fn(Arc<BTreeSet<Backend>>),
208    {
209        let (new_backends, enablement) = self.discovery.discover().await?;
210        self.do_update(new_backends, enablement, callback);
211        Ok(())
212    }
213
214    /// Run health check on all backends if it is set.
215    ///
216    /// When `parallel: true`, all backends are checked in parallel instead of sequentially
217    pub async fn run_health_check(&self, parallel: bool) {
218        use crate::health_check::HealthCheck;
219
220        async fn check_and_report(
221            backend: &Backend,
222            check: &Arc<dyn HealthCheck + Send + Sync>,
223            health_table: &HashMap<u64, Health>,
224        ) {
225            let errored = check.check(backend).await.err();
226            if let Some(health) = health_table.get(&backend.key()) {
227                let flipped =
228                    health.observe_health(errored.is_none(), check.threshold(errored.is_none()));
229                if flipped {
230                    check.status_change(backend, errored.is_none()).await;
231                    if let Some(e) = errored {
232                        warn!("{backend:?} becomes unhealthy, {e}");
233                    } else {
234                        info!("{backend:?} becomes healthy");
235                    }
236                }
237            }
238        }
239
240        let Some(health_check) = self.health_check.as_ref() else {
241            return;
242        };
243
244        let backends = self.backends.load();
245        if parallel {
246            let health_table = self.health.load_full();
247            let jobs = backends.iter().map(|backend| {
248                let backend = backend.clone();
249                let check = health_check.clone();
250                let ht = health_table.clone();
251                spawn(async move {
252                    check_and_report(&backend, &check, &ht).await;
253                })
254            });
255
256            futures::future::join_all(jobs).await;
257        } else {
258            for backend in backends.iter() {
259                check_and_report(backend, health_check, &self.health.load()).await;
260            }
261        }
262    }
263}
satex_load_balancer/lib.rs

satex_load_balancer/
lib.rs
