Skip to main content

rust_web_server/proxy_config/
upstream.rs

1//! What an `[[upstream]]` adds to choosing a backend in config mode
2//! (since 18.6.0): `weights`, cookie stickiness (`[upstream.sticky]`), a
3//! per-backend circuit breaker — passive health checking —
4//! (`[upstream.circuit_breaker]`), retries on another backend
5//! (`[upstream.retry]`), and the per-upstream metrics on `/metrics`.
6//!
7//! One [`UpstreamPolicy`] per declared upstream, shared by every route
8//! that proxies to it, rebuilt on reload.
9
10use std::collections::HashMap;
11use std::sync::{Arc, Mutex, OnceLock, RwLock};
12use std::time::Instant;
13
14use crate::circuit_breaker::{BreakerState, CircuitBreaker};
15
16/// `[upstream.sticky]`: keep a client on one backend with a cookie.
17#[derive(Debug, Clone, PartialEq)]
18pub struct StickyConfig {
19    /// The cookie's name.
20    pub cookie: String,
21    /// `Max-Age` in seconds; `None` = a session cookie.
22    pub max_age_secs: Option<u64>,
23}
24
25/// `[upstream.circuit_breaker]`: take a failing backend out of rotation.
26#[derive(Debug, Clone, PartialEq)]
27pub struct BreakerConfig {
28    /// Consecutive failures (connection errors, 502/503/504) that open it.
29    pub failure_threshold: u32,
30    /// How long an open backend is left alone before a probe request.
31    pub recovery_secs: u64,
32    /// Probe requests allowed at once while half-open.
33    pub half_open_probes: u32,
34}
35
36/// `[upstream.retry]`: try another backend.
37#[derive(Debug, Clone, PartialEq)]
38pub struct RetryConfig {
39    /// Further backends tried after the first (so up to `attempts + 1`).
40    pub attempts: u32,
41    /// Statuses retried besides connection errors (default 502, 503, 504).
42    pub status: Vec<i16>,
43}
44
45/// One upstream's selection state.
46pub(crate) struct UpstreamPolicy {
47    pub(crate) name: String,
48    /// The health-checked backend list (for metrics).
49    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    /// Smooth weighted round-robin's current weights.
55    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            // Every backend, weight 1 unless given (also the gauges' list).
72            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    /// A backend's weight (1 unless configured).
81    pub(crate) fn weight(&self, backend: &str) -> u32 {
82        self.weights.get(backend).copied().unwrap_or(1).max(1)
83    }
84
85    /// Whether any backend has a weight other than 1.
86    pub(crate) fn weighted(&self) -> bool {
87        self.weights.values().any(|w| *w != 1)
88    }
89
90    /// Smooth weighted round-robin (nginx's): each pick raises every
91    /// candidate by its weight and lowers the chosen one by the total, so
92    /// weights 3:1 give a, a, b, a, ... — evenly interleaved.
93    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    /// Weighted pick for `random`: `roll` is any number.
110    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    /// Whether the breaker lets a request through to `backend` now (a
124    /// half-open probe counts as one).
125    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    /// Whether the breaker holds any of `backends` open.
133    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    /// Whether `status` from a backend should be retried elsewhere.
156    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
161/// A backend's id in the sticky cookie: 12 hex digits of its URL's SHA-256
162/// — stable across restarts and replicas, without revealing the address.
163pub(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
167/// The value of cookie `name` in the request's `Cookie` headers.
168pub(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
177/// Whether a request may be sent again to another backend: an idempotent
178/// method and a body still in memory.
179pub(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
183// ── metrics ───────────────────────────────────────────────────────────────────
184
185const 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    /// By status code, `"error"` for a request that got no response.
190    requests: HashMap<String, u64>,
191    buckets: [u64; 11],
192    sum: f64,
193    count: u64,
194}
195
196#[derive(Default)]
197struct Stats {
198    /// `(upstream, backend)` → stats.
199    backends: HashMap<(String, String), BackendStats>,
200    retries: HashMap<String, u64>,
201    /// The current configuration's upstreams, for gauges.
202    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
210/// Make `policies` the ones `/metrics` reports gauges for (each compile).
211pub(crate) fn register(policies: Vec<Arc<UpstreamPolicy>>) {
212    stats().lock().unwrap_or_else(|e| e.into_inner()).policies = policies;
213}
214
215/// One request to `backend` of `upstream`: its status (`None`: no
216/// response) and how long it took.
217pub(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
239/// The `rws_upstream_*` series for `/metrics`; empty without upstreams.
240pub(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;