ready_queue: micro optimize insertion code path

We can save repeated string construction by caching the mapping
from source -> ready queue name, so that's what we do.

While we're in there, swap from tokio::Mutex around the manager
to StdMutex which is a bit cheaper.

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<Instant>,
 }
 
 pub fn fully_qualify(domain_name: &str) -> ResolveResult<Name> {
@@ -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,
 }
 "#
         );

crates/kumo-server-common/src/config_handle.rs

@@ -40,6 +40,10 @@ impl<T: Clone + Send> ConfigHandle<T> {
         }
     }
 
+    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 {

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<Queue>;
 
+struct CachedReadyQueueName {
+    name: ReadyQueueName,
+    /// queue_config.generation()
+    generation: usize,
+}
+
 pub struct Queue {
     name: String,
     queue: StdMutex<TimeQ<Message>>,
@@ -382,6 +388,7 @@ pub struct Queue {
     delayed_gauge: IntGauge,
     activity: Activity,
     rr: EgressPoolRoundRobin,
+    ready_queue_names: StdMutex<HashMap<String, Arc<CachedReadyQueueName>>>,
 }
 
 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<Arc<CachedReadyQueueName>> {
+        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<Arc<CachedReadyQueueName>> {
+        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

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<ReadyQueueManager> = Mutex::new(ReadyQueueManager::new());
+    static ref MANAGER: StdMutex<ReadyQueueManager> = 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<Arc<MailExchanger>>,
 }
 
+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<ReadyQueueHandle> {
+        let manager = MANAGER.lock().unwrap();
+        manager.queues.get(name).cloned()
+    }
+
+    pub fn get_by_ready_queue_name(name: &ReadyQueueName) -> Option<ReadyQueueHandle> {
+        Self::get_by_name(&name.name)
+    }
+
     pub async fn resolve_by_queue_name(
         queue_name: &str,
         queue_config: &ConfigHandle<QueueConfig>,
         egress_source: &str,
         egress_pool: &str,
     ) -> anyhow::Result<ReadyQueueHandle> {
-        // 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();
                                 }