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`. A thread whose
92/// runtime or socket can't be set up reports it and ends; the others keep
93/// serving.
94#[cfg(unix)]
95pub(crate) fn spawn<F, Fut>(
96    addr: SocketAddr,
97    count: usize,
98    stop: watch::Receiver<bool>,
99    accept_loop: F,
100) -> Vec<std::thread::JoinHandle<()>>
101where
102    F: Fn(tokio::net::TcpListener, watch::Receiver<bool>) -> Fut + Send + Clone + 'static,
103    Fut: std::future::Future<Output = ()>,
104{
105    let mut handles = Vec::with_capacity(count);
106    for slot in 1..=count {
107        let stop = stop.clone();
108        let accept_loop = accept_loop.clone();
109        let spawned = std::thread::Builder::new().name(format!("rws-core-{}", slot)).spawn(move || {
110            SLOT.with(|s| s.set(slot));
111            let runtime = match tokio::runtime::Builder::new_current_thread().enable_all().build() {
112                Ok(r) => r,
113                Err(e) => return eprintln!("per-core runtime {}: {}", slot, e),
114            };
115            runtime.block_on(async move {
116                match bind(addr) {
117                    Ok(listener) => accept_loop(listener, stop).await,
118                    Err(e) => eprintln!("per-core runtime {}: unable to listen on {}: {}", slot, addr, e),
119                }
120            });
121        });
122        match spawned {
123            Ok(h) => handles.push(h),
124            Err(e) => eprintln!("per-core runtime {}: {}", slot, e),
125        }
126    }
127    handles
128}
129
130/// Resolves once `stop` is `true` (or its sender is gone).
131pub(crate) async fn stopped(stop: &mut watch::Receiver<bool>) {
132    let _ = stop.wait_for(|s| *s).await;
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn pool_key_is_shared_off_the_per_core_runtimes() {
141        assert!(matches!(pool_key("backend:80"), Cow::Borrowed("backend:80")));
142    }
143
144    #[test]
145    fn pool_key_is_per_runtime_on_a_per_core_thread() {
146        let keys = std::thread::spawn(|| {
147            SLOT.with(|s| s.set(3));
148            pool_key("backend:80").into_owned()
149        })
150        .join()
151        .unwrap();
152        assert_eq!("backend:80#3", keys);
153    }
154
155    #[test]
156    fn enabled_only_for_per_core() {
157        let _g = crate::test_env::lock();
158        let before = std::env::var(RWS_CONFIG_RUNTIME).ok();
159        std::env::remove_var(RWS_CONFIG_RUNTIME);
160        assert!(!enabled());
161        std::env::set_var(RWS_CONFIG_RUNTIME, "shared");
162        assert!(!enabled());
163        std::env::set_var(RWS_CONFIG_RUNTIME, " Per-Core ");
164        assert_eq!(cfg!(unix), enabled());
165        match before {
166            Some(v) => std::env::set_var(RWS_CONFIG_RUNTIME, v),
167            None => std::env::remove_var(RWS_CONFIG_RUNTIME),
168        }
169    }
170
171    #[test]
172    fn the_binary_defaults_to_per_core_unless_configured() {
173        let _g = crate::test_env::lock();
174        let before = std::env::var(RWS_CONFIG_RUNTIME).ok();
175        std::env::remove_var(RWS_CONFIG_RUNTIME);
176        default_to_per_core();
177        assert_eq!(Some("per-core".to_string()), std::env::var(RWS_CONFIG_RUNTIME).ok());
178        std::env::set_var(RWS_CONFIG_RUNTIME, "shared");
179        default_to_per_core();
180        assert_eq!(Some("shared".to_string()), std::env::var(RWS_CONFIG_RUNTIME).ok());
181        match before {
182            Some(v) => std::env::set_var(RWS_CONFIG_RUNTIME, v),
183            None => std::env::remove_var(RWS_CONFIG_RUNTIME),
184        }
185    }
186
187    #[test]
188    fn runtime_count_follows_tokio_worker_threads() {
189        let _g = crate::test_env::lock();
190        let before = std::env::var("TOKIO_WORKER_THREADS").ok();
191        std::env::set_var("TOKIO_WORKER_THREADS", "3");
192        assert_eq!(3, runtime_count());
193        std::env::set_var("TOKIO_WORKER_THREADS", "zero");
194        assert!(runtime_count() >= 1);
195        match before {
196            Some(v) => std::env::set_var("TOKIO_WORKER_THREADS", v),
197            None => std::env::remove_var("TOKIO_WORKER_THREADS"),
198        }
199    }
200
201    /// Two runtimes share one port; every connection is answered by one of
202    /// them, on its own thread, and both stop when told to.
203    #[cfg(unix)]
204    #[test]
205    fn runtimes_share_the_port_and_stop() {
206        use std::io::{Read, Write};
207        use tokio::io::AsyncWriteExt;
208
209        // A free port: bind, read it, release.
210        let addr = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap();
211        let (stop_tx, stop_rx) = watch::channel(false);
212        let accept_loop = |listener: tokio::net::TcpListener, mut stop: watch::Receiver<bool>| async move {
213            loop {
214                tokio::select! {
215                    result = listener.accept() => {
216                        if let Ok((mut stream, _)) = result {
217                            let slot = SLOT.with(|s| s.get());
218                            let _ = stream.write_all(slot.to_string().as_bytes()).await;
219                        }
220                    }
221                    _ = stopped(&mut stop) => break,
222                }
223            }
224        };
225        let handles = spawn(addr, 2, stop_rx, accept_loop);
226        assert_eq!(2, handles.len());
227
228        let mut answered = 0;
229        for _ in 0..50 {
230            let Ok(mut c) = std::net::TcpStream::connect(addr) else {
231                std::thread::sleep(std::time::Duration::from_millis(20)); // not listening yet
232                continue;
233            };
234            let _ = c.write_all(b"x");
235            let mut slot = String::new();
236            let _ = c.read_to_string(&mut slot);
237            if !slot.is_empty() {
238                assert!(slot == "1" || slot == "2", "slot {slot}");
239                answered += 1;
240            }
241        }
242        assert!(answered > 0);
243
244        stop_tx.send(true).unwrap();
245        for h in handles {
246            h.join().unwrap();
247        }
248        assert!(std::net::TcpStream::connect(addr).is_err());
249    }
250
251    #[cfg(unix)]
252    #[test]
253    fn bind_reports_an_address_in_use_without_reuseport() {
254        let taken = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
255        let addr = taken.local_addr().unwrap();
256        let runtime = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();
257        runtime.block_on(async { assert!(bind(addr).is_err()) });
258    }
259}
rust_web_server/server/per_core.rs

rust_web_server/server/
per_core.rs
