rws_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.
44///
45/// A TLS key on a PKCS#11 token keeps the shared runtime (18.25.0): rustls
46/// asks for the signature synchronously inside the handshake, and only on
47/// the multi-threaded runtime can the token's wait be moved off the worker
48/// (`tokio::task::block_in_place`); on a per-core runtime every connection
49/// of that core would wait for the token.
50pub fn default_to_per_core() {
51    if std::env::var_os(RWS_CONFIG_RUNTIME).is_none() {
52        if pkcs11_key_configured() {
53            println!("  runtime: shared (a TLS key is on a PKCS#11 token; set runtime = \"per-core\" to override)");
54            std::env::set_var(RWS_CONFIG_RUNTIME, "shared");
55        } else {
56            std::env::set_var(RWS_CONFIG_RUNTIME, "per-core");
57        }
58        crate::server_config::invalidate();
59    }
60}
61
62/// Whether a `pkcs11:` key is configured, before the config file and flags
63/// are read: in the environment, the arguments or the config file's text.
64/// A false positive (a comment) only keeps the shared runtime.
65fn pkcs11_key_configured() -> bool {
66    let named = |v: &str| v.trim_start_matches(['"', '\'', ' ']).starts_with("pkcs11:");
67    let in_env = std::env::vars().any(|(k, v)| k.starts_with("RWS_CONFIG_") && k.ends_with("KEY_FILE") && named(&v));
68    let in_args = std::env::args().any(|a| a.split_once('=').is_some_and(|(_, v)| named(v)));
69    let file = std::env::var("RWS_CONFIG_FILE").ok().filter(|p| !p.trim().is_empty()).unwrap_or_else(|| "rws.config.toml".to_string());
70    let in_file = std::fs::read_to_string(file).is_ok_and(|text| mentions_pkcs11_key(&text));
71    in_env || in_args || in_file
72}
73
74/// A `*key_file = "pkcs11:…"` line in config text.
75#[doc(hidden)]
76pub fn mentions_pkcs11_key(text: &str) -> bool {
77    text.lines().any(|line| {
78        line.split_once('=').is_some_and(|(k, v)| k.trim().ends_with("key_file") && v.trim().trim_start_matches(['"', '\'']).starts_with("pkcs11:"))
79    })
80}
81
82/// Startup warning for a PKCS#11 key on per-core runtimes, which the
83/// signature wait stalls (`None` when that isn't the case).
84pub fn pkcs11_warning() -> Option<&'static str> {
85    let key_on_token = std::env::vars().any(|(k, v)| k.starts_with("RWS_CONFIG_") && k.ends_with("KEY_FILE") && v.trim().starts_with("pkcs11:"));
86    (enabled() && key_on_token).then_some("[runtime] per-core with a PKCS#11 TLS key: each handshake's signature blocks that core's connections while the token works; runtime = \"shared\" moves the wait off the worker")
87}
88
89/// How many per-core runtimes to start: `TOKIO_WORKER_THREADS` when set
90/// (the same knob as the shared runtime), otherwise one per available CPU.
91pub fn runtime_count() -> usize {
92    std::env::var("TOKIO_WORKER_THREADS")
93        .ok()
94        .and_then(|v| v.trim().parse::<usize>().ok())
95        .filter(|n| *n > 0)
96        .unwrap_or_else(|| std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1))
97}
98
99/// The upstream pool key for `key` on this thread: unchanged on the shared
100/// runtime, suffixed with the runtime's index on a per-core one (its
101/// connections can only be driven by that runtime).
102#[doc(hidden)]
103pub fn pool_key(key: &str) -> Cow<'_, str> {
104    match SLOT.with(|s| s.get()) {
105        0 => Cow::Borrowed(key),
106        // Pushed rather than `format!`ed: this runs for every proxied request.
107        slot => {
108            let slot = crate::ext::decimal::Decimal::new(slot as u128);
109            let mut owned = String::with_capacity(key.len() + 1 + slot.as_str().len());
110            owned.push_str(key);
111            owned.push('#');
112            owned.push_str(slot.as_str());
113            Cow::Owned(owned)
114        }
115    }
116}
117
118/// A listening socket on `addr` that shares the port with the other
119/// runtimes' (`SO_REUSEPORT`: the kernel spreads connections across them).
120#[cfg(unix)]
121#[doc(hidden)]
122pub fn bind(addr: SocketAddr) -> std::io::Result<tokio::net::TcpListener> {
123    // Socket activation: every runtime accepts on a copy of systemd's socket
124    // (a new SO_REUSEPORT socket couldn't share the port with it).
125    if let Some(listener) = crate::systemd::tcp(addr, None) {
126        listener.set_nonblocking(true)?;
127        return tokio::net::TcpListener::from_std(listener);
128    }
129    let socket = if addr.is_ipv4() { tokio::net::TcpSocket::new_v4()? } else { tokio::net::TcpSocket::new_v6()? };
130    socket.set_reuseaddr(true)?;
131    socket.set_reuseport(true)?;
132    socket.bind(addr)?;
133    socket.listen(1024)
134}
135
136/// Start `count` threads, each running `accept_loop(listener)` on its own
137/// single-threaded runtime until `stop` turns `true`. Returns once every
138/// thread has bound its socket (or failed to), so the port accepts
139/// connections again when it returns — the setup listener was closed just
140/// before, and a client let in by it in between was reset. A thread whose
141/// runtime or socket can't be set up reports it and ends; the others keep
142/// serving.
143#[cfg(unix)]
144#[doc(hidden)]
145pub fn spawn<F, Fut>(
146    addr: SocketAddr,
147    count: usize,
148    stop: watch::Receiver<bool>,
149    accept_loop: F,
150) -> Vec<std::thread::JoinHandle<()>>
151where
152    F: Fn(tokio::net::TcpListener, watch::Receiver<bool>) -> Fut + Send + Clone + 'static,
153    Fut: std::future::Future<Output = ()>,
154{
155    if let Some(warning) = pkcs11_warning() {
156        eprintln!("{warning}");
157    }
158    let mut handles = Vec::with_capacity(count);
159    let (bound_tx, bound_rx) = std::sync::mpsc::channel::<()>();
160    for slot in 1..=count {
161        let bound = bound_tx.clone();
162        let stop = stop.clone();
163        let accept_loop = accept_loop.clone();
164        let spawned = std::thread::Builder::new().name(format!("rws-core-{}", slot)).spawn(move || {
165            SLOT.with(|s| s.set(slot));
166            let runtime = match tokio::runtime::Builder::new_current_thread().enable_all().build() {
167                Ok(r) => r,
168                Err(e) => return eprintln!("per-core runtime {}: {}", slot, e),
169            };
170            runtime.block_on(async move {
171                let listener = bind(addr);
172                drop(bound);
173                match listener {
174                    Ok(listener) => accept_loop(listener, stop).await,
175                    Err(e) => eprintln!("per-core runtime {}: unable to listen on {}: {}", slot, addr, e),
176                }
177            });
178        });
179        match spawned {
180            Ok(h) => handles.push(h),
181            Err(e) => eprintln!("per-core runtime {}: {}", slot, e),
182        }
183    }
184    // Every sender dropped: each thread has bound (or given up). Bounded,
185    // in case a thread never gets that far.
186    drop(bound_tx);
187    let _ = bound_rx.recv_timeout(std::time::Duration::from_secs(5));
188    handles
189}
190
191/// Resolves once `stop` is `true` (or its sender is gone).
192#[doc(hidden)]
193pub async fn stopped(stop: &mut watch::Receiver<bool>) {
194    let _ = stop.wait_for(|s| *s).await;
195}
196
197#[cfg(test)]
198mod tests {
199
200    #[test]
201    fn pkcs11_keys_are_spotted_in_config_text() {
202        use super::mentions_pkcs11_key;
203        assert!(mentions_pkcs11_key("tls_key_file = \"pkcs11:token=rws;object=k?module-path=/x.so\""));
204        assert!(mentions_pkcs11_key("[[virtual_host]]\ndomain = \"a\"\nkey_file = 'pkcs11:object=k'"));
205        assert!(!mentions_pkcs11_key("tls_key_file = \"/etc/rws/key.pem\""));
206        assert!(!mentions_pkcs11_key("comment = \"pkcs11:object=k\""));
207        assert!(!mentions_pkcs11_key(""));
208    }
209
210    use super::*;
211
212    #[cfg(unix)]
213    #[test]
214    fn the_port_accepts_connections_when_spawn_returns() {
215        let addr = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap();
216        let (stop_tx, stop_rx) = watch::channel(false);
217        let handles = spawn(addr, 2, stop_rx, |listener: tokio::net::TcpListener, mut stop: watch::Receiver<bool>| async move {
218            tokio::select! {
219                _ = async { while listener.accept().await.is_ok() {} } => {}
220                _ = stopped(&mut stop) => {}
221            }
222        });
223        assert_eq!(2, handles.len());
224        std::net::TcpStream::connect(addr).expect("listening as soon as spawn returns");
225        stop_tx.send(true).unwrap();
226        for handle in handles {
227            handle.join().unwrap();
228        }
229    }
230
231    #[test]
232    fn pool_key_is_shared_off_the_per_core_runtimes() {
233        assert!(matches!(pool_key("backend:80"), Cow::Borrowed("backend:80")));
234    }
235
236    #[test]
237    fn pool_key_is_per_runtime_on_a_per_core_thread() {
238        let keys = std::thread::spawn(|| {
239            SLOT.with(|s| s.set(3));
240            pool_key("backend:80").into_owned()
241        })
242        .join()
243        .unwrap();
244        assert_eq!("backend:80#3", keys);
245    }
246
247    #[test]
248    fn enabled_only_for_per_core() {
249        let _g = crate::test_env::lock();
250        let before = std::env::var(RWS_CONFIG_RUNTIME).ok();
251        std::env::remove_var(RWS_CONFIG_RUNTIME);
252        assert!(!enabled());
253        std::env::set_var(RWS_CONFIG_RUNTIME, "shared");
254        assert!(!enabled());
255        std::env::set_var(RWS_CONFIG_RUNTIME, " Per-Core ");
256        assert_eq!(cfg!(unix), enabled());
257        match before {
258            Some(v) => std::env::set_var(RWS_CONFIG_RUNTIME, v),
259            None => std::env::remove_var(RWS_CONFIG_RUNTIME),
260        }
261    }
262
263    #[test]
264    fn the_binary_defaults_to_per_core_unless_configured() {
265        let _g = crate::test_env::lock();
266        let before = std::env::var(RWS_CONFIG_RUNTIME).ok();
267        std::env::remove_var(RWS_CONFIG_RUNTIME);
268        default_to_per_core();
269        assert_eq!(Some("per-core".to_string()), std::env::var(RWS_CONFIG_RUNTIME).ok());
270        std::env::set_var(RWS_CONFIG_RUNTIME, "shared");
271        default_to_per_core();
272        assert_eq!(Some("shared".to_string()), std::env::var(RWS_CONFIG_RUNTIME).ok());
273        match before {
274            Some(v) => std::env::set_var(RWS_CONFIG_RUNTIME, v),
275            None => std::env::remove_var(RWS_CONFIG_RUNTIME),
276        }
277    }
278
279    #[test]
280    fn runtime_count_follows_tokio_worker_threads() {
281        let _g = crate::test_env::lock();
282        let before = std::env::var("TOKIO_WORKER_THREADS").ok();
283        std::env::set_var("TOKIO_WORKER_THREADS", "3");
284        assert_eq!(3, runtime_count());
285        std::env::set_var("TOKIO_WORKER_THREADS", "zero");
286        assert!(runtime_count() >= 1);
287        match before {
288            Some(v) => std::env::set_var("TOKIO_WORKER_THREADS", v),
289            None => std::env::remove_var("TOKIO_WORKER_THREADS"),
290        }
291    }
292
293    /// Two runtimes share one port; every connection is answered by one of
294    /// them, on its own thread, and both stop when told to.
295    #[cfg(unix)]
296    #[test]
297    fn runtimes_share_the_port_and_stop() {
298        use std::io::{Read, Write};
299        use tokio::io::AsyncWriteExt;
300
301        // A free port: bind, read it, release.
302        let addr = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap();
303        let (stop_tx, stop_rx) = watch::channel(false);
304        let accept_loop = |listener: tokio::net::TcpListener, mut stop: watch::Receiver<bool>| async move {
305            loop {
306                tokio::select! {
307                    result = listener.accept() => {
308                        if let Ok((mut stream, _)) = result {
309                            let slot = SLOT.with(|s| s.get());
310                            let _ = stream.write_all(slot.to_string().as_bytes()).await;
311                        }
312                    }
313                    _ = stopped(&mut stop) => break,
314                }
315            }
316        };
317        let handles = spawn(addr, 2, stop_rx, accept_loop);
318        assert_eq!(2, handles.len());
319
320        let mut answered = 0;
321        for _ in 0..50 {
322            let Ok(mut c) = std::net::TcpStream::connect(addr) else {
323                std::thread::sleep(std::time::Duration::from_millis(20)); // not listening yet
324                continue;
325            };
326            let _ = c.write_all(b"x");
327            let mut slot = String::new();
328            let _ = c.read_to_string(&mut slot);
329            if !slot.is_empty() {
330                assert!(slot == "1" || slot == "2", "slot {slot}");
331                answered += 1;
332            }
333        }
334        assert!(answered > 0);
335
336        stop_tx.send(true).unwrap();
337        for h in handles {
338            h.join().unwrap();
339        }
340        assert!(std::net::TcpStream::connect(addr).is_err());
341    }
342
343    #[cfg(unix)]
344    #[test]
345    fn bind_reports_an_address_in_use_without_reuseport() {
346        let taken = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
347        let addr = taken.local_addr().unwrap();
348        let runtime = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();
349        runtime.block_on(async { assert!(bind(addr).is_err()) });
350    }
351}
rws_server/server/per_core.rs

rws_server/server/
per_core.rs
