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

rust_web_server/server/
per_core.rs
