add dedicated localsets for queue and ready_queue

This commit is contained in:
Wez Furlong
2024-05-01 22:36:25 -07:00
parent c7bf0817f6
commit 79f0701dd3
5 changed files with 224 additions and 179 deletions
+67 -35
View File
@@ -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<Self> {
pub fn new(name_prefix: &'static str) -> anyhow::Result<Self> {
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<F: (FnOnce() -> anyhow::Result<FUT>) + Send + 'static, FUT>(
&self,
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);
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<F: (FnOnce() -> anyhow::Result<FUT>) + 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::<anyhow::Result<JoinHandle<FUT::Output>>>(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<F: (FnOnce() -> anyhow::Result<FUT>) + Send + 'static, FUT>(
@@ -110,19 +154,7 @@ pub fn rt_spawn_non_blocking<F: (FnOnce() -> anyhow::Result<FUT>) + 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.
+5 -5
View File
@@ -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 {
+15 -11
View File
@@ -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:#}");
}
+124 -115
View File
@@ -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<ReadyQueueManager> = 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");
}
}
+13 -13
View File
@@ -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