From ea5a30d68652ce09816644c9a092b0754da5acb2 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Sun, 27 Sep 2026 18:43:43 +0300 Subject: [PATCH 1/3] feat(cli): configure the number of runtime worker threads The worker count of the binary's async runtime, and so how many CPU cores the proxy keeps busy, could only be set through TOKIO_WORKER_THREADS. It is now `runtime.worker_threads` in the config file: - The binary loads the config synchronously, then builds a multi-thread runtime with the configured count, or tokio's default (the variable, else the available parallelism) when the key is unset, and runs the proxy on it. - The count must be a positive integer, and `runtime:` rejects unknown keys; either failure stops startup with an error naming the key. A bad TOKIO_WORKER_THREADS is refused by name instead of panicking in tokio. - The startup log states the count and where it came from. - The library lists `runtime` among its known top-level keys, so a file written for the binary loads without a warning; it does not read it. - README and packaging/config.yaml document the key. Closes #97 --- README.md | 13 +++++ cli/Cargo.toml | 5 +- cli/src/main.rs | 22 ++++++--- cli/src/runtime.rs | 103 +++++++++++++++++++++++++++++++++++++++ cli/src/runtime/tests.rs | 84 +++++++++++++++++++++++++++++++ cli/tests/cli.rs | 101 ++++++++++++++++++++++++++++++++++++++ packaging/config.yaml | 6 +++ src/config.rs | 6 ++- src/config/tests.rs | 19 +++++++- 9 files changed, 349 insertions(+), 10 deletions(-) create mode 100644 cli/src/runtime.rs create mode 100644 cli/src/runtime/tests.rs diff --git a/README.md b/README.md index f1c5464..acb051c 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,11 @@ cargo install structured-proxy-cli structured-proxy --config my-service.yaml ``` +The binary runs the proxy on a multi-thread async runtime. `runtime.worker_threads` +in the config file sets how many worker threads, and so CPU cores, it keeps busy; +unset, `TOKIO_WORKER_THREADS` or the available parallelism decides. The startup +log line (`RUST_LOG=info`) states the count and where it came from. + ## Configuration ```yaml @@ -93,6 +98,14 @@ metrics: enabled: true path: "/metrics" +# Optional: the async runtime of the standalone binary (the library runs on its +# embedder's runtime and ignores this section). +runtime: + # Worker threads, i.e. how many CPU cores the proxy keeps busy. Unset: + # TOKIO_WORKER_THREADS, else the available parallelism (the cgroup CPU quota + # on Linux). Must be at least 1. + worker_threads: 2 + # Optional: maintenance mode (returns 503 except for exempt paths) maintenance: enabled: false diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 7c71f38..64bc715 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -29,8 +29,11 @@ structured-proxy = { path = "..", version = "4.3.0", default-features = false } clap = { version = "4", features = ["derive"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } -tokio = { version = "1", features = ["macros", "rt-multi-thread"] } +tokio = { version = "1", features = ["rt-multi-thread"] } anyhow = "1" +# The binary's own `runtime:` section of the config file. +serde = { version = "1", features = ["derive"] } +serde_yaml = "0.9" [features] # The binary the release packages ship: pure-Rust JWT crypto and the shared diff --git a/cli/src/main.rs b/cli/src/main.rs index 680657b..d187532 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -4,6 +4,8 @@ //! structured-proxy --config proxy.yaml //! ``` +mod runtime; + use anyhow::Context as _; use clap::Parser; use tracing_subscriber::EnvFilter; @@ -20,26 +22,34 @@ struct Cli { config: String, } -#[tokio::main] -async fn main() -> anyhow::Result<()> { +fn main() -> anyhow::Result<()> { tracing_subscriber::fmt() .with_env_filter(EnvFilter::from_default_env()) .init(); let cli = Cli::parse(); + let yaml = std::fs::read_to_string(&cli.config) + .with_context(|| format!("loading config {}", cli.config))?; // Reads the ProxyConfig and the transcoding settings kept outside it - // (error_details, streaming.ndjson_envelope) from the same file. - let server = structured_proxy::ProxyServer::from_file(std::path::Path::new(&cli.config)) + // (error_details, streaming.ndjson_envelope) from the same file. Loading + // does no async work, so the runtime is built afterwards, from the same + // file's `runtime:` section. + let server = structured_proxy::ProxyServer::from_yaml_str(&yaml) .with_context(|| format!("loading config {}", cli.config))?; - let config = server.config(); + let file = runtime::FileConfig::from_yaml_str(&yaml) + .with_context(|| format!("loading config {}", cli.config))?; + let (rt, source) = file.runtime.build().context("starting the async runtime")?; + let config = server.config(); tracing::info!( service = %config.service.name, listen = %config.listen.http, upstream = %config.upstream.default, descriptors = config.descriptors.len(), + worker_threads = rt.metrics().num_workers(), + worker_threads_from = %source, "Starting structured-proxy" ); - server.serve().await + rt.block_on(server.serve()) } diff --git a/cli/src/runtime.rs b/cli/src/runtime.rs new file mode 100644 index 0000000..d438479 --- /dev/null +++ b/cli/src/runtime.rs @@ -0,0 +1,103 @@ +//! The `runtime:` section of the config file: the async runtime the binary +//! runs the proxy on. Only the binary reads it; the library runs on whatever +//! runtime its embedder provides. + +use std::fmt; +use std::num::NonZeroUsize; + +use serde::Deserialize; + +/// The variable tokio reads for its default worker count. +const WORKER_THREADS_ENV: &str = "TOKIO_WORKER_THREADS"; + +/// The keys of the config file this binary reads itself; the library reads +/// the rest and ignores these. +#[derive(Debug, Default, Deserialize)] +pub struct FileConfig { + #[serde(default)] + pub runtime: RuntimeConfig, +} + +/// `runtime:`. +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RuntimeConfig { + /// Worker threads of the async runtime, so the CPU cores the proxy keeps + /// busy. Unset: tokio's default. + #[serde(default)] + pub worker_threads: Option, +} + +/// Where the worker count came from. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WorkerSource { + /// `runtime.worker_threads`. + Config, + /// The `TOKIO_WORKER_THREADS` environment variable. + Environment, + /// The available parallelism (on Linux, the cgroup CPU quota). + AvailableParallelism, +} + +impl fmt::Display for WorkerSource { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Config => "runtime.worker_threads", + Self::Environment => WORKER_THREADS_ENV, + Self::AvailableParallelism => "available parallelism", + }) + } +} + +impl FileConfig { + /// The binary's own settings in a config file. + /// + /// # Errors + /// + /// A `runtime:` section that is not a map, holds an unknown key, or sets + /// `worker_threads` to anything but a positive integer; the message names + /// the key. + pub fn from_yaml_str(yaml: &str) -> Result { + serde_yaml::from_str(yaml) + } +} + +impl RuntimeConfig { + /// A multi-thread runtime with the configured worker count, or tokio's + /// default, and where that count came from. + /// + /// # Errors + /// + /// The runtime cannot be created, or `TOKIO_WORKER_THREADS` is used and + /// is not a positive integer: tokio would panic on it, so it is read here + /// and refused with its name. + pub fn build(&self) -> std::io::Result<(tokio::runtime::Runtime, WorkerSource)> { + let mut builder = tokio::runtime::Builder::new_multi_thread(); + builder.enable_all(); + let source = match self.worker_threads { + Some(n) => { + builder.worker_threads(n.get()); + WorkerSource::Config + } + None => match std::env::var_os(WORKER_THREADS_ENV) { + Some(raw) => { + let n = raw + .to_str() + .and_then(|s| s.trim().parse::().ok()) + .ok_or_else(|| { + std::io::Error::other(format!( + "{WORKER_THREADS_ENV} must be a positive integer, got {raw:?}" + )) + })?; + builder.worker_threads(n.get()); + WorkerSource::Environment + } + None => WorkerSource::AvailableParallelism, + }, + }; + Ok((builder.build()?, source)) + } +} + +#[cfg(test)] +mod tests; diff --git a/cli/src/runtime/tests.rs b/cli/src/runtime/tests.rs new file mode 100644 index 0000000..abac6f0 --- /dev/null +++ b/cli/src/runtime/tests.rs @@ -0,0 +1,84 @@ +use super::*; + +/// The `runtime:` settings of `yaml`, or the error they fail with. +fn runtime(yaml: &str) -> Result { + FileConfig::from_yaml_str(yaml) + .map(|file| file.runtime) + .map_err(|e| e.to_string()) +} + +#[test] +fn configured_worker_threads_start_exactly_that_many_workers() { + let config = runtime("runtime:\n worker_threads: 3\n").unwrap(); + let (rt, source) = config.build().unwrap(); + assert_eq!(rt.metrics().num_workers(), 3); + assert_eq!(source, WorkerSource::Config); +} + +#[test] +fn without_the_key_the_environment_decides() { + // Each test runs in its own process under nextest, so the variable does + // not reach any other test. + std::env::set_var(WORKER_THREADS_ENV, "2"); + let config = runtime("upstream:\n default: \"http://x:1\"\n").unwrap(); + let (rt, source) = config.build().unwrap(); + assert_eq!(rt.metrics().num_workers(), 2); + assert_eq!(source, WorkerSource::Environment); +} + +#[test] +fn without_the_key_or_the_variable_tokio_uses_the_available_parallelism() { + std::env::remove_var(WORKER_THREADS_ENV); + let (rt, source) = RuntimeConfig::default().build().unwrap(); + let parallelism = std::thread::available_parallelism().unwrap().get(); + assert_eq!(rt.metrics().num_workers(), parallelism); + assert_eq!(source, WorkerSource::AvailableParallelism); +} + +#[test] +fn the_config_wins_over_the_environment() { + std::env::set_var(WORKER_THREADS_ENV, "5"); + let (rt, source) = runtime("runtime:\n worker_threads: 1\n") + .unwrap() + .build() + .unwrap(); + assert_eq!(rt.metrics().num_workers(), 1); + assert_eq!(source, WorkerSource::Config); +} + +#[test] +fn an_invalid_worker_count_names_the_key() { + for value in ["0", "-1", "two", "1.5", "[]"] { + let err = runtime(&format!("runtime:\n worker_threads: {value}\n")).unwrap_err(); + assert!(err.contains("runtime.worker_threads"), "{value}: {err}"); + } +} + +#[test] +fn an_unknown_runtime_key_names_it() { + let err = runtime("runtime:\n worker_thread: 2\n").unwrap_err(); + assert!(err.contains("worker_thread"), "{err}"); + assert!(err.starts_with("runtime"), "{err}"); +} + +#[test] +fn an_invalid_environment_value_is_refused_by_name() { + for value in ["0", "many", ""] { + std::env::set_var(WORKER_THREADS_ENV, value); + let err = RuntimeConfig::default().build().unwrap_err().to_string(); + assert!(err.contains(WORKER_THREADS_ENV), "{value:?}: {err}"); + } +} + +#[test] +fn the_source_is_named_for_the_startup_log() { + assert_eq!(WorkerSource::Config.to_string(), "runtime.worker_threads"); + assert_eq!( + WorkerSource::Environment.to_string(), + "TOKIO_WORKER_THREADS" + ); + assert_eq!( + WorkerSource::AvailableParallelism.to_string(), + "available parallelism" + ); +} diff --git a/cli/tests/cli.rs b/cli/tests/cli.rs index 17f4c29..9173571 100644 --- a/cli/tests/cli.rs +++ b/cli/tests/cli.rs @@ -114,3 +114,104 @@ fn a_valid_config_starts_the_proxy() { std::fs::remove_file(&path).unwrap(); assert!(response.starts_with("HTTP/1.1 200"), "{response}"); } + +/// A config file for this test process holding `yaml`, removed on drop. +struct ConfigFile(std::path::PathBuf); + +impl ConfigFile { + fn new(name: &str, yaml: &str) -> Self { + let path = std::env::temp_dir().join(format!( + "structured-proxy-cli-test-{name}-{}.yaml", + std::process::id() + )); + std::fs::write(&path, yaml).unwrap(); + Self(path) + } +} + +impl Drop for ConfigFile { + fn drop(&mut self) { + if let Err(e) = std::fs::remove_file(&self.0) { + eprintln!("could not remove {}: {e}", self.0.display()); + } + } +} + +#[test] +fn configured_worker_threads_run_the_proxy_and_are_logged() { + let port = TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() + .port(); + let config = ConfigFile::new( + "workers", + &format!( + "listen:\n http: \"127.0.0.1:{port}\"\n\ + upstream:\n default: \"http://127.0.0.1:9\"\n\ + descriptors: []\n\ + runtime:\n worker_threads: 2\n" + ), + ); + let mut child = Command::new(BIN) + .arg("--config") + .arg(&config.0) + .env("RUST_LOG", "info") + .env("NO_COLOR", "1") + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + let stdout = child.stdout.take().unwrap(); + let _proxy = Running(child); + + // Read the log on its own thread, so a proxy that never logs fails the + // test at the deadline instead of blocking it. + let (lines, logged) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + use std::io::BufRead; + for line in std::io::BufReader::new(stdout).lines() { + let Ok(line) = line else { break }; + if lines.send(line).is_err() { + break; + } + } + }); + let deadline = Instant::now() + Duration::from_secs(20); + let start = loop { + let left = deadline.saturating_duration_since(Instant::now()); + let line = logged + .recv_timeout(left) + .expect("the proxy never logged its start"); + if line.contains("Starting structured-proxy") { + break line; + } + }; + assert!(start.contains("worker_threads=2"), "{start}"); + assert!( + start.contains("worker_threads_from=runtime.worker_threads"), + "{start}" + ); +} + +#[test] +fn a_bad_runtime_section_fails_at_startup_and_names_the_key() { + for (name, runtime, key) in [ + ("zero", "worker_threads: 0", "runtime.worker_threads"), + ("text", "worker_threads: two", "runtime.worker_threads"), + ("unknown", "worker_thread: 2", "worker_thread"), + ] { + let config = ConfigFile::new( + name, + &format!("upstream:\n default: \"http://127.0.0.1:9\"\nruntime:\n {runtime}\n"), + ); + let out = Command::new(BIN) + .arg("--config") + .arg(&config.0) + .output() + .unwrap(); + assert!(!out.status.success(), "{name}"); + let stderr = String::from_utf8(out.stderr).unwrap(); + assert!(stderr.contains(key), "{name}: {stderr}"); + } +} diff --git a/packaging/config.yaml b/packaging/config.yaml index 966c45b..e8ca3c5 100644 --- a/packaging/config.yaml +++ b/packaging/config.yaml @@ -38,6 +38,12 @@ descriptors: [] service: name: "structured-proxy" +# Worker threads of the async runtime, i.e. how many CPU cores the proxy keeps +# busy. Unset: TOKIO_WORKER_THREADS, else the available parallelism (the cgroup +# CPU quota on Linux). Must be at least 1. +# runtime: +# worker_threads: 2 + # CORS: empty list = permissive (dev). List exact origins for production. cors: origins: [] diff --git a/src/config.rs b/src/config.rs index 1b2544d..8da5d7e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -202,8 +202,9 @@ struct StreamingFileConfig { ndjson_envelope: bool, } -/// Every top-level key a config file may use: the [`ProxyConfig`] fields plus -/// the transcoding settings kept outside it. +/// Every top-level key a config file may use: the [`ProxyConfig`] fields, the +/// transcoding settings kept outside it, and `runtime`, which the standalone +/// binary reads to build its async runtime (an embedder brings its own). pub(crate) const KNOWN_TOP_LEVEL_KEYS: &[&str] = &[ "upstream", "descriptors", @@ -224,6 +225,7 @@ pub(crate) const KNOWN_TOP_LEVEL_KEYS: &[&str] = &[ "streaming", "error_details", "response_headers", + "runtime", ]; /// Every `streaming:` key: the [`StreamingConfig`] fields plus the ones diff --git a/src/config/tests.rs b/src/config/tests.rs index eaa33d4..a3d8db0 100644 --- a/src/config/tests.rs +++ b/src/config/tests.rs @@ -478,10 +478,27 @@ fn known_top_level_keys_cover_every_proxy_config_field() { "streaming", "error_details", "response_headers", + "runtime", ] { assert!(KNOWN_TOP_LEVEL_KEYS.contains(&key), "{key}"); } - assert_eq!(KNOWN_TOP_LEVEL_KEYS.len(), 19); + assert_eq!(KNOWN_TOP_LEVEL_KEYS.len(), 20); +} + +#[test] +fn runtime_key_is_known() { + // `runtime:` is the standalone binary's section: the library does not + // read it, but a file written for the binary must load through the + // library without a warning, while a misspelling of it is reported. + let yaml = r#" +upstream: + default: "grpc://x:1" +runtime: + worker_threads: 2 +runtimes: + worker_threads: 2 +"#; + assert_eq!(unknown_config_keys(yaml), vec!["runtimes".to_string()]); } #[test] From 990f92d636e6e2850024a98e64d9bd5828e5ea73 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Sun, 27 Sep 2026 18:57:28 +0300 Subject: [PATCH 2/3] fix(cli): refuse a null worker count, keep tests off the environment - An explicit `runtime.worker_threads: null` or empty value was taken for a left-out key, so a config mistake silently fell back to tokio's default. A present key must now hold a positive integer; only an absent one falls back. - The runtime tests set and removed TOKIO_WORKER_THREADS, a process-wide variable, so tests running in parallel in one process could read each other's value. The count is now computed by build_with from the variable's value, which build reads once; the tests pass values to it and leave the environment alone. With neither the key nor the variable the available parallelism is set explicitly, as tokio would. Regression test: an_invalid_worker_count_names_the_key (null, ~, empty) Part of #97 --- cli/src/runtime.rs | 69 ++++++++++++++++++++++++++-------------- cli/src/runtime/tests.rs | 31 ++++++++++-------- 2 files changed, 62 insertions(+), 38 deletions(-) diff --git a/cli/src/runtime.rs b/cli/src/runtime.rs index d438479..e86c8d9 100644 --- a/cli/src/runtime.rs +++ b/cli/src/runtime.rs @@ -2,6 +2,7 @@ //! runs the proxy on. Only the binary reads it; the library runs on whatever //! runtime its embedder provides. +use std::ffi::OsStr; use std::fmt; use std::num::NonZeroUsize; @@ -23,11 +24,21 @@ pub struct FileConfig { #[serde(deny_unknown_fields)] pub struct RuntimeConfig { /// Worker threads of the async runtime, so the CPU cores the proxy keeps - /// busy. Unset: tokio's default. - #[serde(default)] + /// busy. Left out: tokio's default. Present, it must be a positive + /// integer: an explicit null or empty value is refused, not taken for + /// "left out". + #[serde(default, deserialize_with = "present_count")] pub worker_threads: Option, } +/// A key that is present holds a count; only an absent one (`default`) is +/// `None`. +fn present_count<'de, D: serde::Deserializer<'de>>( + de: D, +) -> Result, D::Error> { + NonZeroUsize::deserialize(de).map(Some) +} + /// Where the worker count came from. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum WorkerSource { @@ -72,30 +83,40 @@ impl RuntimeConfig { /// is not a positive integer: tokio would panic on it, so it is read here /// and refused with its name. pub fn build(&self) -> std::io::Result<(tokio::runtime::Runtime, WorkerSource)> { - let mut builder = tokio::runtime::Builder::new_multi_thread(); - builder.enable_all(); - let source = match self.worker_threads { - Some(n) => { - builder.worker_threads(n.get()); - WorkerSource::Config + self.build_with(std::env::var_os(WORKER_THREADS_ENV).as_deref()) + } + + /// [`build`](Self::build) with the value of `TOKIO_WORKER_THREADS` given + /// rather than read, so the count depends on the arguments alone. + fn build_with( + &self, + env: Option<&OsStr>, + ) -> std::io::Result<(tokio::runtime::Runtime, WorkerSource)> { + let (workers, source) = match (self.worker_threads, env) { + (Some(n), _) => (n.get(), WorkerSource::Config), + (None, Some(raw)) => { + let n = raw + .to_str() + .and_then(|s| s.trim().parse::().ok()) + .ok_or_else(|| { + std::io::Error::other(format!( + "{WORKER_THREADS_ENV} must be a positive integer, got {raw:?}" + )) + })?; + (n.get(), WorkerSource::Environment) } - None => match std::env::var_os(WORKER_THREADS_ENV) { - Some(raw) => { - let n = raw - .to_str() - .and_then(|s| s.trim().parse::().ok()) - .ok_or_else(|| { - std::io::Error::other(format!( - "{WORKER_THREADS_ENV} must be a positive integer, got {raw:?}" - )) - })?; - builder.worker_threads(n.get()); - WorkerSource::Environment - } - None => WorkerSource::AvailableParallelism, - }, + // What tokio falls back to itself; set here so tokio does not + // read the variable a second time. + (None, None) => ( + std::thread::available_parallelism().map_or(1, NonZeroUsize::get), + WorkerSource::AvailableParallelism, + ), }; - Ok((builder.build()?, source)) + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(workers) + .enable_all() + .build()?; + Ok((runtime, source)) } } diff --git a/cli/src/runtime/tests.rs b/cli/src/runtime/tests.rs index abac6f0..4adcccd 100644 --- a/cli/src/runtime/tests.rs +++ b/cli/src/runtime/tests.rs @@ -7,29 +7,29 @@ fn runtime(yaml: &str) -> Result { .map_err(|e| e.to_string()) } +// The tests pass the value of TOKIO_WORKER_THREADS to `build_with` instead of +// setting the process-wide variable, so they can run in parallel in one +// process. + #[test] fn configured_worker_threads_start_exactly_that_many_workers() { let config = runtime("runtime:\n worker_threads: 3\n").unwrap(); - let (rt, source) = config.build().unwrap(); + let (rt, source) = config.build_with(None).unwrap(); assert_eq!(rt.metrics().num_workers(), 3); assert_eq!(source, WorkerSource::Config); } #[test] fn without_the_key_the_environment_decides() { - // Each test runs in its own process under nextest, so the variable does - // not reach any other test. - std::env::set_var(WORKER_THREADS_ENV, "2"); let config = runtime("upstream:\n default: \"http://x:1\"\n").unwrap(); - let (rt, source) = config.build().unwrap(); + let (rt, source) = config.build_with(Some(OsStr::new("2"))).unwrap(); assert_eq!(rt.metrics().num_workers(), 2); assert_eq!(source, WorkerSource::Environment); } #[test] -fn without_the_key_or_the_variable_tokio_uses_the_available_parallelism() { - std::env::remove_var(WORKER_THREADS_ENV); - let (rt, source) = RuntimeConfig::default().build().unwrap(); +fn without_the_key_or_the_variable_the_available_parallelism_decides() { + let (rt, source) = RuntimeConfig::default().build_with(None).unwrap(); let parallelism = std::thread::available_parallelism().unwrap().get(); assert_eq!(rt.metrics().num_workers(), parallelism); assert_eq!(source, WorkerSource::AvailableParallelism); @@ -37,10 +37,9 @@ fn without_the_key_or_the_variable_tokio_uses_the_available_parallelism() { #[test] fn the_config_wins_over_the_environment() { - std::env::set_var(WORKER_THREADS_ENV, "5"); let (rt, source) = runtime("runtime:\n worker_threads: 1\n") .unwrap() - .build() + .build_with(Some(OsStr::new("5"))) .unwrap(); assert_eq!(rt.metrics().num_workers(), 1); assert_eq!(source, WorkerSource::Config); @@ -48,7 +47,9 @@ fn the_config_wins_over_the_environment() { #[test] fn an_invalid_worker_count_names_the_key() { - for value in ["0", "-1", "two", "1.5", "[]"] { + // An explicit null or empty value is a mistake to report, not the same + // as leaving the key out. + for value in ["0", "-1", "two", "1.5", "[]", "null", "~", ""] { let err = runtime(&format!("runtime:\n worker_threads: {value}\n")).unwrap_err(); assert!(err.contains("runtime.worker_threads"), "{value}: {err}"); } @@ -63,9 +64,11 @@ fn an_unknown_runtime_key_names_it() { #[test] fn an_invalid_environment_value_is_refused_by_name() { - for value in ["0", "many", ""] { - std::env::set_var(WORKER_THREADS_ENV, value); - let err = RuntimeConfig::default().build().unwrap_err().to_string(); + for value in ["0", "many", "", "-2"] { + let err = RuntimeConfig::default() + .build_with(Some(OsStr::new(value))) + .unwrap_err() + .to_string(); assert!(err.contains(WORKER_THREADS_ENV), "{value:?}: {err}"); } } From 2cd49f15a29953897e2366392397709dc50bb3d8 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Sun, 27 Sep 2026 19:14:38 +0300 Subject: [PATCH 3/3] fix(cli): refuse worker counts that overflow the thread limit tokio adds the blocking-thread limit to the worker count unchecked, so a count near usize::MAX from runtime.worker_threads or TOKIO_WORKER_THREADS panicked inside tokio's builder. The limit (tokio's default of 512) is now set explicitly and a count whose sum with it overflows is refused with the name of its source. A regression test covers both sources. The available-parallelism test expects the same one-worker fallback as the builder instead of unwrapping the query. --- cli/src/runtime.rs | 17 ++++++++++++++--- cli/src/runtime/tests.rs | 19 ++++++++++++++++++- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/cli/src/runtime.rs b/cli/src/runtime.rs index e86c8d9..b5ec8da 100644 --- a/cli/src/runtime.rs +++ b/cli/src/runtime.rs @@ -11,6 +11,10 @@ use serde::Deserialize; /// The variable tokio reads for its default worker count. const WORKER_THREADS_ENV: &str = "TOKIO_WORKER_THREADS"; +/// Threads the runtime may add for blocking work (tokio's default), set +/// explicitly because tokio adds it to the worker count unchecked. +const MAX_BLOCKING_THREADS: usize = 512; + /// The keys of the config file this binary reads itself; the library reads /// the rest and ignores these. #[derive(Debug, Default, Deserialize)] @@ -79,9 +83,10 @@ impl RuntimeConfig { /// /// # Errors /// - /// The runtime cannot be created, or `TOKIO_WORKER_THREADS` is used and - /// is not a positive integer: tokio would panic on it, so it is read here - /// and refused with its name. + /// The runtime cannot be created, `TOKIO_WORKER_THREADS` is used and is + /// not a positive integer, or the worker count is so large that tokio's + /// thread limit overflows: tokio would panic on either, so they are + /// refused here with the name of the setting. pub fn build(&self) -> std::io::Result<(tokio::runtime::Runtime, WorkerSource)> { self.build_with(std::env::var_os(WORKER_THREADS_ENV).as_deref()) } @@ -112,8 +117,14 @@ impl RuntimeConfig { WorkerSource::AvailableParallelism, ), }; + if workers.checked_add(MAX_BLOCKING_THREADS).is_none() { + return Err(std::io::Error::other(format!( + "{source} is too large: {workers} worker threads" + ))); + } let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(workers) + .max_blocking_threads(MAX_BLOCKING_THREADS) .enable_all() .build()?; Ok((runtime, source)) diff --git a/cli/src/runtime/tests.rs b/cli/src/runtime/tests.rs index 4adcccd..dd92a32 100644 --- a/cli/src/runtime/tests.rs +++ b/cli/src/runtime/tests.rs @@ -30,7 +30,7 @@ fn without_the_key_the_environment_decides() { #[test] fn without_the_key_or_the_variable_the_available_parallelism_decides() { let (rt, source) = RuntimeConfig::default().build_with(None).unwrap(); - let parallelism = std::thread::available_parallelism().unwrap().get(); + let parallelism = std::thread::available_parallelism().map_or(1, NonZeroUsize::get); assert_eq!(rt.metrics().num_workers(), parallelism); assert_eq!(source, WorkerSource::AvailableParallelism); } @@ -73,6 +73,23 @@ fn an_invalid_environment_value_is_refused_by_name() { } } +#[test] +fn a_worker_count_that_overflows_the_thread_limit_is_refused_by_source() { + // tokio adds the blocking-thread limit to the worker count; a sum past + // usize::MAX panicked inside tokio instead of naming the setting. + let too_many = usize::MAX - MAX_BLOCKING_THREADS + 1; + let config = runtime(&format!("runtime:\n worker_threads: {too_many}\n")).unwrap(); + let err = config.build_with(None).unwrap_err().to_string(); + assert!(err.contains("runtime.worker_threads"), "{err}"); + + let env = too_many.to_string(); + let err = RuntimeConfig::default() + .build_with(Some(OsStr::new(&env))) + .unwrap_err() + .to_string(); + assert!(err.contains(WORKER_THREADS_ENV), "{err}"); +} + #[test] fn the_source_is_named_for_the_startup_log() { assert_eq!(WorkerSource::Config.to_string(), "runtime.worker_threads");