rust_web_server/proxy_config/
upstream.rs1use std::collections::HashMap;
11use std::sync::{Arc, Mutex, OnceLock, RwLock};
12use std::time::Instant;
13
14use crate::circuit_breaker::{BreakerState, CircuitBreaker};
15
16#[derive(Debug, Clone, PartialEq)]
18pub struct StickyConfig {
19 pub cookie: String,
21 pub max_age_secs: Option<u64>,
23}
24
25#[derive(Debug, Clone, PartialEq)]
27pub struct BreakerConfig {
28 pub failure_threshold: u32,
30 pub recovery_secs: u64,
32 pub half_open_probes: u32,
34}
35
36#[derive(Debug, Clone, PartialEq)]
38pub struct RetryConfig {
39 pub attempts: u32,
41 pub status: Vec<i16>,
43}
44
45pub(crate) struct UpstreamPolicy {
47 pub(crate) name: String,
48 live: Arc<RwLock<Vec<String>>>,
50 weights: HashMap<String, u32>,
51 pub(crate) sticky: Option<StickyConfig>,
52 breaker: Option<Mutex<CircuitBreaker>>,
53 pub(crate) retry: Option<RetryConfig>,
54 current: Mutex<HashMap<String, i64>>,
56}
57
58impl UpstreamPolicy {
59 pub(crate) fn new(
60 name: &str,
61 live: Arc<RwLock<Vec<String>>>,
62 backends: &[String],
63 weights: &[u32],
64 sticky: Option<StickyConfig>,
65 breaker: Option<&BreakerConfig>,
66 retry: Option<RetryConfig>,
67 ) -> Self {
68 UpstreamPolicy {
69 name: name.to_string(),
70 live,
71 weights: backends.iter().enumerate().map(|(i, b)| (b.clone(), weights.get(i).copied().unwrap_or(1))).collect(),
73 sticky,
74 breaker: breaker.map(|b| Mutex::new(CircuitBreaker::new(b.failure_threshold.max(1), b.recovery_secs).max_half_open_probes(b.half_open_probes))),
75 retry,
76 current: Mutex::new(HashMap::new()),
77 }
78 }
79
80 pub(crate) fn weight(&self, backend: &str) -> u32 {
82 self.weights.get(backend).copied().unwrap_or(1).max(1)
83 }
84
85 pub(crate) fn weighted(&self) -> bool {
87 self.weights.values().any(|w| *w != 1)
88 }
89
90 pub(crate) fn weighted_round_robin(&self, candidates: &[String]) -> Option<String> {
94 let mut current = self.current.lock().unwrap_or_else(|e| e.into_inner());
95 let total: i64 = candidates.iter().map(|b| self.weight(b) as i64).sum();
96 let mut best: Option<(&String, i64)> = None;
97 for backend in candidates {
98 let value = current.entry(backend.clone()).or_insert(0);
99 *value += self.weight(backend) as i64;
100 if best.is_none_or(|(_, top)| *value > top) {
101 best = Some((backend, *value));
102 }
103 }
104 let (chosen, _) = best?;
105 *current.get_mut(chosen.as_str()).expect("just inserted") -= total;
106 Some(chosen.clone())
107 }
108
109 pub(crate) fn weighted_index(&self, candidates: &[String], roll: usize) -> usize {
111 let total: usize = candidates.iter().map(|b| self.weight(b) as usize).sum();
112 let mut point = roll % total.max(1);
113 for (i, backend) in candidates.iter().enumerate() {
114 let weight = self.weight(backend) as usize;
115 if point < weight {
116 return i;
117 }
118 point -= weight;
119 }
120 0
121 }
122
123 pub(crate) fn admits(&self, backend: &str) -> bool {
126 match &self.breaker {
127 Some(breaker) => breaker.lock().unwrap_or_else(|e| e.into_inner()).is_available(backend),
128 None => true,
129 }
130 }
131
132 pub(crate) fn any_open(&self, backends: &[String]) -> bool {
134 self.breaker.as_ref().is_some_and(|breaker| {
135 let breaker = breaker.lock().unwrap_or_else(|e| e.into_inner());
136 backends.iter().any(|b| breaker.state(b) != BreakerState::Closed)
137 })
138 }
139
140 pub(crate) fn record(&self, backend: &str, failed: bool) {
141 if let Some(breaker) = &self.breaker {
142 let mut breaker = breaker.lock().unwrap_or_else(|e| e.into_inner());
143 if failed {
144 breaker.record_failure(backend);
145 } else {
146 breaker.record_success(backend);
147 }
148 }
149 }
150
151 fn breaker_state(&self, backend: &str) -> Option<BreakerState> {
152 self.breaker.as_ref().map(|b| b.lock().unwrap_or_else(|e| e.into_inner()).state(backend))
153 }
154
155 pub(crate) fn retries_status(&self, status: i16) -> bool {
157 self.retry.as_ref().is_some_and(|r| r.status.contains(&status))
158 }
159}
160
161pub(crate) fn backend_id(backend: &str) -> String {
164 crate::core::sha256::sha256(backend.as_bytes())[..6].iter().map(|b| format!("{b:02x}")).collect()
165}
166
167pub(crate) fn cookie_value(request: &crate::request::Request, name: &str) -> Option<String> {
169 request
170 .headers
171 .iter()
172 .filter(|h| h.name.eq_ignore_ascii_case("cookie"))
173 .flat_map(|h| h.value.split(';').map(str::trim).map(str::to_string).collect::<Vec<_>>())
174 .find_map(|pair| pair.split_once('=').filter(|(n, _)| *n == name).map(|(_, v)| v.to_string()))
175}
176
177pub(crate) fn replayable(request: &crate::request::Request) -> bool {
180 matches!(request.method.as_str(), "GET" | "HEAD" | "OPTIONS" | "PUT" | "DELETE" | "TRACE") && !request.body.is_stream()
181}
182
183const BUCKETS: [f64; 11] = [0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0];
186
187#[derive(Default)]
188struct BackendStats {
189 requests: HashMap<String, u64>,
191 buckets: [u64; 11],
192 sum: f64,
193 count: u64,
194}
195
196#[derive(Default)]
197struct Stats {
198 backends: HashMap<(String, String), BackendStats>,
200 retries: HashMap<String, u64>,
201 policies: Vec<Arc<UpstreamPolicy>>,
203}
204
205fn stats() -> &'static Mutex<Stats> {
206 static STATS: OnceLock<Mutex<Stats>> = OnceLock::new();
207 STATS.get_or_init(Default::default)
208}
209
210pub(crate) fn register(policies: Vec<Arc<UpstreamPolicy>>) {
212 stats().lock().unwrap_or_else(|e| e.into_inner()).policies = policies;
213}
214
215pub(crate) fn observe(upstream: &str, backend: &str, status: Option<i16>, started: Instant) {
218 let secs = started.elapsed().as_secs_f64();
219 let mut stats = stats().lock().unwrap_or_else(|e| e.into_inner());
220 let entry = stats.backends.entry((upstream.to_string(), backend.to_string())).or_default();
221 *entry.requests.entry(status.map_or_else(|| "error".to_string(), |s| s.to_string())).or_default() += 1;
222 for (i, upper) in BUCKETS.iter().enumerate() {
223 if secs <= *upper {
224 entry.buckets[i] += 1;
225 }
226 }
227 entry.sum += secs;
228 entry.count += 1;
229}
230
231pub(crate) fn observe_retry(upstream: &str) {
232 *stats().lock().unwrap_or_else(|e| e.into_inner()).retries.entry(upstream.to_string()).or_default() += 1;
233}
234
235fn label(value: &str) -> String {
236 value.replace('\\', "\\\\").replace('"', "\\\"").replace('\n', "\\n")
237}
238
239pub(crate) fn prometheus_text() -> String {
241 let stats = stats().lock().unwrap_or_else(|e| e.into_inner());
242 let mut out = String::new();
243 if !stats.backends.is_empty() {
244 let mut keys: Vec<&(String, String)> = stats.backends.keys().collect();
245 keys.sort();
246 out.push_str("# HELP rws_upstream_requests_total Requests sent to each upstream backend, by response status (\"error\": no response)\n");
247 out.push_str("# TYPE rws_upstream_requests_total counter\n");
248 for key in &keys {
249 let entry = &stats.backends[key];
250 let mut codes: Vec<&String> = entry.requests.keys().collect();
251 codes.sort();
252 for code in codes {
253 out.push_str(&format!(
254 "rws_upstream_requests_total{{upstream=\"{}\",backend=\"{}\",code=\"{}\"}} {}\n",
255 label(&key.0), label(&key.1), code, entry.requests[code]
256 ));
257 }
258 }
259 out.push_str("# HELP rws_upstream_request_duration_seconds Time from sending a request to an upstream backend to its response head\n");
260 out.push_str("# TYPE rws_upstream_request_duration_seconds histogram\n");
261 for key in &keys {
262 let entry = &stats.backends[key];
263 let labels = format!("upstream=\"{}\",backend=\"{}\"", label(&key.0), label(&key.1));
264 for (i, upper) in BUCKETS.iter().enumerate() {
265 out.push_str(&format!("rws_upstream_request_duration_seconds_bucket{{{labels},le=\"{upper}\"}} {}\n", entry.buckets[i]));
266 }
267 out.push_str(&format!("rws_upstream_request_duration_seconds_bucket{{{labels},le=\"+Inf\"}} {}\n", entry.count));
268 out.push_str(&format!("rws_upstream_request_duration_seconds_sum{{{labels}}} {:.6}\n", entry.sum));
269 out.push_str(&format!("rws_upstream_request_duration_seconds_count{{{labels}}} {}\n", entry.count));
270 }
271 }
272 if !stats.retries.is_empty() {
273 let mut names: Vec<&String> = stats.retries.keys().collect();
274 names.sort();
275 out.push_str("# HELP rws_upstream_retries_total Requests sent again to another backend\n");
276 out.push_str("# TYPE rws_upstream_retries_total counter\n");
277 for name in names {
278 out.push_str(&format!("rws_upstream_retries_total{{upstream=\"{}\"}} {}\n", label(name), stats.retries[name]));
279 }
280 }
281 if !stats.policies.is_empty() {
282 out.push_str("# HELP rws_upstream_backend_up Whether a backend is in rotation: passing health checks and its circuit breaker not open (1) or not (0)\n");
283 out.push_str("# TYPE rws_upstream_backend_up gauge\n");
284 let mut breaker_lines = String::new();
285 for policy in &stats.policies {
286 let live = policy.live.read().unwrap_or_else(|e| e.into_inner()).clone();
287 let mut backends: Vec<&String> = policy.weights.keys().collect();
288 backends.sort();
289 for backend in backends {
290 let state = policy.breaker_state(backend);
291 let up = live.contains(backend) && state.as_ref().is_none_or(|s| *s != BreakerState::Open);
292 let labels = format!("upstream=\"{}\",backend=\"{}\"", label(&policy.name), label(backend));
293 out.push_str(&format!("rws_upstream_backend_up{{{labels}}} {}\n", u8::from(up)));
294 if let Some(state) = state {
295 let value = match state {
296 BreakerState::Closed => 0,
297 BreakerState::HalfOpen => 1,
298 BreakerState::Open => 2,
299 };
300 breaker_lines.push_str(&format!("rws_upstream_circuit_breaker_state{{{labels}}} {value}\n"));
301 }
302 }
303 }
304 if !breaker_lines.is_empty() {
305 out.push_str("# HELP rws_upstream_circuit_breaker_state Each backend's circuit breaker (0=closed, 1=half_open, 2=open)\n");
306 out.push_str("# TYPE rws_upstream_circuit_breaker_state gauge\n");
307 out.push_str(&breaker_lines);
308 }
309 }
310 out
311}
312
313#[cfg(test)]
314mod tests;