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

rws_proxy/proxy/
mirror.rs
