mirror/kumomta
1
0
Fork 0
mirror of https://github.com/mailscope/kumomta.git synced 2026-10-29 01:36:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: make task spawning less boilerplatey

Browse Source
This commit is contained in:
Wez Furlong
2023-03-06 11:01:31 -07:00
parent 43155c809b
commit 265698e4c3
10 changed files with 227 additions and 216 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/kumod/src/egress_path.rs
+74 -95
View File
@@ -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<Message> = 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();
}
crates/kumod/src/http_server/auth.rs
+3 -5
View File
@@ -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<bool> {
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?
}
crates/kumod/src/http_server/inject_v1.rs
+6 -10
View File
@@ -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?
}
crates/kumod/src/http_server/mod.rs
+11 -14
View File
@@ -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::<SocketAddr>())
.await
})?;
spawn(format!("https {addr:?}"), async move {
server
.serve(app.into_make_service_with_connect_info::<SocketAddr>())
.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::<SocketAddr>())
.await
})?;
spawn(format!("http {addr:?}"), async move {
server
.serve(app.into_make_service_with_connect_info::<SocketAddr>())
.await
})?;
}
Ok(())
}
crates/kumod/src/main.rs
+9 -11
View File
@@ -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;
crates/kumod/src/mod_kumo.rs
+16 -19
View File
@@ -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)
})?,
)?;
crates/kumod/src/queue.rs
+9 -12
View File
@@ -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)
crates/kumod/src/runtime.rs
+65 -12
View File
@@ -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<F: FnOnce() + Send + 'static>(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<F: (FnOnce() -> anyhow::Result<FUT>) + Send + 'static, FUT>(
name: String,
func: F,
) -> anyhow::Result<JoinHandle<FUT::Output>>
where
FUT: Future + 'static,
FUT::Output: Send,
{
let (tx, rx) = bounded::<anyhow::Result<JoinHandle<FUT::Output>>>(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<FUT, N: AsRef<str>>(name: N, fut: FUT) -> std::io::Result<JoinHandle<FUT::Output>>
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<FUT, N: AsRef<str>>(
name: N,
fut: FUT,
) -> std::io::Result<JoinHandle<FUT::Output>>
where
FUT: Future + 'static,
FUT::Output: Send,
{
tokio::task::Builder::new()
.name(name.as_ref())
.spawn_local(fut)
}
crates/kumod/src/smtp_server.rs
+14 -14
View File
@@ -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(" ");
crates/kumod/src/spool.rs
+20 -24
View File
@@ -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);
}