feat(cli): configure the number of runtime worker threads

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

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

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())
 }

cli/src/runtime.rs

@@ -0,0 +1,135 @@
+//! 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::ffi::OsStr;
+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";
+
+/// 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)]
+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. 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<NonZeroUsize>,
+}
+
+/// 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<Option<NonZeroUsize>, D::Error> {
+    NonZeroUsize::deserialize(de).map(Some)
+}
+
+/// 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<Self, serde_yaml::Error> {
+        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, `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())
+    }
+
+    /// [`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::<NonZeroUsize>().ok())
+                    .ok_or_else(|| {
+                        std::io::Error::other(format!(
+                            "{WORKER_THREADS_ENV} must be a positive integer, got {raw:?}"
+                        ))
+                    })?;
+                (n.get(), WorkerSource::Environment)
+            }
+            // 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,
+            ),
+        };
+        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))
+    }
+}
+
+#[cfg(test)]
+mod tests;

cli/src/runtime/tests.rs

@@ -0,0 +1,104 @@
+use super::*;
+
+/// The `runtime:` settings of `yaml`, or the error they fail with.
+fn runtime(yaml: &str) -> Result<RuntimeConfig, String> {
+    FileConfig::from_yaml_str(yaml)
+        .map(|file| file.runtime)
+        .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_with(None).unwrap();
+    assert_eq!(rt.metrics().num_workers(), 3);
+    assert_eq!(source, WorkerSource::Config);
+}
+
+#[test]
+fn without_the_key_the_environment_decides() {
+    let config = runtime("upstream:\n  default: \"http://x:1\"\n").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_the_available_parallelism_decides() {
+    let (rt, source) = RuntimeConfig::default().build_with(None).unwrap();
+    let parallelism = std::thread::available_parallelism().map_or(1, NonZeroUsize::get);
+    assert_eq!(rt.metrics().num_workers(), parallelism);
+    assert_eq!(source, WorkerSource::AvailableParallelism);
+}
+
+#[test]
+fn the_config_wins_over_the_environment() {
+    let (rt, source) = runtime("runtime:\n  worker_threads: 1\n")
+        .unwrap()
+        .build_with(Some(OsStr::new("5")))
+        .unwrap();
+    assert_eq!(rt.metrics().num_workers(), 1);
+    assert_eq!(source, WorkerSource::Config);
+}
+
+#[test]
+fn an_invalid_worker_count_names_the_key() {
+    // 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}");
+    }
+}
+
+#[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", "", "-2"] {
+        let err = RuntimeConfig::default()
+            .build_with(Some(OsStr::new(value)))
+            .unwrap_err()
+            .to_string();
+        assert!(err.contains(WORKER_THREADS_ENV), "{value:?}: {err}");
+    }
+}
+
+#[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");
+    assert_eq!(
+        WorkerSource::Environment.to_string(),
+        "TOKIO_WORKER_THREADS"
+    );
+    assert_eq!(
+        WorkerSource::AvailableParallelism.to_string(),
+        "available parallelism"
+    );
+}