diff --git a/Cargo.lock b/Cargo.lock index a36ca55c..533d614e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2738,6 +2738,7 @@ dependencies = [ "async-channel 1.9.0", "kumo-server-memory", "lazy_static", + "prometheus", "tokio", "tracing", ] diff --git a/crates/integration-tests/src/kumod.rs b/crates/integration-tests/src/kumod.rs index 4282f5ee..d27524bc 100644 --- a/crates/integration-tests/src/kumod.rs +++ b/crates/integration-tests/src/kumod.rs @@ -299,7 +299,10 @@ impl KumoDaemon { let mut cmd = Command::new(&path); cmd.args(["--policy", &args.policy_file, "--user", &user.name]) - .env("KUMOD_LOG", "kumod=trace,kumo_server_common=info") + .env( + "KUMOD_LOG", + "kumod=trace,kumo_server_common=info,kumo_server_runtime=info", + ) .env("KUMOD_TEST_DIR", dir.path()) .envs(args.env.iter().cloned()) .stdout(Stdio::piped()) diff --git a/crates/integration-tests/src/tsa.rs b/crates/integration-tests/src/tsa.rs index fb41e98c..87aba992 100644 --- a/crates/integration-tests/src/tsa.rs +++ b/crates/integration-tests/src/tsa.rs @@ -33,7 +33,10 @@ impl TsaDaemon { let mut cmd = Command::new(&path); cmd.args(["--policy", &args.policy_file]) - .env("KUMO_TSA_LOG", "tsa_daemon=trace,kumo_server_common=info") + .env( + "KUMO_TSA_LOG", + "tsa_daemon=trace,kumo_server_common=info,kumo_server_runtime=info", + ) .env("KUMO_TSA_TEST_DIR", dir.path()) .envs(args.env.iter().cloned()) .stdout(Stdio::piped()) diff --git a/crates/kumo-server-runtime/Cargo.toml b/crates/kumo-server-runtime/Cargo.toml index ed156b7e..a199e144 100644 --- a/crates/kumo-server-runtime/Cargo.toml +++ b/crates/kumo-server-runtime/Cargo.toml @@ -10,5 +10,6 @@ anyhow = "1.0" async-channel = "1.8" kumo-server-memory = {path="../kumo-server-memory"} lazy_static = "1.4" +prometheus = "0.13" tokio = {workspace=true, features=["full", "tracing"]} tracing = "0.1" diff --git a/crates/kumo-server-runtime/src/lib.rs b/crates/kumo-server-runtime/src/lib.rs index c07bea4e..95f6199e 100644 --- a/crates/kumo-server-runtime/src/lib.rs +++ b/crates/kumo-server-runtime/src/lib.rs @@ -12,11 +12,33 @@ //! For example, when accepting new connections, we use this to //! spawn the server processing future. use async_channel::{bounded, unbounded, Sender}; +use prometheus::IntGaugeVec; use std::future::Future; +use std::sync::atomic::{AtomicUsize, Ordering}; use tokio::task::{JoinHandle, LocalSet}; lazy_static::lazy_static! { - static ref RUNTIME: Runtime = Runtime::new("localset").unwrap(); + static ref RUNTIME: Runtime = Runtime::new( + "localset", |cpus| cpus/4, &LOCALSET_THREADS).unwrap(); + + static ref PARKED_THREADS: IntGaugeVec = { + prometheus::register_int_gauge_vec!( + "thread_pool_parked", + "number of parked(idle) threads in a thread pool", + &["pool"]).unwrap() + }; + static ref NUM_THREADS: IntGaugeVec = { + prometheus::register_int_gauge_vec!( + "thread_pool_size", + "number of threads in a thread pool", + &["pool"]).unwrap() + }; +} + +static LOCALSET_THREADS: AtomicUsize = AtomicUsize::new(0); + +pub fn set_localset_threads(n: usize) { + LOCALSET_THREADS.store(n, Ordering::SeqCst); } enum Command { @@ -28,31 +50,63 @@ pub struct Runtime { } impl Runtime { - pub fn new(name_prefix: &'static str) -> anyhow::Result { + pub fn new( + name_prefix: &'static str, + default_size: F, + configured_size: &AtomicUsize, + ) -> anyhow::Result + where + F: FnOnce(usize) -> usize, + { let env_name = format!("KUMOD_{}_THREADS", name_prefix.to_uppercase()); let n_threads = match std::env::var(env_name) { Ok(n) => n.parse()?, - Err(_) => match std::env::var("TOKIO_WORKER_THREADS") { - Ok(n) => n.parse()?, - Err(_) => std::thread::available_parallelism()?, - }, + Err(_) => { + let configured = configured_size.load(Ordering::SeqCst); + if configured == 0 { + let cpus = std::thread::available_parallelism()?.get(); + (default_size)(cpus).max(1) + } else { + configured + } + } }; let (tx, rx) = unbounded::(); + let num_parked = PARKED_THREADS.get_metric_with_label_values(&[name_prefix])?; + let num_threads = NUM_THREADS.get_metric_with_label_values(&[name_prefix])?; + num_threads.set(n_threads as i64); + for n in 0..n_threads.into() { let rx = rx.clone(); + let num_parked = num_parked.clone(); std::thread::Builder::new() .name(format!("{name_prefix}-{n}")) .spawn(move || { let runtime = tokio::runtime::Builder::new_current_thread() .enable_io() .enable_time() - .on_thread_park(|| kumo_server_memory::purge_thread_cache()) + .on_thread_park({ + let num_parked = num_parked.clone(); + move || { + kumo_server_memory::purge_thread_cache(); + num_parked.inc(); + } + }) + .on_thread_unpark({ + let num_parked = num_parked.clone(); + move || { + num_parked.dec(); + } + }) .build() .unwrap(); let local_set = LocalSet::new(); local_set.block_on(&runtime, async move { + if n == 0 { + tracing::info!("{name_prefix} pool starting with {n_threads} threads"); + } tracing::trace!("{name_prefix}-{n} started up!"); while let Ok(cmd) = rx.recv().await { match cmd { diff --git a/crates/kumod/src/main.rs b/crates/kumod/src/main.rs index ce148908..23b39a86 100644 --- a/crates/kumod/src/main.rs +++ b/crates/kumod/src/main.rs @@ -184,7 +184,7 @@ async fn run(opts: Opt) -> anyhow::Result<()> { diag_format: opts.diag_format, tokio_console: opts.tokio_console, filter_env_var: "KUMOD_LOG", - default_filter: "kumod=info,kumo_server_common=info", + default_filter: "kumod=info,kumo_server_common=info,kumo_server_runtime=info", }, lua_funcs: &[ kumo_server_common::register, diff --git a/crates/kumod/src/mod_kumo.rs b/crates/kumod/src/mod_kumo.rs index 79bf2362..807c82ac 100644 --- a/crates/kumod/src/mod_kumo.rs +++ b/crates/kumod/src/mod_kumo.rs @@ -49,6 +49,30 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { })?, )?; + kumo_mod.set( + "set_smtpsrv_threads", + lua.create_function(move |_, limit: usize| { + crate::smtp_server::set_smtpsrv_threads(limit); + Ok(()) + })?, + )?; + + kumo_mod.set( + "set_qmaint_threads", + lua.create_function(move |_, limit: usize| { + crate::queue::set_qmaint_threads(limit); + Ok(()) + })?, + )?; + + kumo_mod.set( + "set_readyq_threads", + lua.create_function(move |_, limit: usize| { + crate::ready_queue::set_readyq_threads(limit); + Ok(()) + })?, + )?; + kumo_mod.set( "reject", lua.create_function(move |_lua, (code, message): (u16, String)| { diff --git a/crates/kumod/src/queue.rs b/crates/kumod/src/queue.rs index debc2301..7770890c 100644 --- a/crates/kumod/src/queue.rs +++ b/crates/kumod/src/queue.rs @@ -20,6 +20,7 @@ use prometheus::{IntGauge, IntGaugeVec}; use rfc5321::{EnhancedStatusCode, Response}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use thiserror::Error; @@ -33,7 +34,10 @@ lazy_static::lazy_static! { static ref DELAY_GAUGE: IntGaugeVec = { prometheus::register_int_gauge_vec!("scheduled_count", "number of messages in the scheduled queue", &["queue"]).unwrap() }; - static ref MAINT_RUNTIME: Runtime = Runtime::new("schedqmaint").unwrap(); + + pub static ref QMAINT_RUNTIME: Runtime = Runtime::new( + "qmaint", |cpus| cpus/4, &QMAINT_THREADS).unwrap(); + pub static ref GET_Q_CONFIG_SIG: CallbackSignature::<'static, (&'static str, Option<&'static str>, Option<&'static str>, Option<&'static str>), QueueConfig> = CallbackSignature::new_with_multiple("get_queue_config"); @@ -42,6 +46,12 @@ lazy_static::lazy_static! { ()> = CallbackSignature::new_with_multiple("throttle_insert_ready_queue"); } +static QMAINT_THREADS: AtomicUsize = AtomicUsize::new(0); + +pub fn set_qmaint_threads(n: usize) { + QMAINT_THREADS.store(n, Ordering::SeqCst); +} + #[derive(Deserialize, Serialize, Debug, Clone)] #[serde(untagged)] pub enum DeliveryProto { @@ -437,7 +447,7 @@ impl Queue { }); let queue_clone = handle.clone(); - MAINT_RUNTIME + QMAINT_RUNTIME .spawn(format!("maintain {name}"), move || { Ok(async move { if let Err(err) = maintain_named_queue(&queue_clone).await { @@ -476,7 +486,7 @@ impl Queue { // reported numbers shown the to initial bounce request will // likely be lower, but it is better for the server to be // healthy than for that command to block and show 100% stats. - let result = MAINT_RUNTIME.spawn_non_blocking( + let result = QMAINT_RUNTIME.spawn_non_blocking( "bounce_all remove_from_spool".to_string(), move || { Ok(async move { diff --git a/crates/kumod/src/ready_queue.rs b/crates/kumod/src/ready_queue.rs index 0bcdb4f9..6eabcbdc 100644 --- a/crates/kumod/src/ready_queue.rs +++ b/crates/kumod/src/ready_queue.rs @@ -4,7 +4,7 @@ use crate::http_server::admin_bounce_v1::AdminBounceEntry; use crate::http_server::admin_suspend_ready_q_v1::AdminSuspendReadyQEntry; use crate::logging::{log_disposition, LogDisposition, RecordType}; use crate::lua_deliver::LuaQueueDispatcher; -use crate::queue::{DeliveryProto, Queue, QueueConfig, QueueManager}; +use crate::queue::{DeliveryProto, Queue, QueueConfig, QueueManager, QMAINT_RUNTIME}; use crate::smtp_dispatcher::{MxListEntry, SmtpDispatcher}; use crate::spool::SpoolManager; use anyhow::Context; @@ -36,7 +36,14 @@ lazy_static::lazy_static! { static ref MANAGER: StdMutex = StdMutex::new(ReadyQueueManager::new()); pub static ref REQUEUE_MESSAGE_SIG: CallbackSignature::<'static, Message, ()> = CallbackSignature::new_with_multiple("message_requeued"); - pub static ref READYQ_RUNTIME: Runtime = Runtime::new("readyq").unwrap(); + pub static ref READYQ_RUNTIME: Runtime = Runtime::new( + "readyq", |cpus| cpus / 2, &READYQ_THREADS).unwrap(); +} + +static READYQ_THREADS: AtomicUsize = AtomicUsize::new(0); + +pub fn set_readyq_threads(n: usize) { + READYQ_THREADS.store(n, Ordering::SeqCst); } pub struct Fifo { @@ -254,7 +261,7 @@ impl ReadyQueueManager { let handle = manager.queues.entry(name.clone()).or_insert_with(|| { let notify_maintainer = Arc::new(Notify::new()); - READYQ_RUNTIME + QMAINT_RUNTIME .spawn_non_blocking(format!("maintain {name}"), { let name = name.clone(); let notify_maintainer = notify_maintainer.clone(); diff --git a/crates/kumod/src/smtp_server.rs b/crates/kumod/src/smtp_server.rs index 3d2a88f2..8375b03c 100644 --- a/crates/kumod/src/smtp_server.rs +++ b/crates/kumod/src/smtp_server.rs @@ -12,7 +12,7 @@ use data_encoding::BASE64; use data_loader::KeySource; use kumo_log_types::ResolvedAddress; use kumo_server_lifecycle::{Activity, ShutdownSubcription}; -use kumo_server_runtime::rt_spawn; +use kumo_server_runtime::Runtime; use lruttl::LruCacheWithTtl; use mailparsing::ConformanceDisposition; use memchr::memmem::Finder; @@ -30,6 +30,7 @@ use spool::SpoolId; use std::fmt::Debug; use std::net::SocketAddr; use std::str::FromStr; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use thiserror::Error; @@ -49,6 +50,15 @@ struct DomainAndListener { static DOMAINS: Lazy>>> = Lazy::new(|| Mutex::new(LruCacheWithTtl::new(1024))); +static SMTPSRV: Lazy = + Lazy::new(|| Runtime::new("smtpsrv", |cpus| cpus * 3 / 8, &SMTPSRV_THREADS).unwrap()); + +static SMTPSRV_THREADS: AtomicUsize = AtomicUsize::new(0); + +pub fn set_smtpsrv_threads(n: usize) { + SMTPSRV_THREADS.store(n, Ordering::SeqCst); +} + #[derive(Deserialize, Clone, Debug, Default, Serialize)] #[serde(deny_unknown_fields)] pub struct EsmtpDomain { @@ -268,7 +278,7 @@ impl EsmtpListenerParams { socket.set_nodelay(true)?; let my_address = socket.local_addr()?; let params = self.clone(); - rt_spawn( + SMTPSRV.spawn( format!("SmtpServer {peer_address:?}"), move || Ok(async move { if let Err(err) = diff --git a/crates/tsa-daemon/src/main.rs b/crates/tsa-daemon/src/main.rs index 9c2ffd7c..0d978fd3 100644 --- a/crates/tsa-daemon/src/main.rs +++ b/crates/tsa-daemon/src/main.rs @@ -99,7 +99,7 @@ async fn run(opts: Opt) -> anyhow::Result<()> { diag_format: opts.diag_format, tokio_console: opts.tokio_console, filter_env_var: "KUMO_TSA_LOG", - default_filter: "tsa_daemon=info,kumo_server_common=info", + default_filter: "tsa_daemon=info,kumo_server_common=info,kumo_server_runtime=info", }, lua_funcs: &[kumo_server_common::register, mod_auto::register], policy: &opts.policy,