diff --git a/crates/kumo-server-runtime/src/lib.rs b/crates/kumo-server-runtime/src/lib.rs index 86a3a0a8..141627c6 100644 --- a/crates/kumo-server-runtime/src/lib.rs +++ b/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( + name_prefix: &'static str, + default_size: SIZE, + configured_size: &AtomicUsize, + func_factory: FUNC, +) -> anyhow::Result +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( 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::(); - 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, diff --git a/crates/kumod/src/spool.rs b/crates/kumod/src/spool.rs index 84f71c2d..0ebf1831 100644 --- a/crates/kumod/src/spool.rs +++ b/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) => {