diff --git a/crates/kumod/src/egress_path.rs b/crates/kumod/src/egress_path.rs index 58e43b08..aa745b0f 100644 --- a/crates/kumod/src/egress_path.rs +++ b/crates/kumod/src/egress_path.rs @@ -5,6 +5,7 @@ use crate::lifecycle::{Activity, ShutdownSubcription}; use crate::logging::{log_disposition, LogDisposition, RecordType}; use crate::mx::{MailExchanger, ResolvedAddress}; use crate::queue::{Queue, QueueManager}; +use crate::runtime::{rt_spawn, spawn}; use crate::spool::SpoolManager; use anyhow::Context; use config::load_config; @@ -203,11 +204,10 @@ impl EgressPathManager { let mut manager = Self::get().await; let activity = Activity::get()?; let handle = manager.paths.entry(name.clone()).or_insert_with(|| { - tokio::task::Builder::new() - .name(&format!("maintain {name}")) - .spawn({ - let name = name.clone(); - async move { + rt_spawn(format!("maintain {name}"), { + let name = name.clone(); + move || { + Ok(async move { let mut shutdown = ShutdownSubcription::get(); let mut interval = Duration::from_secs(60); loop { @@ -234,9 +234,10 @@ impl EgressPathManager { } } } - } - }) - .expect("failed to spawn maintainer"); + }) + } + }) + .expect("failed to spawn maintainer"); let service = format!("smtp_client:{name}"); let metrics = DeliveryMetrics { @@ -367,16 +368,14 @@ impl EgressPath { let msgs: Vec = self.ready.lock().unwrap().drain(..).collect(); if !msgs.is_empty() { let activity = self.activity.clone(); - tokio::task::Builder::new() - .name(&format!("saving messages for {}", self.name)) - .spawn(async move { - for msg in msgs { - Queue::save_if_needed_and_log(&msg).await; - drop(msg); - } - drop(activity); - }) - .expect("failed to spawn save_if_needed_and_log"); + spawn(format!("saving messages for {}", self.name), async move { + for msg in msgs { + Queue::save_if_needed_and_log(&msg).await; + drop(msg); + } + drop(activity); + }) + .expect("failed to spawn save_if_needed_and_log"); } return; @@ -396,31 +395,28 @@ impl EgressPath { let egress_source = self.egress_source.clone(); let egress_pool = self.egress_pool.clone(); let consecutive_connection_failures = self.consecutive_connection_failures.clone(); - self.connections.push( - tokio::task::Builder::new() - .name(&format!("smtp client {name}")) - .spawn(async move { - if let Err(err) = Dispatcher::run( - &name, - mx, - ready, - notify, - path_config, - metrics, - consecutive_connection_failures.clone(), - egress_source, - egress_pool, - ) - .await - { - tracing::debug!( - "Error in dispatch_queue for {name}: {err:#} \ + if let Ok(handle) = spawn(format!("smtp client {name}"), async move { + if let Err(err) = Dispatcher::run( + &name, + mx, + ready, + notify, + path_config, + metrics, + consecutive_connection_failures.clone(), + egress_source, + egress_pool, + ) + .await + { + tracing::debug!( + "Error in dispatch_queue for {name}: {err:#} \ (consecutive_connection_failures={consecutive_connection_failures:?})" - ); - } - }) - .expect("failed to spawn connection"), - ); + ); + } + }) { + self.connections.push(handle); + } } } @@ -596,20 +592,17 @@ impl Dispatcher { ); chrono::Duration::seconds(60) }); - crate::runtime::Runtime::run(move || { - tokio::task::Builder::new() - .name("requeue for throttle") - .spawn_local(async move { - for msg in msgs { - if let Err(err) = Self::requeue_message(msg, false, Some(delay)).await { - tracing::error!("error requeuing message: {err:#}"); - } + 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:#}"); } - drop(activity); - }) - .expect("failed to requeue for throttle"); + } + drop(activity); + }) }) - .ok(); + .expect("failed to spawn requeue"); } } @@ -623,13 +616,10 @@ impl Dispatcher { let name = self.name.clone(); let egress_pool = self.egress_pool.clone(); let egress_source = self.egress_source.name.clone(); - crate::runtime::Runtime::run(move || { - tokio::task::Builder::new() - .name(&format!( - "bulk queue op for {} msgs {name} {response:?}", - msgs.len(), - )) - .spawn_local(async move { + 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 { @@ -655,19 +645,17 @@ impl Dispatcher { tracing::error!("error requeuing message: {err:#}"); } } else if response.is_permanent() { - tokio::task::Builder::new() - .name("remove msg from spool") - .spawn(async move { - SpoolManager::remove_from_spool(*msg.id()).await - }) - .ok(); + spawn("remove msg from spool", async move { + SpoolManager::remove_from_spool(*msg.id()).await + }) + .ok(); } } drop(activity); }) - .expect("bulk queue spawned"); - }) - .ok(); + }, + ) + .expect("bulk queue spawned"); } } @@ -948,13 +936,8 @@ impl Dispatcher { relay_disposition: None, }) .await; - crate::runtime::Runtime::run(move || { - tokio::task::Builder::new() - .name("requeue message") - .spawn_local( - async move { Self::requeue_message(msg, true, None).await }, - ) - .expect("failed to requeue message"); + rt_spawn("requeue message".to_string(), move || { + Ok(async move { Self::requeue_message(msg, true, None).await }) })?; } self.metrics.msgs_transfail.inc(); @@ -980,9 +963,9 @@ impl Dispatcher { relay_disposition: None, }) .await; - tokio::task::Builder::new() - .name("remove from spool") - .spawn(async move { SpoolManager::remove_from_spool(*msg.id()).await })?; + spawn("remove from spool", async move { + SpoolManager::remove_from_spool(*msg.id()).await + })?; } } Err(err) => { @@ -1008,9 +991,9 @@ impl Dispatcher { relay_disposition: None, }) .await; - tokio::task::Builder::new() - .name("remove from spool") - .spawn(async move { SpoolManager::remove_from_spool(*msg.id()).await })?; + spawn("remove from spool", async move { + SpoolManager::remove_from_spool(*msg.id()).await + })?; } self.metrics.msgs_delivered.inc(); self.metrics.global_msgs_delivered.inc(); @@ -1028,18 +1011,14 @@ 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(); - crate::runtime::Runtime::run(move || { - tokio::task::Builder::new() - .name("Dispatcher::drop") - .spawn_local(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:#}"); - } - }) - .expect("failed to make messages safe in Dispatcher::drop"); + rt_spawn("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(); } diff --git a/crates/kumod/src/http_server/auth.rs b/crates/kumod/src/http_server/auth.rs index 2f141360..dac9ece2 100644 --- a/crates/kumod/src/http_server/auth.rs +++ b/crates/kumod/src/http_server/auth.rs @@ -1,4 +1,5 @@ use crate::http_server::AppState; +use crate::runtime::rt_spawn; use axum::async_trait; use axum::extract::{FromRequestParts, State}; use axum::http::{Request, StatusCode}; @@ -62,11 +63,8 @@ impl AuthKind { pub async fn validate(&self) -> anyhow::Result { let kind = self.clone(); let (tx, rx) = tokio::sync::oneshot::channel(); - crate::runtime::Runtime::run(move || { - tokio::task::Builder::new() - .name(&format!("http auth validate {kind:?}")) - .spawn_local(async move { tx.send(kind.validate_impl().await) }) - .expect("spawned auth validator"); + rt_spawn(format!("http auth validate {kind:?}"), move || { + Ok(async move { tx.send(kind.validate_impl().await) }) })?; rx.await? } diff --git a/crates/kumod/src/http_server/inject_v1.rs b/crates/kumod/src/http_server/inject_v1.rs index 5da4d4d4..5ddd0812 100644 --- a/crates/kumod/src/http_server/inject_v1.rs +++ b/crates/kumod/src/http_server/inject_v1.rs @@ -3,6 +3,7 @@ use crate::http_server::AppError; use crate::logging::{log_disposition, LogDisposition, RecordType}; use crate::mx::ResolvedAddress; use crate::queue::QueueManager; +use crate::runtime::rt_spawn; use anyhow::Context; use axum::extract::Json; use axum_client_ip::InsecureClientIp; @@ -433,9 +434,9 @@ async fn process_recipient<'a>( relay_disposition: None, }) .await; - tokio::task::Builder::new() - .name(&format!("http inject for {peer_address:?}")) - .spawn_local(async move { QueueManager::insert(&queue_name, message).await })?; + rt_spawn(format!("http inject for {peer_address:?}"), move || { + Ok(async move { QueueManager::insert(&queue_name, message).await }) + })?; } Ok(()) @@ -494,13 +495,8 @@ pub async fn inject_v1( let (tx, rx) = tokio::sync::oneshot::channel(); // Bounce to the thread pool where we can run async lua - crate::runtime::Runtime::run(move || { - tokio::task::Builder::new() - .name(&format!("http inject_v1 for {peer_address:?}")) - .spawn_local(async move { - tx.send(inject_v1_impl(auth, sender, peer_address, request).await) - }) - .expect("spawned injection task"); + rt_spawn(format!("http inject_v1 for {peer_address:?}"), move || { + Ok(async move { tx.send(inject_v1_impl(auth, sender, peer_address, request).await) }) })?; rx.await? } diff --git a/crates/kumod/src/http_server/mod.rs b/crates/kumod/src/http_server/mod.rs index 42cc6521..6956803b 100644 --- a/crates/kumod/src/http_server/mod.rs +++ b/crates/kumod/src/http_server/mod.rs @@ -1,3 +1,4 @@ +use crate::runtime::spawn; use anyhow::Context; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; @@ -104,23 +105,19 @@ impl HttpListenerParams { let config = self.tls_config().await?; tracing::debug!("https listener on {addr:?}"); let server = axum_server::from_tcp_rustls(socket, config); - tokio::task::Builder::new() - .name(&format!("https {addr:?}")) - .spawn(async move { - server - .serve(app.into_make_service_with_connect_info::()) - .await - })?; + spawn(format!("https {addr:?}"), async move { + server + .serve(app.into_make_service_with_connect_info::()) + .await + })?; } else { tracing::debug!("http listener on {addr:?}"); let server = axum_server::from_tcp(socket); - tokio::task::Builder::new() - .name(&format!("http {addr:?}")) - .spawn(async move { - server - .serve(app.into_make_service_with_connect_info::()) - .await - })?; + spawn(format!("http {addr:?}"), async move { + server + .serve(app.into_make_service_with_connect_info::()) + .await + })?; } Ok(()) } diff --git a/crates/kumod/src/main.rs b/crates/kumod/src/main.rs index 5e709342..ad08ae77 100644 --- a/crates/kumod/src/main.rs +++ b/crates/kumod/src/main.rs @@ -1,4 +1,5 @@ use crate::lifecycle::LifeCycle; +use crate::runtime::rt_spawn; use anyhow::Context; use caps::{CapSet, Capability, CapsHashSet}; use clap::{Parser, ValueEnum}; @@ -170,17 +171,14 @@ async fn run(opts: Opt) -> anyhow::Result<()> { let mut life_cycle = LifeCycle::new(); - crate::runtime::Runtime::run(move || { - tokio::task::Builder::new() - .name("initialize") - .spawn_local(async move { - if let Err(err) = perform_init().await { - tracing::error!("problem initializing: {err:#}"); - LifeCycle::request_shutdown().await; - } - Ok::<(), anyhow::Error>(()) - }) - .expect("initialize spawned"); + rt_spawn("initialize".to_string(), move || { + Ok(async move { + if let Err(err) = perform_init().await { + tracing::error!("problem initializing: {err:#}"); + LifeCycle::request_shutdown().await; + } + Ok::<(), anyhow::Error>(()) + }) })?; life_cycle.wait_for_shutdown().await; diff --git a/crates/kumod/src/mod_kumo.rs b/crates/kumod/src/mod_kumo.rs index 09381b63..7a387817 100644 --- a/crates/kumod/src/mod_kumo.rs +++ b/crates/kumod/src/mod_kumo.rs @@ -4,6 +4,7 @@ use crate::http_server::HttpListenerParams; use crate::lifecycle::LifeCycle; use crate::logging::{ClassifierParams, LogFileParams}; use crate::queue::QueueConfig; +use crate::runtime::spawn; use crate::smtp_server::{EsmtpListenerParams, RejectError}; use config::{any_err, get_or_create_module}; use mlua::{Function, Lua, LuaSerdeExt, Value}; @@ -52,14 +53,12 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { "start_esmtp_listener", lua.create_async_function(|lua, params: Value| async move { let params: EsmtpListenerParams = lua.from_value(params)?; - tokio::task::Builder::new() - .name("start_esmtp_listener") - .spawn(async move { - if let Err(err) = params.run().await { - tracing::error!("Error in SmtpServer: {err:#}"); - } - }) - .map_err(any_err)?; + spawn("start_esmtp_listener", async move { + if let Err(err) = params.run().await { + tracing::error!("Error in SmtpServer: {err:#}"); + } + }) + .map_err(any_err)?; Ok(()) })?, )?; @@ -68,17 +67,15 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { "define_spool", lua.create_async_function(|lua, params: Value| async move { let params = lua.from_value(params)?; - tokio::task::Builder::new() - .name("define_spool") - .spawn(async move { - if let Err(err) = define_spool(params).await { - tracing::error!("Error in spool: {err:#}"); - LifeCycle::request_shutdown().await; - } - }) - .map_err(any_err)? - .await - .map_err(any_err) + spawn("define_spool", async move { + if let Err(err) = define_spool(params).await { + tracing::error!("Error in spool: {err:#}"); + LifeCycle::request_shutdown().await; + } + }) + .map_err(any_err)? + .await + .map_err(any_err) })?, )?; diff --git a/crates/kumod/src/queue.rs b/crates/kumod/src/queue.rs index 664f5ada..d8469032 100644 --- a/crates/kumod/src/queue.rs +++ b/crates/kumod/src/queue.rs @@ -351,18 +351,15 @@ impl Queue { }))); let queue_clone = handle.clone(); - crate::runtime::Runtime::run(move || { - tokio::task::Builder::new() - .name(&format!("maintain {name}")) - .spawn_local(async move { - if let Err(err) = maintain_named_queue(&queue_clone).await { - tracing::error!( - "maintain_named_queue {}: {err:#}", - queue_clone.lock().await.name - ); - } - }) - .expect("failed to spawn queue maintainer"); + crate::runtime::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.lock().await.name + ); + } + }) })?; Ok(handle) diff --git a/crates/kumod/src/runtime.rs b/crates/kumod/src/runtime.rs index 8be6e99e..62131756 100644 --- a/crates/kumod/src/runtime.rs +++ b/crates/kumod/src/runtime.rs @@ -11,8 +11,9 @@ //! is to use tokio::task::spawn_local to spawn the local future. //! For example, when accepting new connections, we use this to //! spawn the server processing future. -use async_channel::{unbounded, Sender}; -use tokio::task::LocalSet; +use async_channel::{bounded, unbounded, Sender}; +use std::future::Future; +use tokio::task::{JoinHandle, LocalSet}; lazy_static::lazy_static! { static ref RUNTIME: Runtime = Runtime::new().unwrap(); @@ -56,15 +57,67 @@ impl Runtime { Ok(Self { jobs: tx }) } +} - /// Schedule a future on the pool. - /// Will return once the future is scheduled. - /// Does not wait for the future to complete. - pub fn run(func: F) -> anyhow::Result<()> { - RUNTIME - .jobs - .send_blocking(Command::Run(Box::new(func))) - .map_err(|err| anyhow::anyhow!("failed to send func to runtime thread: {err:#}"))?; - Ok(()) - } +/// 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 fn rt_spawn anyhow::Result) + Send + 'static, FUT>( + name: String, + func: F, +) -> anyhow::Result> +where + FUT: Future + 'static, + FUT::Output: Send, +{ + let (tx, rx) = bounded::>>(1); + RUNTIME + .jobs + .send_blocking(Command::Run(Box::new(move || match (func)() { + Ok(future) => { + tx.send_blocking( + tokio::task::Builder::new() + .name(&name) + .spawn_local(future) + .map_err(|e| e.into()), + ) + .ok(); + } + Err(err) => { + tx.send_blocking(Err(err)).ok(); + } + }))) + .map_err(|err| anyhow::anyhow!("failed to send func to runtime thread: {err:#}"))?; + rx.recv_blocking()? +} + +/// Spawn a future as a task with a name. +pub fn spawn>(name: N, fut: FUT) -> std::io::Result> +where + FUT: Future + Send + 'static, + FUT::Output: Send, +{ + tokio::task::Builder::new().name(name.as_ref()).spawn(fut) +} + +/// Spawn a local future as a task with a name. +pub fn spawn_local>( + name: N, + fut: FUT, +) -> std::io::Result> +where + FUT: Future + 'static, + FUT::Output: Send, +{ + tokio::task::Builder::new() + .name(name.as_ref()) + .spawn_local(fut) } diff --git a/crates/kumod/src/smtp_server.rs b/crates/kumod/src/smtp_server.rs index 1046879f..a5570463 100644 --- a/crates/kumod/src/smtp_server.rs +++ b/crates/kumod/src/smtp_server.rs @@ -3,6 +3,7 @@ use crate::lifecycle::{Activity, ShutdownSubcription}; use crate::logging::{log_disposition, LogDisposition, RecordType}; use crate::mx::ResolvedAddress; use crate::queue::QueueManager; +use crate::runtime::{rt_spawn, spawn_local}; use crate::spool::SpoolManager; use anyhow::{anyhow, Context}; use chrono::Utc; @@ -211,17 +212,15 @@ impl EsmtpListenerParams { let (socket, peer_address) = result?; let my_address = socket.local_addr()?; let params = self.clone(); - crate::runtime::Runtime::run(move || { - tokio::task::Builder::new() - .name(&format!("SmtpServer {peer_address:?}")) - .spawn_local(async move { - if let Err(err) = - SmtpServer::run(socket, my_address, peer_address, params).await + rt_spawn( + format!("SmtpServer {peer_address:?}"), + move || Ok(async move { + if let Err(err) = + SmtpServer::run(socket, my_address, peer_address, params).await { tracing::error!("SmtpServer::run: {err:#}"); - } - }).ok(); - })?; + } + }))?; } }; } @@ -884,18 +883,19 @@ impl SmtpServer { } if !messages.is_empty() { - tokio::task::Builder::new() - .name(&format!( + spawn_local( + format!( "SmtpServer: insert {} msgs for {:?}", messages.len(), self.peer_address - )) - .spawn_local(async move { + ), + async move { for (queue_name, msg) in messages { QueueManager::insert(&queue_name, msg).await?; } Ok::<(), anyhow::Error>(()) - })?; + }, + )?; } let ids = ids.join(" "); diff --git a/crates/kumod/src/spool.rs b/crates/kumod/src/spool.rs index 083a5526..b47bda17 100644 --- a/crates/kumod/src/spool.rs +++ b/crates/kumod/src/spool.rs @@ -157,35 +157,31 @@ impl SpoolManager { tracing::debug!("starting maintainer for spool {name} is_meta={is_meta}"); - let maintainer = tokio::task::Builder::new() - .name(&format!("maintain spool {name}")) - .spawn({ - let name = name.clone(); - let spool = spool.clone(); - let tx = if is_meta { Some(tx.clone()) } else { None }; - { - async move { - // start enumeration - if let Some(tx) = tx { - if let Err(err) = spool.enumerate(tx) { - tracing::error!( - "error during spool enumeration for {name}: {err:#}" - ); - } + let maintainer = crate::runtime::spawn(format!("maintain spool {name}"), { + let name = name.clone(); + let spool = spool.clone(); + let tx = if is_meta { Some(tx.clone()) } else { None }; + { + async move { + // start enumeration + if let Some(tx) = tx { + if let Err(err) = spool.enumerate(tx) { + tracing::error!( + "error during spool enumeration for {name}: {err:#}" + ); } + } - // And maintain it every 10 minutes - loop { - tokio::time::sleep(Duration::from_secs(10 * 60)).await; - if let Err(err) = spool.cleanup().await { - tracing::error!( - "error doing spool cleanup for {name}: {err:#}" - ); - } + // And maintain it every 10 minutes + loop { + tokio::time::sleep(Duration::from_secs(10 * 60)).await; + if let Err(err) = spool.cleanup().await { + tracing::error!("error doing spool cleanup for {name}: {err:#}"); } } } - })?; + } + })?; spool.0.maintainer.lock().unwrap().replace(maintainer); }