diff --git a/crates/kumo-server-runtime/src/lib.rs b/crates/kumo-server-runtime/src/lib.rs index aa6bf107..ae1a49ff 100644 --- a/crates/kumo-server-runtime/src/lib.rs +++ b/crates/kumo-server-runtime/src/lib.rs @@ -16,7 +16,7 @@ use std::future::Future; use tokio::task::{JoinHandle, LocalSet}; lazy_static::lazy_static! { - static ref RUNTIME: Runtime = Runtime::new().unwrap(); + static ref RUNTIME: Runtime = Runtime::new("localset").unwrap(); } enum Command { @@ -28,7 +28,7 @@ pub struct Runtime { } impl Runtime { - pub fn new() -> anyhow::Result { + pub fn new(name_prefix: &'static str) -> anyhow::Result { let n_threads = match std::env::var("TOKIO_WORKER_THREADS") { Ok(n) => n.parse()?, Err(_) => std::thread::available_parallelism()?, @@ -38,7 +38,7 @@ impl Runtime { for n in 0..n_threads.into() { let rx = rx.clone(); std::thread::Builder::new() - .name(format!("localset-{n}")) + .name(format!("{name_prefix}-{n}")) .spawn(move || { let runtime = tokio::runtime::Builder::new_current_thread() .enable_io() @@ -61,6 +61,68 @@ impl Runtime { Ok(Self { jobs: tx }) } + + /// Schedule func to run in the runtime pool. + /// func must return a future; that future will be spawned into the thread-local + /// executor. + /// This function will return the result of the spawn attempt, but will not + /// wait for the future it spawns to complete. + /// + /// This function is useful for getting into a localset environment where + /// !Send futures can be scheduled when you are not already in such an environment. + /// + /// If you are already in a !Send future, then using spawn_local below has + /// less overhead. + pub async fn spawn anyhow::Result) + Send + 'static, FUT>( + &self, + name: String, + func: F, + ) -> anyhow::Result> + where + FUT: Future + 'static, + FUT::Output: Send, + { + let (tx, rx) = bounded::>>(1); + self.jobs + .try_send(Command::Run(Box::new(move || match (func)() { + Ok(future) => { + tx.try_send( + tokio::task::Builder::new() + .name(&name) + .spawn_local(future) + .map_err(|e| e.into()), + ) + .ok(); + } + Err(err) => { + tx.try_send(Err(err)).ok(); + } + }))) + .map_err(|err| anyhow::anyhow!("failed to send func to runtime thread: {err:#}"))?; + rx.recv().await? + } + + pub fn spawn_non_blocking anyhow::Result) + Send + 'static, FUT>( + &self, + name: String, + func: F, + ) -> anyhow::Result<()> + where + FUT: Future + 'static, + { + self.jobs + .try_send(Command::Run(Box::new(move || match (func)() { + Ok(future) => { + if let Err(err) = tokio::task::Builder::new().name(&name).spawn_local(future) { + tracing::error!("rt_spawn_non_blocking: error: {err:#}"); + } + } + Err(err) => { + tracing::error!("rt_spawn_non_blocking: error: {err:#}"); + } + }))) + .map_err(|err| anyhow::anyhow!("failed to send func to runtime thread: {err:#}")) + } } /// Schedule func to run in the runtime pool. @@ -82,25 +144,7 @@ where FUT: Future + 'static, FUT::Output: Send, { - let (tx, rx) = bounded::>>(1); - RUNTIME - .jobs - .try_send(Command::Run(Box::new(move || match (func)() { - Ok(future) => { - tx.try_send( - tokio::task::Builder::new() - .name(&name) - .spawn_local(future) - .map_err(|e| e.into()), - ) - .ok(); - } - Err(err) => { - tx.try_send(Err(err)).ok(); - } - }))) - .map_err(|err| anyhow::anyhow!("failed to send func to runtime thread: {err:#}"))?; - rx.recv().await? + RUNTIME.spawn(name, func).await } pub fn rt_spawn_non_blocking anyhow::Result) + Send + 'static, FUT>( @@ -110,19 +154,7 @@ pub fn rt_spawn_non_blocking anyhow::Result) + Send + 'stat where FUT: Future + 'static, { - RUNTIME - .jobs - .try_send(Command::Run(Box::new(move || match (func)() { - Ok(future) => { - if let Err(err) = tokio::task::Builder::new().name(&name).spawn_local(future) { - tracing::error!("rt_spawn_non_blocking: error: {err:#}"); - } - } - Err(err) => { - tracing::error!("rt_spawn_non_blocking: error: {err:#}"); - } - }))) - .map_err(|err| anyhow::anyhow!("failed to send func to runtime thread: {err:#}")) + RUNTIME.spawn_non_blocking(name, func) } /// Spawn a future as a task with a name. diff --git a/crates/kumod/src/lua_deliver.rs b/crates/kumod/src/lua_deliver.rs index fbf35cd2..453b5c59 100644 --- a/crates/kumod/src/lua_deliver.rs +++ b/crates/kumod/src/lua_deliver.rs @@ -6,7 +6,7 @@ use crate::spool::SpoolManager; use async_trait::async_trait; use config::{CallbackSignature, LuaConfig}; use kumo_log_types::{RecordType, ResolvedAddress}; -use kumo_server_runtime::{rt_spawn, spawn}; +use kumo_server_runtime::{spawn, spawn_local}; use message::message::QueueNameComponents; use message::Message; use mlua::{RegistryKey, Value}; @@ -145,10 +145,10 @@ impl QueueDispatcher for LuaQueueDispatcher { tls_info: None, }) .await; - rt_spawn("requeue message".to_string(), move || { - Ok(async move { Dispatcher::requeue_message(msg, true, None).await }) - }) - .await?; + spawn_local( + "requeue message".to_string(), + Dispatcher::requeue_message(msg, true, None), + )?; } dispatcher.metrics.inc_transfail(); } else { diff --git a/crates/kumod/src/queue.rs b/crates/kumod/src/queue.rs index d1c0f074..dd46e26a 100644 --- a/crates/kumod/src/queue.rs +++ b/crates/kumod/src/queue.rs @@ -11,7 +11,7 @@ use chrono::Utc; use config::{load_config, CallbackSignature, LuaConfig}; use kumo_server_common::config_handle::ConfigHandle; use kumo_server_lifecycle::{Activity, ShutdownSubcription}; -use kumo_server_runtime::{rt_spawn, rt_spawn_non_blocking, spawn, spawn_blocking}; +use kumo_server_runtime::{spawn, spawn_blocking, Runtime}; use message::message::QueueNameComponents; use message::Message; use mlua::prelude::*; @@ -32,6 +32,7 @@ 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 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"); @@ -431,14 +432,15 @@ impl Queue { }); let queue_clone = handle.clone(); - rt_spawn(format!("maintain {name}"), move || { - Ok(async move { - if let Err(err) = maintain_named_queue(&queue_clone).await { - tracing::error!("maintain_named_queue {}: {err:#}", queue_clone.name); - } + MAINT_RUNTIME + .spawn(format!("maintain {name}"), move || { + Ok(async move { + if let Err(err) = maintain_named_queue(&queue_clone).await { + tracing::error!("maintain_named_queue {}: {err:#}", queue_clone.name); + } + }) }) - }) - .await?; + .await?; Ok(handle) } @@ -457,8 +459,9 @@ 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 = - rt_spawn_non_blocking("bounce_all remove_from_spool".to_string(), move || { + let result = MAINT_RUNTIME.spawn_non_blocking( + "bounce_all remove_from_spool".to_string(), + move || { Ok(async move { for msg in msgs { let msg = (*msg).clone(); @@ -467,7 +470,8 @@ impl Queue { SpoolManager::remove_from_spool(id).await.ok(); } }) - }); + }, + ); if let Err(err) = result { tracing::error!("Unable to schedule spool removal for {count} messages! {err:#}"); } diff --git a/crates/kumod/src/ready_queue.rs b/crates/kumod/src/ready_queue.rs index 5e3cd108..3ffdf603 100644 --- a/crates/kumod/src/ready_queue.rs +++ b/crates/kumod/src/ready_queue.rs @@ -15,7 +15,7 @@ use kumo_api_types::egress_path::EgressPathConfig; use kumo_server_common::config_handle::ConfigHandle; use kumo_server_lifecycle::{Activity, ShutdownSubcription}; use kumo_server_memory::{get_headroom, low_memory, subscribe_to_memory_status_changes}; -use kumo_server_runtime::{rt_spawn, rt_spawn_non_blocking, spawn}; +use kumo_server_runtime::{spawn, Runtime}; use message::message::QueueNameComponents; use message::Message; use rfc5321::{EnhancedStatusCode, Response}; @@ -33,6 +33,7 @@ lazy_static::lazy_static! { static ref MANAGER: Mutex = Mutex::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 struct ReadyQueueName { @@ -201,12 +202,13 @@ impl ReadyQueueManager { let handle = manager.queues.entry(name.clone()).or_insert_with(|| { let notify = Arc::new(Notify::new()); - rt_spawn_non_blocking(format!("maintain {name}"), { - let name = name.clone(); - let notify = notify.clone(); - move || Ok(async move { Self::maintainer_task(name, notify).await }) - }) - .expect("failed to spawn maintainer"); + READYQ_RUNTIME + .spawn_non_blocking(format!("maintain {name}"), { + let name = name.clone(); + let notify = notify.clone(); + move || Ok(async move { Self::maintainer_task(name, notify).await }) + }) + .expect("failed to spawn maintainer"); let proto = queue_config.borrow().protocol.metrics_protocol_name(); let service = format!("{proto}:{name}"); let metrics = DeliveryMetrics::new(&service, &proto); @@ -433,21 +435,22 @@ impl ReadyQueue { let msgs = Self::take_ready_queue(&self.ready, &self.metrics); if !msgs.is_empty() { let activity = self.activity.clone(); - rt_spawn( - format!("reinserting {} due to {reason}", self.name), - move || { - Ok(async move { - for msg in msgs { - if let Err(err) = Dispatcher::reinsert_message(msg).await { - tracing::error!("error reinserting message: {err:#}"); + READYQ_RUNTIME + .spawn( + format!("reinserting {} due to {reason}", self.name), + move || { + Ok(async move { + for msg in msgs { + if let Err(err) = Dispatcher::reinsert_message(msg).await { + tracing::error!("error reinserting message: {err:#}"); + } } - } - drop(activity); - }) - }, - ) - .await - .expect("failed to spawn reinsertion"); + drop(activity); + }) + }, + ) + .await + .expect("failed to spawn reinsertion"); } } @@ -529,32 +532,33 @@ impl ReadyQueue { self.consecutive_connection_failures.clone(); tracing::trace!("spawning client for {name}"); - if let Ok(handle) = rt_spawn(format!("smtp client {name}"), move || { - Ok(async move { - if let Err(err) = Dispatcher::run( - &name, - queue_name_for_config_change_purposes_only, - mx, - ready, - notify, - queue_config, - path_config, - metrics, - consecutive_connection_failures.clone(), - egress_source, - egress_pool, - lease, - ) - .await - { - tracing::debug!( - "Error in Dispatcher::run for {name}: {err:#} \ + if let Ok(handle) = READYQ_RUNTIME + .spawn(format!("smtp client {name}"), move || { + Ok(async move { + if let Err(err) = Dispatcher::run( + &name, + queue_name_for_config_change_purposes_only, + mx, + ready, + notify, + queue_config, + path_config, + metrics, + consecutive_connection_failures.clone(), + egress_source, + egress_pool, + lease, + ) + .await + { + tracing::debug!( + "Error in Dispatcher::run for {name}: {err:#} \ (consecutive_connection_failures={consecutive_connection_failures:?})" - ); - } + ); + } + }) }) - }) - .await + .await { connections.push(handle); } @@ -623,16 +627,18 @@ impl Drop for Dispatcher { // Ensure that we re-queue any message that we had popped if let Some(msg) = self.msg.take() { let activity = self.activity.clone(); - rt_spawn_non_blocking("Dispatcher::drop".to_string(), move || { - Ok(async move { - if activity.is_shutting_down() { - Queue::save_if_needed_and_log(&msg).await; - } else if let Err(err) = Dispatcher::requeue_message(msg, false, None).await { - tracing::error!("error requeuing message: {err:#}"); - } + READYQ_RUNTIME + .spawn_non_blocking("Dispatcher::drop".to_string(), move || { + Ok(async move { + if activity.is_shutting_down() { + Queue::save_if_needed_and_log(&msg).await; + } else if let Err(err) = Dispatcher::requeue_message(msg, false, None).await + { + tracing::error!("error requeuing message: {err:#}"); + } + }) }) - }) - .ok(); + .ok(); } } } @@ -967,18 +973,19 @@ impl Dispatcher { msgs.len() ); let activity = self.activity.clone(); - rt_spawn("reinserting".to_string(), move || { - Ok(async move { - for msg in msgs { - if let Err(err) = Self::reinsert_message(msg).await { - tracing::error!("error reinserting message: {err:#}"); + READYQ_RUNTIME + .spawn("reinserting".to_string(), move || { + Ok(async move { + for msg in msgs { + if let Err(err) = Self::reinsert_message(msg).await { + tracing::error!("error reinserting message: {err:#}"); + } } - } - drop(activity); + drop(activity); + }) }) - }) - .await - .expect("failed to spawn reinsertion"); + .await + .expect("failed to spawn reinsertion"); } } @@ -1000,18 +1007,19 @@ impl Dispatcher { ); kumo_chrono_helper::MINUTE }); - rt_spawn("requeue for throttle".to_string(), move || { - Ok(async move { - for msg in msgs { - if let Err(err) = Self::requeue_message(msg, false, Some(delay)).await { - tracing::error!("error requeuing message: {err:#}"); + READYQ_RUNTIME + .spawn("requeue for throttle".to_string(), move || { + Ok(async move { + for msg in msgs { + if let Err(err) = Self::requeue_message(msg, false, Some(delay)).await { + tracing::error!("error requeuing message: {err:#}"); + } } - } - drop(activity); + drop(activity); + }) }) - }) - .await - .expect("failed to spawn requeue"); + .await + .expect("failed to spawn requeue"); } } @@ -1035,49 +1043,50 @@ impl Dispatcher { } else { self.metrics.inc_fail_by(msgs.len()); } - rt_spawn( - format!("bulk queue op for {} msgs {name} {response:?}", msgs.len()), - move || { - Ok(async move { - let increment_attempts = true; - for msg in msgs { - log_disposition(LogDisposition { - kind: if response.is_transient() { - RecordType::TransientFailure - } else { - RecordType::Bounce - }, - msg: msg.clone(), - site: &name, - peer_address: None, - response: response.clone(), - egress_pool: Some(&egress_pool), - egress_source: Some(&egress_source), - relay_disposition: None, - delivery_protocol: None, - tls_info: None, - }) - .await; - - if response.is_transient() { - if let Err(err) = - Self::requeue_message(msg, increment_attempts, None).await - { - tracing::error!("error requeuing message: {err:#}"); - } - } else if response.is_permanent() { - spawn("remove msg from spool", async move { - SpoolManager::remove_from_spool(*msg.id()).await + READYQ_RUNTIME + .spawn( + format!("bulk queue op for {} msgs {name} {response:?}", msgs.len()), + move || { + Ok(async move { + let increment_attempts = true; + for msg in msgs { + log_disposition(LogDisposition { + kind: if response.is_transient() { + RecordType::TransientFailure + } else { + RecordType::Bounce + }, + msg: msg.clone(), + site: &name, + peer_address: None, + response: response.clone(), + egress_pool: Some(&egress_pool), + egress_source: Some(&egress_source), + relay_disposition: None, + delivery_protocol: None, + tls_info: None, }) - .ok(); + .await; + + if response.is_transient() { + if let Err(err) = + Self::requeue_message(msg, increment_attempts, None).await + { + tracing::error!("error requeuing message: {err:#}"); + } + } else if response.is_permanent() { + spawn("remove msg from spool", async move { + SpoolManager::remove_from_spool(*msg.id()).await + }) + .ok(); + } } - } - drop(activity); - }) - }, - ) - .await - .expect("bulk queue spawned"); + drop(activity); + }) + }, + ) + .await + .expect("bulk queue spawned"); } } diff --git a/crates/kumod/src/smtp_dispatcher.rs b/crates/kumod/src/smtp_dispatcher.rs index 26eef695..d766d56c 100644 --- a/crates/kumod/src/smtp_dispatcher.rs +++ b/crates/kumod/src/smtp_dispatcher.rs @@ -9,7 +9,7 @@ use dns_resolver::{resolve_a_or_aaaa, ResolvedMxAddresses}; use kumo_api_types::egress_path::Tls; use kumo_log_types::ResolvedAddress; use kumo_server_lifecycle::ShutdownSubcription; -use kumo_server_runtime::{rt_spawn, spawn}; +use kumo_server_runtime::{spawn, spawn_local}; use message::message::QueueNameComponents; use message::Message; use mta_sts::policy::PolicyMode; @@ -546,10 +546,10 @@ impl QueueDispatcher for SmtpDispatcher { tls_info: self.tls_info.as_ref(), }) .await; - rt_spawn("requeue message".to_string(), move || { - Ok(async move { Dispatcher::requeue_message(msg, true, None).await }) - }) - .await?; + spawn_local( + "requeue message".to_string(), + Dispatcher::requeue_message(msg, true, None), + )?; } dispatcher.metrics.inc_transfail(); } else if response.code >= 200 && response.code < 300 { @@ -631,10 +631,10 @@ impl QueueDispatcher for SmtpDispatcher { tls_info: self.tls_info.as_ref(), }) .await; - rt_spawn("requeue message".to_string(), move || { - Ok(async move { Dispatcher::requeue_message(msg, true, None).await }) - }) - .await?; + spawn_local( + "requeue message".to_string(), + Dispatcher::requeue_message(msg, true, None), + )?; } dispatcher.metrics.inc_transfail(); // Move on to the next host @@ -672,10 +672,10 @@ impl QueueDispatcher for SmtpDispatcher { tls_info: self.tls_info.as_ref(), }) .await; - rt_spawn("requeue message".to_string(), move || { - Ok(async move { Dispatcher::requeue_message(msg, true, None).await }) - }) - .await?; + spawn_local( + "requeue message".to_string(), + Dispatcher::requeue_message(msg, true, None), + )?; } dispatcher.metrics.inc_transfail(); // Move on to the next host