rust_web_server/server/
per_core.rs

1//! Per-core runtimes (`RWS_CONFIG_RUNTIME=per-core`, Unix only): one
2//! single-threaded tokio runtime per core, each accepting on its own
3//! `SO_REUSEPORT` socket, the way nginx runs one worker per core. A
4//! connection stays on the thread the kernel gave it to, so there's no work
5//! stealing, no cross-thread wake-ups and no idle workers searching for work.
6//!
7//! The `rws` binary runs per-core unless configured otherwise (since 18.0.0,
8//! [`default_to_per_core`]); for an application built on the library the
9//! default stays `shared`, tokio's multi-threaded runtime. `run_tls` and
10//! `run_plain` start the per-core accept loops ([`spawn`]) and keep the
11//! signals: SIGINT/SIGTERM stop every loop, SIGHUP reloads the config and
12//! hands the rebuilt TLS acceptor to every loop. HTTP/3, the redirect
13//! listener and ACME stay on the main runtime.
14//!
15//! Tokio I/O belongs to the runtime that created it, so pooled upstream
16//! connections are kept per thread ([`pool_key`]).
17
18use std::borrow::Cow;
19use std::cell::Cell;
20use std::net::SocketAddr;
21
22use tokio::sync::watch;
23
24/// `RWS_CONFIG_RUNTIME`: `shared` or `per-core` (the `rws` binary's default).
25pub const RWS_CONFIG_RUNTIME: &str = "RWS_CONFIG_RUNTIME";
26
27thread_local! {
28    /// 1-based index of this thread's per-core runtime; 0 elsewhere.
29    static SLOT: Cell<usize> = const { Cell::new(0) };
30}
31
32/// Whether `RWS_CONFIG_RUNTIME` asks for per-core runtimes (always `false`
33/// off Unix, where there is no `SO_REUSEPORT`).
34pub fn enabled() -> bool {
35    cfg!(unix)
36        && std::env::var(RWS_CONFIG_RUNTIME).is_ok_and(|v| v.trim().eq_ignore_ascii_case("per-core"))
37}
38
39/// Make `per-core` the runtime unless `RWS_CONFIG_RUNTIME` is already set —
40/// what the `rws` binary does before reading its config file and flags, so
41/// `runtime = "shared"` there (or in the environment) still wins. Library
42/// applications don't call it: their handlers may keep tokio I/O in
43/// process-wide structures, which per-core runtimes don't share.
44pub fn default_to_per_core() {
45    if std::env::var_os(RWS_CONFIG_RUNTIME).is_none() {
46        std::env::set_var(RWS_CONFIG_RUNTIME, "per-core");
47        crate::server_config::invalidate();
48    }
49}
50
51/// How many per-core runtimes to start: `TOKIO_WORKER_THREADS` when set
52/// (the same knob as the shared runtime), otherwise one per available CPU.
53pub fn runtime_count() -> usize {
54    std::env::var("TOKIO_WORKER_THREADS")
55        .ok()
56        .and_then(|v| v.trim().parse::<usize>().ok())
57        .filter(|n| *n > 0)
58        .unwrap_or_else(|| std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1))
59}
60
61/// The upstream pool key for `key` on this thread: unchanged on the shared
62/// runtime, suffixed with the runtime's index on a per-core one (its
63/// connections can only be driven by that runtime).
64pub(crate) fn pool_key(key: &str) -> Cow<'_, str> {
65    match SLOT.with(|s| s.get()) {
66        0 => Cow::Borrowed(key),
67        // Pushed rather than `format!`ed: this runs for every proxied request.
68        slot => {
69            let slot = crate::ext::decimal::Decimal::new(slot as u128);
70            let mut owned = String::with_capacity(key.len() + 1 + slot.as_str().len());
71            owned.push_str(key);
72            owned.push('#');
73            owned.push_str(slot.as_str());
74            Cow::Owned(owned)
75        }
76    }
77}
78
79/// A listening socket on `addr` that shares the port with the other
80/// runtimes' (`SO_REUSEPORT`: the kernel spreads connections across them).
81#[cfg(unix)]
82pub(crate) fn bind(addr: SocketAddr) -> std::io::Result<tokio::net::TcpListener> {
83    let socket = if addr.is_ipv4() { tokio::net::TcpSocket::new_v4()? } else { tokio::net::TcpSocket::new_v6()? };
84    socket.set_reuseaddr(true)?;
85    socket.set_reuseport(true)?;
86    socket.bind(addr)?;
87    socket.listen(1024)
88}
89
90/// Start `count` threads, each running `accept_loop(listener)` on its own
91/// single-threaded runtime until `stop` turns `true`. Returns once every
92/// thread has bound its socket (or failed to), so the port accepts
93/// connections again when it returns — the setup listener was closed just
94/// before, and a client let in by it in between was reset. A thread whose
95/// runtime or socket can't be set up reports it and ends; the others keep
96/// serving.
97#[cfg(unix)]
98pub(crate) fn spawn<F, Fut>(
99    addr: SocketAddr,
100    count: usize,
101    stop: watch::Receiver<bool>,
102    accept_loop: F,
103) -> Vec<std::thread::JoinHandle<()>>
104where
105    F: Fn(tokio::net::TcpListener, watch::Receiver<bool>) -> Fut + Send + Clone + 'static,
106    Fut: std::future::Future<Output = ()>,
107{
108    let mut handles = Vec::with_capacity(count);
109    let (bound_tx, bound_rx) = std::sync::mpsc::channel::<()>();
110    for slot in 1..=count {
111        let bound = bound_tx.clone();
112        let stop = stop.clone();
113        let accept_loop = accept_loop.clone();
114        let spawned = std::thread::Builder::new().name(format!("rws-core-{}", slot)).spawn(move || {
115            SLOT.with(|s| s.set(slot));
116            let runtime = match tokio::runtime::Builder::new_current_thread().enable_all().build() {
117                Ok(r) => r,
118                Err(e) => return eprintln!("per-core runtime {}: {}", slot, e),
119            };
120            runtime.block_on(async move {
121                let listener = bind(addr);
122                drop(bound);
123                match listener {
124                    Ok(listener) => accept_loop(listener, stop).await,
125                    Err(e) => eprintln!("per-core runtime {}: unable to listen on {}: {}", slot, addr, e),
126                }
127            });
128        });
129        match spawned {
130            Ok(h) => handles.push(h),
131            Err(e) => eprintln!("per-core runtime {}: {}", slot, e),
132        }
133    }
134    // Every sender dropped: each thread has bound (or given up). Bounded,
135    // in case a thread never gets that far.
136    drop(bound_tx);
137    let _ = bound_rx.recv_timeout(std::time::Duration::from_secs(5));
138    handles
139}
140
141/// Resolves once `stop` is `true` (or its sender is gone).
142pub(crate) async fn stopped(stop: &mut watch::Receiver<bool>) {
143    let _ = stop.wait_for(|s| *s).await;
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    #[cfg(unix)]
151    #[test]
152    fn the_port_accepts_connections_when_spawn_returns() {
153        let addr = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap();
154        let (stop_tx, stop_rx) = watch::channel(false);
155        let handles = spawn(addr, 2, stop_rx, |listener: tokio::net::TcpListener, mut stop: watch::Receiver<bool>| async move {
156            tokio::select! {
157                _ = async { while listener.accept().await.is_ok() {} } => {}
158                _ = stopped(&mut stop) => {}
159            }
160        });
161        assert_eq!(2, handles.len());
162        std::net::TcpStream::connect(addr).expect("listening as soon as spawn returns");
163        stop_tx.send(true).unwrap();
164        for handle in handles {
165            handle.join().unwrap();
166        }
167    }
168
169    #[test]
170    fn pool_key_is_shared_off_the_per_core_runtimes() {
171        assert!(matches!(pool_key("backend:80"), Cow::Borrowed("backend:80")));
172    }
173
174    #[test]
175    fn pool_key_is_per_runtime_on_a_per_core_thread() {
176        let keys = std::thread::spawn(|| {
177            SLOT.with(|s| s.set(3));
178            pool_key("backend:80").into_owned()
179        })
180        .join()
181        .unwrap();
182        assert_eq!("backend:80#3", keys);
183    }
184
185    #[test]
186    fn enabled_only_for_per_core() {
187        let _g = crate::test_env::lock();
188        let before = std::env::var(RWS_CONFIG_RUNTIME).ok();
189        std::env::remove_var(RWS_CONFIG_RUNTIME);
190        assert!(!enabled());
191        std::env::set_var(RWS_CONFIG_RUNTIME, "shared");
192        assert!(!enabled());
193        std::env::set_var(RWS_CONFIG_RUNTIME, " Per-Core ");
194        assert_eq!(cfg!(unix), enabled());
195        match before {
196            Some(v) => std::env::set_var(RWS_CONFIG_RUNTIME, v),
197            None => std::env::remove_var(RWS_CONFIG_RUNTIME),
198        }
199    }
200
201    #[test]
202    fn the_binary_defaults_to_per_core_unless_configured() {
203        let _g = crate::test_env::lock();
204        let before = std::env::var(RWS_CONFIG_RUNTIME).ok();
205        std::env::remove_var(RWS_CONFIG_RUNTIME);
206        default_to_per_core();
207        assert_eq!(Some("per-core".to_string()), std::env::var(RWS_CONFIG_RUNTIME).ok());
208        std::env::set_var(RWS_CONFIG_RUNTIME, "shared");
209        default_to_per_core();
210        assert_eq!(Some("shared".to_string()), std::env::var(RWS_CONFIG_RUNTIME).ok());
211        match before {
212            Some(v) => std::env::set_var(RWS_CONFIG_RUNTIME, v),
213            None => std::env::remove_var(RWS_CONFIG_RUNTIME),
214        }
215    }
216
217    #[test]
218    fn runtime_count_follows_tokio_worker_threads() {
219        let _g = crate::test_env::lock();
220        let before = std::env::var("TOKIO_WORKER_THREADS").ok();
221        std::env::set_var("TOKIO_WORKER_THREADS", "3");
222        assert_eq!(3, runtime_count());
223        std::env::set_var("TOKIO_WORKER_THREADS", "zero");
224        assert!(runtime_count() >= 1);
225        match before {
226            Some(v) => std::env::set_var("TOKIO_WORKER_THREADS", v),
227            None => std::env::remove_var("TOKIO_WORKER_THREADS"),
228        }
229    }
230
231    /// Two runtimes share one port; every connection is answered by one of
232    /// them, on its own thread, and both stop when told to.
233    #[cfg(unix)]
234    #[test]
235    fn runtimes_share_the_port_and_stop() {
236        use std::io::{Read, Write};
237        use tokio::io::AsyncWriteExt;
238
239        // A free port: bind, read it, release.
240        let addr = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap();
241        let (stop_tx, stop_rx) = watch::channel(false);
242        let accept_loop = |listener: tokio::net::TcpListener, mut stop: watch::Receiver<bool>| async move {
243            loop {
244                tokio::select! {
245                    result = listener.accept() => {
246                        if let Ok((mut stream, _)) = result {
247                            let slot = SLOT.with(|s| s.get());
248                            let _ = stream.write_all(slot.to_string().as_bytes()).await;
249                        }
250                    }
251                    _ = stopped(&mut stop) => break,
252                }
253            }
254        };
255        let handles = spawn(addr, 2, stop_rx, accept_loop);
256        assert_eq!(2, handles.len());
257
258        let mut answered = 0;
259        for _ in 0..50 {
260            let Ok(mut c) = std::net::TcpStream::connect(addr) else {
261                std::thread::sleep(std::time::Duration::from_millis(20)); // not listening yet
262                continue;
263            };
264            let _ = c.write_all(b"x");
265            let mut slot = String::new();
266            let _ = c.read_to_string(&mut slot);
267            if !slot.is_empty() {
268                assert!(slot == "1" || slot == "2", "slot {slot}");
269                answered += 1;
270            }
271        }
272        assert!(answered > 0);
273
274        stop_tx.send(true).unwrap();
275        for h in handles {
276            h.join().unwrap();
277        }
278        assert!(std::net::TcpStream::connect(addr).is_err());
279    }
280
281    #[cfg(unix)]
282    #[test]
283    fn bind_reports_an_address_in_use_without_reuseport() {
284        let taken = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
285        let addr = taken.local_addr().unwrap();
286        let runtime = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();
287        runtime.block_on(async { assert!(bind(addr).is_err()) });
288    }
289}
rust_web_server/server/per_core.rs

rust_web_server/server/
per_core.rs
