From 89d98cd3ad5e7eac8d0e854043b45c2ca777bd69 Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Fri, 3 May 2024 18:05:57 -0700 Subject: [PATCH] expand thread pool/runtime concept some more As much as I'd love for tokio to be magic and figure everything out, it doesn't and cannot do that. In any complex system it is important to be able to control and provision the various stages so that a certain QoS is met. This commit expands the use of the various thread pools, and adjusts them from always spawning num-cpus to a more balanced set of sizes, so that not everything is competing with everything. These sizes are configurable via env vars, but a preferably configured during the `init` event via a handful of new lua functions that must be called very early on. More docs on these will appear in a later commit. The default pool sizes are a function of the number of cores, and are intended to balance receiving and sending functions such that a throughput test should run without any significant build up of the scheduled queue. This commit also exposes the sizes of the thread pools and the number of parked (idle) threads via prometheus. --- Cargo.lock | 1 + crates/integration-tests/src/kumod.rs | 5 +- crates/integration-tests/src/tsa.rs | 5 +- crates/kumo-server-runtime/Cargo.toml | 1 + crates/kumo-server-runtime/src/lib.rs | 68 ++++++++++++++++++++++++--- crates/kumod/src/main.rs | 2 +- crates/kumod/src/mod_kumo.rs | 24 ++++++++++ crates/kumod/src/queue.rs | 16 +++++-- crates/kumod/src/ready_queue.rs | 13 +++-- crates/kumod/src/smtp_server.rs | 14 +++++- crates/tsa-daemon/src/main.rs | 2 +- 11 files changed, 132 insertions(+), 19 deletions(-) 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,