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

@@ -2738,6 +2738,7 @@ dependencies = [
  "async-channel 1.9.0",
  "kumo-server-memory",
  "lazy_static",
+ "prometheus",
  "tokio",
  "tracing",
 ]

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

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

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"

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<Self> {
+    pub fn new<F>(
+        name_prefix: &'static str,
+        default_size: F,
+        configured_size: &AtomicUsize,
+    ) -> anyhow::Result<Self>
+    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::<Command>();
 
+        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 {

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,

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)| {

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 {

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<ReadyQueueManager> = 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();

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<Mutex<LruCacheWithTtl<DomainAndListener, Option<EsmtpDomain>>>> =
     Lazy::new(|| Mutex::new(LruCacheWithTtl::new(1024)));
 
+static SMTPSRV: Lazy<Runtime> =
+    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) =

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,