diff --git a/crates/kumo-server-common/src/lib.rs b/crates/kumo-server-common/src/lib.rs index 19fb7460..a38e212c 100644 --- a/crates/kumo-server-common/src/lib.rs +++ b/crates/kumo-server-common/src/lib.rs @@ -2,6 +2,7 @@ use config::{ any_err, decorate_callback_name, from_lua_value, get_or_create_module, load_config, CallbackSignature, }; +use kumo_server_runtime::available_parallelism; use mlua::{Function, Lua, LuaSerdeExt, Value, Variadic}; use mod_redis::RedisConnKey; use serde::{Deserialize, Serialize}; @@ -327,9 +328,7 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { kumo_mod.set( "available_parallelism", - lua.create_function(move |_, _: ()| { - Ok(std::thread::available_parallelism().map_err(any_err)?.get()) - })?, + lua.create_function(move |_, _: ()| available_parallelism().map_err(any_err))?, )?; kumo_mod.set( diff --git a/crates/kumo-server-runtime/src/lib.rs b/crates/kumo-server-runtime/src/lib.rs index 92cd3547..08d8218e 100644 --- a/crates/kumo-server-runtime/src/lib.rs +++ b/crates/kumo-server-runtime/src/lib.rs @@ -1,3 +1,4 @@ +use anyhow::Context; use parking_lot::Mutex; use prometheus::IntGaugeVec; use std::collections::HashMap; @@ -84,7 +85,7 @@ impl Runtime { Err(_) => { let configured = configured_size.load(Ordering::SeqCst); if configured == 0 { - let cpus = std::thread::available_parallelism()?.get(); + let cpus = available_parallelism()?; (default_size)(cpus).max(1) } else { configured @@ -239,3 +240,14 @@ where .name(name.as_ref()) .spawn_blocking_on(func, runtime) } + +pub fn available_parallelism() -> anyhow::Result { + match std::env::var("KUMO_AVAILABLE_PARALLELISM") { + Ok(n) => n + .parse() + .context("failed to parse KUMO_AVAILABLE_PARALLELISM as a number"), + Err(_) => Ok(std::thread::available_parallelism() + .context("failed to get available_parallelism")? + .get()), + } +} diff --git a/crates/kumod/src/logging/classify.rs b/crates/kumod/src/logging/classify.rs index 7f71fb60..92675f51 100644 --- a/crates/kumod/src/logging/classify.rs +++ b/crates/kumod/src/logging/classify.rs @@ -4,6 +4,7 @@ use bounce_classify::{ }; use config::epoch::{get_current_epoch, ConfigEpoch}; use kumo_log_types::JsonLogRecord; +use kumo_server_runtime::available_parallelism; use lru_cache::LruCache; use parking_lot::Mutex; use prometheus::Histogram; @@ -45,9 +46,7 @@ impl ClassifierParams { } fn default_pool_size() -> usize { - std::thread::available_parallelism() - .map(|p| (p.get() / 4).max(1)) - .unwrap_or(4) + available_parallelism().map(|p| (p / 4).max(1)).unwrap_or(4) } fn default_cache_size() -> usize { diff --git a/crates/kumod/src/main.rs b/crates/kumod/src/main.rs index bcd048a3..949f7562 100644 --- a/crates/kumod/src/main.rs +++ b/crates/kumod/src/main.rs @@ -7,6 +7,7 @@ use config::{declare_event, CallbackSignature}; use kumo_server_common::diagnostic_logging::{DiagnosticFormat, LoggingConfig}; use kumo_server_common::start::StartConfig; use kumo_server_lifecycle::LifeCycle; +use kumo_server_runtime::available_parallelism; use nix::sys::resource::{getrlimit, setrlimit, Resource}; use nix::unistd::{Uid, User}; use std::path::PathBuf; @@ -179,7 +180,7 @@ fn main() -> anyhow::Result<()> { let n_threads = match std::env::var("KUMOD_MAIN_THREADS") { Ok(n) => n.parse()?, - Err(_) => std::thread::available_parallelism()?.get(), + Err(_) => available_parallelism()?, }; tokio::runtime::Builder::new_multi_thread() @@ -221,9 +222,7 @@ fn main() -> anyhow::Result<()> { async fn perform_init(opts: Opt) -> anyhow::Result<()> { let nodeid = kumo_server_common::nodeid::NodeId::get(); tracing::info!("NodeId is {nodeid}"); - let num_cores = std::thread::available_parallelism() - .context("failed to get available_parallelism")? - .get(); + let num_cores = available_parallelism()?; tracing::info!("available_parallelism={num_cores}"); if num_cores < 4 { tracing::error!( diff --git a/crates/tsa-daemon/src/publish.rs b/crates/tsa-daemon/src/publish.rs index a4afa1c4..2409403c 100644 --- a/crates/tsa-daemon/src/publish.rs +++ b/crates/tsa-daemon/src/publish.rs @@ -1,5 +1,6 @@ use crate::http_server::{open_history_db, publish_log_batch}; use kumo_log_types::JsonLogRecord; +use kumo_server_runtime::available_parallelism; use parking_lot::Mutex; use std::sync::{Arc, LazyLock}; use tokio::sync::Notify; @@ -139,7 +140,7 @@ async fn flush_batches() { } fn start_processor_pool() -> anyhow::Result> { - let n_threads: usize = std::thread::available_parallelism()?.into(); + let n_threads = available_parallelism()?; for i in 0..n_threads { std::thread::Builder::new() diff --git a/docs/changelog/main.md b/docs/changelog/main.md index 68f06aae..97711e8f 100644 --- a/docs/changelog/main.md +++ b/docs/changelog/main.md @@ -34,6 +34,12 @@ changes in behavior around DNS. * Added example and recommended default shaping configuration for the TSA daemon to the default `shaping.toml` file. +* If you are running on a system where + [kumo.available_parallelism](../reference/kumo/available_parallelism.md) + returns an inaccurate value (such as an overcommitted VM), then you may + now export `KUMO_AVAILABLE_PARALLELISM` into the launching environment to + override the value to something more appropriate, which helps to scale + the various thread pools more appropriately. ## Fixes diff --git a/docs/reference/kumo/available_parallelism.md b/docs/reference/kumo/available_parallelism.md index 22606044..fc07918c 100644 --- a/docs/reference/kumo/available_parallelism.md +++ b/docs/reference/kumo/available_parallelism.md @@ -58,3 +58,10 @@ The following platform limitations currently apply to `kumo.available_parallelis *This documentation is excerpted from the underlying [Rust function](https://doc.rust-lang.org/stable/std/thread/fn.available_parallelism.html).* + +!!! note + If `kumo.available_parallelism` returns an incorrect estimate for your + system (for the reasons mentioned above), you can force it to return a + different by value by setting the `KUMO_AVAILABLE_PARALLELISM` + environment variable to the desired value in the environment when you + launch `kumod` or `tsa-daemon`. {{since('dev', inline=True)}}.