diff --git a/crates/dns-resolver/src/lib.rs b/crates/dns-resolver/src/lib.rs index 5c07ad0e..2905e1e0 100644 --- a/crates/dns-resolver/src/lib.rs +++ b/crates/dns-resolver/src/lib.rs @@ -49,6 +49,8 @@ pub struct MailExchanger { /// DNSSEC verified pub is_secure: bool, pub is_mx: bool, + #[serde(skip)] + expires: Option, } pub fn fully_qualify(domain_name: &str) -> ResolveResult { @@ -198,6 +200,7 @@ impl MailExchanger { is_domain_literal: true, is_secure: false, is_mx: false, + expires: None, })); } Err(err) => { @@ -221,6 +224,7 @@ impl MailExchanger { is_domain_literal: true, is_secure: false, is_mx: false, + expires: None, })); } Err(err) => { @@ -263,6 +267,7 @@ impl MailExchanger { is_domain_literal: false, is_secure, is_mx, + expires: Some(expires), }; let mx = Arc::new(mx); @@ -273,6 +278,13 @@ impl MailExchanger { Ok(mx) } + pub fn has_expired(&self) -> bool { + match self.expires { + Some(deadline) => deadline <= Instant::now(), + None => false, + } + } + pub async fn resolve_addresses(&self) -> ResolvedMxAddresses { let mut result = vec![]; @@ -562,6 +574,7 @@ MailExchanger { is_domain_literal: true, is_secure: false, is_mx: false, + expires: None, } "# ); @@ -597,6 +610,7 @@ MailExchanger { is_domain_literal: true, is_secure: false, is_mx: false, + expires: None, } "# ); @@ -632,6 +646,7 @@ MailExchanger { is_domain_literal: true, is_secure: false, is_mx: false, + expires: None, } "# ); diff --git a/crates/kumo-server-common/src/config_handle.rs b/crates/kumo-server-common/src/config_handle.rs index 303749be..679b2cb4 100644 --- a/crates/kumo-server-common/src/config_handle.rs +++ b/crates/kumo-server-common/src/config_handle.rs @@ -40,6 +40,10 @@ impl ConfigHandle { } } + pub fn generation(&self) -> usize { + self.inner.generation.load(Ordering::SeqCst) + } + /// Updates the upstream, shared value. /// Other config handles will notice the change when borrow() is subsequently called. pub fn update(&self, new_value: T) -> usize { diff --git a/crates/kumod/src/queue.rs b/crates/kumod/src/queue.rs index 5f316210..c5d704d9 100644 --- a/crates/kumod/src/queue.rs +++ b/crates/kumod/src/queue.rs @@ -3,7 +3,7 @@ use crate::http_server::admin_bounce_v1::AdminBounceEntry; use crate::http_server::admin_suspend_v1::AdminSuspendEntry; use crate::logging::{log_disposition, LogDisposition, RecordType}; use crate::lua_deliver::LuaDeliveryProtocol; -use crate::ready_queue::ReadyQueueManager; +use crate::ready_queue::{ReadyQueueManager, ReadyQueueName}; use crate::smtp_dispatcher::SmtpProtocol; use crate::spool::SpoolManager; use anyhow::Context; @@ -374,6 +374,12 @@ struct ReadyQueueFull; type QueueHandle = Arc; +struct CachedReadyQueueName { + name: ReadyQueueName, + /// queue_config.generation() + generation: usize, +} + pub struct Queue { name: String, queue: StdMutex>, @@ -382,6 +388,7 @@ pub struct Queue { delayed_gauge: IntGauge, activity: Activity, rr: EgressPoolRoundRobin, + ready_queue_names: StdMutex>>, } impl Queue { @@ -425,6 +432,7 @@ impl Queue { delayed_gauge, activity, rr, + ready_queue_names: StdMutex::new(HashMap::new()), }); let queue_clone = handle.clone(); @@ -684,6 +692,41 @@ impl Queue { Ok(()) } + fn get_ready_queue_for_source(&self, source: &str) -> Option> { + let mut ready_queue_names = self.ready_queue_names.lock().unwrap(); + let name = ready_queue_names.get(source)?; + + if self.queue_config.generation() != name.generation || name.name.has_expired() { + ready_queue_names.remove(source); + return None; + } + + Some(Arc::clone(name)) + } + + async fn compute_ready_queue_name( + &self, + source: &str, + ) -> anyhow::Result> { + if let Some(entry) = self.get_ready_queue_for_source(source) { + return Ok(entry); + } + + let generation = self.queue_config.generation(); + + let name = + ReadyQueueManager::compute_queue_name(&self.name, &self.queue_config, source).await?; + + let cached = Arc::new(CachedReadyQueueName { name, generation }); + + self.ready_queue_names + .lock() + .unwrap() + .insert(source.to_string(), cached.clone()); + + Ok(cached) + } + #[instrument(skip(self, msg))] async fn insert_ready_impl(&self, msg: Message) -> anyhow::Result<()> { tracing::trace!("insert_ready {}", msg.id()); @@ -751,6 +794,13 @@ impl Queue { } }; + // Hot path: use cached source -> ready queue mapping + let ready_name = self.compute_ready_queue_name(&egress_source).await?; + if let Some(site) = ReadyQueueManager::get_by_ready_queue_name(&ready_name.name) { + return site.insert(msg).map_err(|_| ReadyQueueFull.into()); + } + + // Miss: compute and establish a new queue match ReadyQueueManager::resolve_by_queue_name( &self.name, &self.queue_config, @@ -1035,8 +1085,7 @@ async fn maintain_named_queue(q: &QueueHandle) -> anyhow::Result<()> { drop(msg); } - let queue_mgr = ReadyQueueManager::get().await; - if queue_mgr.number_of_queues() == 0 { + if ReadyQueueManager::number_of_queues() == 0 { tracing::debug!( "{}: there are no more queues and the delayed queue is empty, reaping", q.name diff --git a/crates/kumod/src/ready_queue.rs b/crates/kumod/src/ready_queue.rs index ac3ba0f7..3e453308 100644 --- a/crates/kumod/src/ready_queue.rs +++ b/crates/kumod/src/ready_queue.rs @@ -25,12 +25,12 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex as StdMutex}; use std::time::{Duration, Instant}; use throttle::limit::{LimitLease, LimitSpec}; -use tokio::sync::{Mutex, MutexGuard, Notify}; +use tokio::sync::Notify; use tokio::task::JoinHandle; use tracing::instrument; // TODO move to here lazy_static::lazy_static! { - static ref MANAGER: Mutex = Mutex::new(ReadyQueueManager::new()); + static ref MANAGER: StdMutex = StdMutex::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(); @@ -42,6 +42,15 @@ pub struct ReadyQueueName { pub mx: Option>, } +impl ReadyQueueName { + pub fn has_expired(&self) -> bool { + match &self.mx { + Some(mx) => mx.has_expired(), + None => false, + } + } +} + pub struct ReadyQueueConfig { pub name: String, pub site_name: String, @@ -60,12 +69,8 @@ impl ReadyQueueManager { Self::default() } - pub fn number_of_queues(&self) -> usize { - self.queues.len() - } - - pub async fn get() -> MutexGuard<'static, Self> { - MANAGER.lock().await + pub fn number_of_queues() -> usize { + MANAGER.lock().unwrap().queues.len() } pub async fn compute_queue_name( @@ -179,29 +184,21 @@ impl ReadyQueueManager { }) } + pub fn get_by_name(name: &str) -> Option { + let manager = MANAGER.lock().unwrap(); + manager.queues.get(name).cloned() + } + + pub fn get_by_ready_queue_name(name: &ReadyQueueName) -> Option { + Self::get_by_name(&name.name) + } + pub async fn resolve_by_queue_name( queue_name: &str, queue_config: &ConfigHandle, egress_source: &str, egress_pool: &str, ) -> anyhow::Result { - // Assumption: that the queue likely already exists. - // We do a smaller amount of work for that case, even - // though the creation case below may need to repeat - // some of it. - let ReadyQueueName { - name, - site_name: _, - mx: _, - } = Self::compute_queue_name(queue_name, queue_config, egress_source).await?; - - { - let manager = MANAGER.lock().await; - if let Some(handle) = manager.queues.get(&name) { - return Ok(handle.clone()); - } - } - let ReadyQueueConfig { name, site_name: _, @@ -210,7 +207,7 @@ impl ReadyQueueManager { mx, } = Self::compute_config(queue_name, queue_config, egress_source).await?; - let mut manager = Self::get().await; + let mut manager = MANAGER.lock().unwrap(); let activity = Activity::get(format!("ReadyQueueHandle {name}"))?; let handle = manager.queues.entry(name.clone()).or_insert_with(|| { @@ -257,12 +254,9 @@ impl ReadyQueueManager { let mut reap_deadline = None; let mut done_abort = false; - let queue = { - let mgr = Self::get().await; - mgr.queues.get(&name).cloned().ok_or_else(|| { - anyhow::anyhow!("ready_queue {name} not found when starting up maintainer_task") - })? - }; + let queue = ReadyQueueManager::get_by_name(&name).ok_or_else(|| { + anyhow::anyhow!("ready_queue {name} not found when starting up maintainer_task") + })?; loop { tokio::select! { @@ -307,11 +301,13 @@ impl ReadyQueueManager { queue.maintain().await; if queue.reapable(&last_change) { - let mut mgr = Self::get().await; + let mut mgr = MANAGER.lock().unwrap(); if queue.reapable(&last_change) { tracing::debug!("reaping site {name}"); - queue.reinsert_ready_queue("reap").await; mgr.queues.remove(&name); + drop(mgr); + + queue.reinsert_ready_queue("reap").await; crate::metrics_helper::remove_metrics_for_service(&format!( "smtp_client:{name}" )); @@ -664,11 +660,7 @@ impl Drop for Dispatcher { tracing::error!("error requeuing message: {err:#}"); } else { tokio::time::sleep(Duration::from_secs(1)).await; - let ready_queue = { - let mgr = ReadyQueueManager::get().await; - mgr.queues.get(&name).cloned() - }; - + let ready_queue = ReadyQueueManager::get_by_name(&name); if let Some(q) = ready_queue { q.notify_maintainer.notify_one(); }