spoolin: improve thread pool saturation

Previously, it was possible for the spool in tasks to end up
being spawned on just a subset of the available spool in threads
if there was some idleness in the task processing startup, for
example, if the dns for the first few messages returned from
spool enumeration is slow to resolve.

In that situation we can end up with no effective concurrency
during spool enumeration, leading to a very slow startup.

What I'd like to see to resolve this wholistically is adopting
the main tokio work stealing task runner, but we are prevented
from doing this until mlua 0.10 is released.

What this commit does is refactor the core of the Runtime
code to extract the function that sets up the thread pool so
that we can directly spawn the spool in thread logic into
each of the worker threads, guaranteeing that they are spread
out one to a thread.

With this change in place, I always observe 100% utilization
of spoolin on startup where I previously would see only around
60 or 70%.

crates/kumo-server-runtime/src/lib.rs

@@ -15,7 +15,7 @@ use async_channel::{bounded, unbounded, Sender};
 use prometheus::IntGaugeVec;
 use std::future::Future;
 use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::LazyLock;
+use std::sync::{Arc, LazyLock};
 use tokio::runtime::Handle;
 use tokio::task::{JoinHandle, LocalSet};
 
@@ -77,6 +77,99 @@ impl Drop for Runtime {
     }
 }
 
+pub fn spawn_simple_worker_pool<SIZE, FUNC, FUT>(
+    name_prefix: &'static str,
+    default_size: SIZE,
+    configured_size: &AtomicUsize,
+    func_factory: FUNC,
+) -> anyhow::Result<usize>
+where
+    SIZE: FnOnce(usize) -> usize,
+    FUNC: (Fn() -> FUT) + Send + Sync + 'static,
+    FUT: Future + 'static,
+    FUT::Output: Send,
+{
+    let env_name = format!("KUMOD_{}_THREADS", name_prefix.to_uppercase());
+    let n_threads = match std::env::var(env_name) {
+        Ok(n) => n.parse()?,
+        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 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);
+
+    let func_factory = Arc::new(func_factory);
+
+    for n in 0..n_threads.into() {
+        let num_parked = num_parked.clone();
+        let func_factory = func_factory.clone();
+        std::thread::Builder::new()
+            .name(format!("{name_prefix}-{n}"))
+            .spawn(move || {
+                let runtime = tokio::runtime::Builder::new_current_thread()
+                    .enable_io()
+                    .enable_time()
+                    .event_interval(
+                        std::env::var("KUMOD_EVENT_INTERVAL")
+                            .ok()
+                            .and_then(|n| n.parse().ok())
+                            .unwrap_or(61),
+                    )
+                    .max_io_events_per_tick(
+                        std::env::var("KUMOD_IO_EVENTS_PER_TICK")
+                            .ok()
+                            .and_then(|n| n.parse().ok())
+                            .unwrap_or(1024),
+                    )
+                    .on_thread_park({
+                        let num_parked = num_parked.clone();
+                        move || {
+                            kumo_server_memory::purge_thread_cache();
+                            num_parked.inc();
+                        }
+                    })
+                    .thread_name(format!("{name_prefix}-blocking"))
+                    .max_blocking_threads(
+                        std::env::var(format!(
+                            "KUMOD_{}_MAX_BLOCKING_THREADS",
+                            name_prefix.to_uppercase()
+                        ))
+                        .ok()
+                        .and_then(|n| n.parse().ok())
+                        .unwrap_or(512),
+                    )
+                    .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!");
+                    let func = (func_factory)();
+                    (func).await
+                });
+            })?;
+    }
+    Ok(n_threads)
+}
+
 impl Runtime {
     pub fn new<F>(
         name_prefix: &'static str,
@@ -86,86 +179,20 @@ impl Runtime {
     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(_) => {
-                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()
-                        .event_interval(
-                            std::env::var("KUMOD_EVENT_INTERVAL")
-                                .ok()
-                                .and_then(|n| n.parse().ok())
-                                .unwrap_or(61),
-                        )
-                        .max_io_events_per_tick(
-                            std::env::var("KUMOD_IO_EVENTS_PER_TICK")
-                                .ok()
-                                .and_then(|n| n.parse().ok())
-                                .unwrap_or(1024),
-                        )
-                        .on_thread_park({
-                            let num_parked = num_parked.clone();
-                            move || {
-                                kumo_server_memory::purge_thread_cache();
-                                num_parked.inc();
-                            }
-                        })
-                        .thread_name(format!("{name_prefix}-blocking"))
-                        .max_blocking_threads(
-                            std::env::var(format!(
-                                "KUMOD_{}_MAX_BLOCKING_THREADS",
-                                name_prefix.to_uppercase()
-                            ))
-                            .ok()
-                            .and_then(|n| n.parse().ok())
-                            .unwrap_or(512),
-                        )
-                        .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");
+        let n_threads =
+            spawn_simple_worker_pool(name_prefix, default_size, configured_size, move || {
+                let rx = rx.clone();
+                async move {
+                    while let Ok(cmd) = rx.recv().await {
+                        match cmd {
+                            Command::Run(func) => (func)(),
                         }
-                        tracing::trace!("{name_prefix}-{n} started up!");
-                        while let Ok(cmd) = rx.recv().await {
-                            match cmd {
-                                Command::Run(func) => (func)(),
-                            }
-                        }
-                    });
-                })?;
-        }
+                    }
+                    Ok::<(), anyhow::Error>(())
+                }
+            })?;
 
         Ok(Self {
             jobs: tx,

crates/kumod/src/spool.rs

@@ -5,7 +5,7 @@ use chrono::Utc;
 use config::{any_err, from_lua_value, get_or_create_module, CallbackSignature};
 use kumo_server_common::disk_space::{MinFree, MonitoredPath};
 use kumo_server_lifecycle::{Activity, LifeCycle, ShutdownSubcription};
-use kumo_server_runtime::{spawn, Runtime};
+use kumo_server_runtime::{spawn, spawn_simple_worker_pool};
 use message::Message;
 use mlua::{Lua, Value};
 use rfc5321::{EnhancedStatusCode, Response};
@@ -472,29 +472,25 @@ impl SpoolManager {
 
         let (complete_tx, complete_rx) = flume::bounded(1);
 
-        let mut num_tasks = 0;
-        let spool_in = Runtime::new("spoolin", |cpus| cpus / 2, &SPOOLIN_THREADS)?;
-        let n_threads = spool_in.get_num_threads();
+        let spooled_in_clone = Arc::clone(&spooled_in);
+        let n_threads = spawn_simple_worker_pool(
+            "spoolin",
+            |cpus| cpus / 2,
+            &SPOOLIN_THREADS,
+            move || {
+                let spooled_in = Arc::clone(&spooled_in_clone);
+                let rx = rx.clone();
+                let complete_tx = complete_tx.clone();
+                async move {
+                    let mgr = Self::get();
+                    let result = mgr.spool_in_thread(rx, spooled_in).await;
+                    complete_tx.send_async(result).await
+                }
+            },
+        )?;
+        let mut num_tasks = n_threads;
         tracing::info!("Using concurrency {n_threads} for spooling in");
 
-        for n in 0..n_threads {
-            let spooled_in = Arc::clone(&spooled_in);
-            let rx = rx.clone();
-            let complete_tx = complete_tx.clone();
-            spool_in
-                .spawn(format!("spool_in-{n}"), move || {
-                    Ok(async move {
-                        let mgr = Self::get();
-                        let result = mgr.spool_in_thread(rx, spooled_in).await;
-                        complete_tx.send_async(result).await
-                    })
-                })
-                .await?;
-            num_tasks += 1;
-        }
-
-        drop(complete_tx);
-
         while num_tasks > 0 {
             tokio::select! {
                 _ = tokio::time::sleep(interval) => {