rust_web_server/proxy/
mirror.rs

1//! Request mirroring (18.11.0, #39): send a copy of a share of requests to
2//! other backends — a new version under real traffic, a shadow service —
3//! and throw the answers away.
4//!
5//! Fire-and-forget: the client's request never waits for, or depends on, a
6//! mirror. Copies go through a bounded queue to a few worker threads (not a
7//! thread per copy), so a slow or dead mirror can't use up threads or
8//! memory; when the queue is full the copy is dropped and counted. Only
9//! requests whose body is in memory and at most [`MAX_MIRRORED_BODY`] are
10//! mirrored: a streamed upload is read once, by the real backend.
11//!
12//! ```rust,no_run
13//! use rust_web_server::app::App;
14//! use rust_web_server::core::New;
15//! use rust_web_server::proxy::mirror::{Mirror, MirrorLayer};
16//!
17//! let shadow = Mirror::new("v2-shadow", vec!["http://10.0.0.9:8080".into()]).percent(10.0);
18//! let app = App::new().wrap(MirrorLayer::new(shadow));
19//! ```
20//!
21//! Counts: `rws_mirror_requests_total{mirror,outcome}` with outcome `sent`
22//! (the mirror answered), `failed` (no answer), `dropped` (queue full) and
23//! `skipped` (body streamed or too large).
24
25use std::collections::BTreeMap;
26use std::sync::atomic::{AtomicU64, Ordering};
27use std::sync::mpsc::{sync_channel, SyncSender, TrySendError};
28use std::sync::{Arc, Mutex, OnceLock, RwLock};
29use std::time::Duration;
30
31use crate::application::Application;
32use crate::middleware::Middleware;
33use crate::request::Request;
34use crate::response::Response;
35use crate::server::ConnectionInfo;
36
37/// Copies waiting for a worker; past this, copies are dropped.
38pub const QUEUE: usize = 1024;
39/// Worker threads sending copies (shared by every mirror).
40pub const WORKERS: usize = 4;
41/// Larger request bodies aren't mirrored.
42pub const MAX_MIRRORED_BODY: usize = 64 * 1024;
43
44const CONNECT_TIMEOUT: Duration = Duration::from_secs(2);
45const READ_TIMEOUT: Duration = Duration::from_secs(10);
46
47/// One mirror target: backends (round-robin) and the share of requests
48/// copied to it. Cheap to clone; clones share their state.
49#[derive(Clone)]
50pub struct Mirror {
51    inner: Arc<Inner>,
52}
53
54struct Inner {
55    name: String,
56    /// `http://` / `https://` URLs or `host:port`; a live list (health
57    /// checks may change it).
58    backends: Arc<RwLock<Vec<String>>>,
59    /// Share of requests copied, 0–100.
60    percent: f64,
61    seen: AtomicU64,
62    next: AtomicU64,
63}
64
65impl std::fmt::Debug for Mirror {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        f.debug_struct("Mirror").field("name", &self.inner.name).field("percent", &self.inner.percent).finish_non_exhaustive()
68    }
69}
70
71impl Mirror {
72    /// Copy every request to `backends` (round-robin).
73    pub fn new(name: impl Into<String>, backends: Vec<String>) -> Mirror {
74        Mirror::with_live_backends(name, Arc::new(RwLock::new(backends)))
75    }
76
77    /// Like [`Mirror::new`], with a backend list another component updates
78    /// (a config-file upstream's health-checked list).
79    pub fn with_live_backends(name: impl Into<String>, backends: Arc<RwLock<Vec<String>>>) -> Mirror {
80        Mirror { inner: Arc::new(Inner { name: name.into(), backends, percent: 100.0, seen: AtomicU64::new(0), next: AtomicU64::new(0) }) }
81    }
82
83    /// Copy only this share of requests (0–100), spread evenly: 25 copies
84    /// every 4th request, not 25 in a row.
85    pub fn percent(self, percent: f64) -> Mirror {
86        let percent = if percent.is_finite() { percent.clamp(0.0, 100.0) } else { 0.0 };
87        let inner = &self.inner;
88        Mirror {
89            inner: Arc::new(Inner {
90                name: inner.name.clone(),
91                backends: Arc::clone(&inner.backends),
92                percent,
93                seen: AtomicU64::new(0),
94                next: AtomicU64::new(0),
95            }),
96        }
97    }
98
99    pub fn name(&self) -> &str {
100        &self.inner.name
101    }
102
103    /// Whether the `n`th request (from 0) is in the sampled share.
104    fn sampled(&self) -> bool {
105        let p = self.inner.percent;
106        if p >= 100.0 {
107            return true;
108        }
109        if p <= 0.0 {
110            return false;
111        }
112        let n = self.inner.seen.fetch_add(1, Ordering::Relaxed) as f64;
113        ((n + 1.0) * p / 100.0).floor() > (n * p / 100.0).floor()
114    }
115
116    /// Queue a copy of `request` if it's in the sampled share; never blocks.
117    /// `true` when a copy was queued.
118    pub fn send(&self, request: &Request, connection: &ConnectionInfo) -> bool {
119        if !self.sampled() {
120            return false;
121        }
122        if request.body.is_stream() || request.body.len() > MAX_MIRRORED_BODY {
123            count(&self.inner.name, "skipped");
124            return false;
125        }
126        let backend = {
127            let backends = self.inner.backends.read().unwrap_or_else(|e| e.into_inner());
128            if backends.is_empty() {
129                count(&self.inner.name, "failed");
130                return false;
131            }
132            let i = self.inner.next.fetch_add(1, Ordering::Relaxed) as usize % backends.len();
133            backends[i].clone()
134        };
135        let job = Job { mirror: self.inner.name.clone(), backend, request: request.clone(), client_ip: connection.client.ip.clone(), tls: connection.tls };
136        match queue().try_send(job) {
137            Ok(()) => true,
138            Err(TrySendError::Full(job)) | Err(TrySendError::Disconnected(job)) => {
139                count(&job.mirror, "dropped");
140                false
141            }
142        }
143    }
144}
145
146struct Job {
147    mirror: String,
148    backend: String,
149    request: Request,
150    client_ip: String,
151    tls: bool,
152}
153
154fn queue() -> &'static SyncSender<Job> {
155    static QUEUE_TX: OnceLock<SyncSender<Job>> = OnceLock::new();
156    QUEUE_TX.get_or_init(|| {
157        let (tx, rx) = sync_channel::<Job>(QUEUE);
158        let rx = Arc::new(Mutex::new(rx));
159        for i in 0..WORKERS {
160            let rx = Arc::clone(&rx);
161            let _ = std::thread::Builder::new().name(format!("rws-mirror-{i}")).spawn(move || loop {
162                let job = match rx.lock().unwrap_or_else(|e| e.into_inner()).recv() {
163                    Ok(job) => job,
164                    Err(_) => return,
165                };
166                let outcome = if deliver(&job) { "sent" } else { "failed" };
167                count(&job.mirror, outcome);
168            });
169        }
170        tx
171    })
172}
173
174/// Send one copy and wait for (and drop) the answer.
175fn deliver(job: &Job) -> bool {
176    let Some((host, port, tls)) = crate::proxy_config::health::parse_backend_url(&job.backend) else { return false };
177    crate::proxy::forward_upstream(
178        crate::proxy::shared_pool(),
179        &job.request,
180        crate::proxy::Peer { ip: &job.client_ip, tls: job.tls },
181        &host,
182        port,
183        tls,
184        CONNECT_TIMEOUT,
185        READ_TIMEOUT,
186    )
187    .map(|response| {
188        // A streamed answer is read to the end so the connection isn't
189        // left half-used; then it's dropped.
190        if let Some(mut pipe) = response.stream_pipe {
191            let _ = std::io::copy(&mut pipe, &mut std::io::sink());
192        }
193    })
194    .is_ok()
195}
196
197fn counts() -> &'static Mutex<BTreeMap<(String, &'static str), u64>> {
198    static COUNTS: OnceLock<Mutex<BTreeMap<(String, &'static str), u64>>> = OnceLock::new();
199    COUNTS.get_or_init(Mutex::default)
200}
201
202fn count(mirror: &str, outcome: &'static str) {
203    *counts().lock().unwrap_or_else(|e| e.into_inner()).entry((mirror.to_string(), outcome)).or_default() += 1;
204}
205
206/// The `rws_mirror_requests_total` series; empty until a mirror is used.
207pub(crate) fn prometheus_text() -> String {
208    let counts = counts().lock().unwrap_or_else(|e| e.into_inner());
209    if counts.is_empty() {
210        return String::new();
211    }
212    let mut out = String::from(
213        "# HELP rws_mirror_requests_total Request copies sent to mirrors, by outcome: sent, failed, dropped (queue full), skipped (body streamed or too large)\n# TYPE rws_mirror_requests_total counter\n",
214    );
215    for ((mirror, outcome), n) in counts.iter() {
216        let mirror = mirror.replace('\\', "\\\\").replace('"', "\\\"").replace('\n', "\\n");
217        out.push_str(&format!("rws_mirror_requests_total{{mirror=\"{mirror}\",outcome=\"{outcome}\"}} {n}\n"));
218    }
219    out
220}
221
222/// A [`Mirror`] as middleware: copies the request, then hands it on.
223#[derive(Clone, Debug)]
224pub struct MirrorLayer {
225    mirror: Mirror,
226}
227
228impl MirrorLayer {
229    pub fn new(mirror: Mirror) -> MirrorLayer {
230        MirrorLayer { mirror }
231    }
232}
233
234impl Middleware for MirrorLayer {
235    fn handle(&self, request: &Request, connection: &ConnectionInfo, next: &dyn Application) -> Result<Response, String> {
236        self.mirror.send(request, connection);
237        next.execute(request, connection)
238    }
239
240    /// Queueing a copy never blocks.
241    fn runs_inline(&self, _request: &Request, _connection: &ConnectionInfo) -> bool {
242        true
243    }
244
245    /// A streamed body passes through untouched (and isn't mirrored).
246    fn accepts_streamed_body(&self, _request: &Request, _connection: &ConnectionInfo) -> bool {
247        true
248    }
249}
250
251#[cfg(test)]
252mod tests;
rust_web_server/proxy/mirror.rs

rust_web_server/proxy/
mirror.rs
