diff --git a/crates/kumod/src/queue/maintainer.rs b/crates/kumod/src/queue/maintainer.rs index b150427c..70ab575a 100644 --- a/crates/kumod/src/queue/maintainer.rs +++ b/crates/kumod/src/queue/maintainer.rs @@ -2,14 +2,20 @@ use crate::http_server::admin_bounce_v1::AdminBounceEntry; use crate::queue::insert_context::InsertReason; use crate::queue::manager::{QueueManager, MANAGER}; use crate::queue::queue::{Queue, QueueHandle}; +use crate::queue::strategy::{QueueStructure, SINGLETON_WHEEL, SINGLETON_WHEEL_V2}; use crate::queue::wait_for_message_batch; use crate::ready_queue::ReadyQueueManager; use kumo_server_lifecycle::{Activity, ShutdownSubcription}; use kumo_server_runtime::Runtime; +use message::message::{MessageList, WeakMessage}; +use message::Message; +use parking_lot::FairMutex; use prometheus::IntCounter; +use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::LazyLock; use std::time::{Duration, Instant}; +use timeq::PopResult; use tracing::instrument; const ONE_SECOND: Duration = Duration::from_secs(1); @@ -152,3 +158,198 @@ pub async fn maintain_named_queue(q: &QueueHandle) -> anyhow::Result<()> { } } } + +pub async fn run_singleton_wheel_v2() -> anyhow::Result<()> { + let mut shutdown = ShutdownSubcription::get(); + + tracing::debug!("singleton_wheel_v2: starting up"); + + async fn reinsert_ready( + msg: Message, + to_shrink: &mut HashMap, + ) -> anyhow::Result<()> { + if !msg.is_meta_loaded() { + msg.load_meta().await?; + } + let queue_name = msg.get_queue_name()?; + // Use get_opt rather than resolve here. If the queue is not currently + // tracked in the QueueManager then this message cannot possibly belong + // to it. Using resolve would have the side effect of creating an empty + // queue for it, which will then age out later. It's a waste to do that, + // so we just check and skip. + let queue = QueueManager::get_opt(&queue_name) + .ok_or_else(|| anyhow::anyhow!("no scheduled queue"))?; + + if let Some(b) = AdminBounceEntry::get_for_queue_name(&queue.name) { + // Note that this will cause the msg to be removed from the + // queue so the remove() check below will return false + queue.bounce_all(&b).await; + } + + // Verify that the message is still in the queue + match &queue.queue { + QueueStructure::SingletonTimerWheelV2(q) => { + fn remove(q: &FairMutex>, msg: &Message) -> bool { + q.lock().remove(msg) + } + + if remove(q, &msg) { + queue.metrics().sub(1); + queue + .insert_ready(msg, InsertReason::DueTimeWasReached.into(), None) + .await?; + if !to_shrink.contains_key(&queue_name) { + to_shrink.insert(queue_name, queue); + } + } + } + _ => { + anyhow::bail!("impossible queue strategy"); + } + } + + Ok(()) + } + + let mut to_shrink = HashMap::new(); + + loop { + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(3)) => { + TOTAL_QMAINT_RUNS.inc(); + + fn pop() -> (MessageList, usize) { + let mut wheel = SINGLETON_WHEEL_V2.lock(); + + let ready = wheel.pop(); + if !ready.is_empty() { + tracing::debug!("singleton_wheel_v2: popped {} messages", ready.len()); + } + + (ready, wheel.len()) + } + + let mut reinserted = 0; + let (msgs, len) = pop(); + for msg in msgs { + reinserted += 1; + if let Err(err) = reinsert_ready(msg, &mut to_shrink).await { + tracing::error!("singleton_wheel_v2: reinsert_ready: {err:#}"); + } + } + tracing::debug!("singleton_wheel_v2: done reinserting {reinserted}. total scheduled={len}"); + + for (_queue_name, queue) in to_shrink.drain() { + queue.queue.shrink(); + } + to_shrink.shrink_to_fit(); + } + _ = shutdown.shutting_down() => { + tracing::info!("singleton_wheel: stopping"); + return Ok(()); + } + } + } +} + +pub async fn run_singleton_wheel_v1() -> anyhow::Result<()> { + let mut shutdown = ShutdownSubcription::get(); + + tracing::debug!("singleton_wheel_v1: starting up"); + + async fn reinsert_ready( + msg: Message, + to_shrink: &mut HashMap, + ) -> anyhow::Result<()> { + if !msg.is_meta_loaded() { + msg.load_meta().await?; + } + let queue_name = msg.get_queue_name()?; + // Use get_opt rather than resolve here. If the queue is not currently + // tracked in the QueueManager then this message cannot possibly belong + // to it. Using resolve would have the side effect of creating an empty + // queue for it, which will then age out later. It's a waste to do that, + // so we just check and skip. + let queue = QueueManager::get_opt(&queue_name) + .ok_or_else(|| anyhow::anyhow!("no scheduled queue"))?; + + if let Some(b) = AdminBounceEntry::get_for_queue_name(&queue.name) { + // Note that this will cause the msg to be removed from the + // queue so the remove() check below will return false + queue.bounce_all(&b).await; + } + + // Verify that the message is still in the queue + match &queue.queue { + QueueStructure::SingletonTimerWheel(q) => { + fn remove(q: &FairMutex>, msg: &Message) -> bool { + q.lock().remove(msg) + } + + if remove(q, &msg) { + queue.metrics().sub(1); + queue + .insert_ready(msg, InsertReason::DueTimeWasReached.into(), None) + .await?; + if !to_shrink.contains_key(&queue_name) { + to_shrink.insert(queue_name, queue); + } + } + } + _ => { + anyhow::bail!("impossible queue strategy"); + } + } + + Ok(()) + } + + let mut to_shrink = HashMap::new(); + + loop { + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(3)) => { + TOTAL_QMAINT_RUNS.inc(); + + fn pop() -> (Vec, usize) { + let mut wheel = SINGLETON_WHEEL.lock(); + + let msgs = if let PopResult::Items(weak_messages) = wheel.pop() { + tracing::debug!("singleton_wheel: popped {} messages", weak_messages.len()); + weak_messages + } else { + vec![] + }; + + (msgs, wheel.len()) + } + + let mut reinserted = 0; + let (msgs, len) = pop(); + let mut messages = vec![]; + for weak_message in msgs { + if let Some(msg) = weak_message.upgrade() { + messages.push(msg); + } + } + wait_for_message_batch(&messages).await; + for msg in messages { + reinserted += 1; + if let Err(err) = reinsert_ready(msg, &mut to_shrink).await { + tracing::error!("singleton_wheel: reinsert_ready: {err:#}"); + } + } + tracing::debug!("singleton_wheel: done reinserting {reinserted}. total scheduled={len}"); + + for (_queue_name, queue) in to_shrink.drain() { + queue.queue.shrink(); + } + to_shrink.shrink_to_fit(); + } + _ = shutdown.shutting_down() => { + tracing::info!("singleton_wheel: stopping"); + return Ok(()); + } + } + } +} diff --git a/crates/kumod/src/queue/queue.rs b/crates/kumod/src/queue/queue.rs index 452d0bdd..c3237bc3 100644 --- a/crates/kumod/src/queue/queue.rs +++ b/crates/kumod/src/queue/queue.rs @@ -7,15 +7,11 @@ use crate::logging::disposition::{log_disposition, LogDisposition, RecordType}; use crate::queue::config::QueueConfig; use crate::queue::delivery_proto::DeliveryProto; use crate::queue::insert_context::{InsertContext, InsertReason}; -use crate::queue::maintainer::{maintain_named_queue, QMAINT_RUNTIME, TOTAL_QMAINT_RUNS}; +use crate::queue::maintainer::{maintain_named_queue, QMAINT_RUNTIME}; use crate::queue::manager::{QueueManager, MANAGER, SCHEDULED_QUEUE_COUNT}; use crate::queue::metrics::ScheduledMetrics; -use crate::queue::strategy::{ - QueueInsertResult, QueueStrategy, QueueStructure, SINGLETON_WHEEL, SINGLETON_WHEEL_V2, -}; -use crate::queue::{ - opt_timeout_at, wait_for_message_batch, IncrementAttempts, InsertResult, ReadyQueueFull, -}; +use crate::queue::strategy::{QueueInsertResult, QueueStrategy, QueueStructure}; +use crate::queue::{opt_timeout_at, IncrementAttempts, InsertResult, ReadyQueueFull}; use crate::ready_queue::ReadyQueueManager; use crate::smtp_server::{make_deferred_queue_config, DEFERRED_QUEUE_NAME}; use crate::spool::SpoolManager; @@ -30,18 +26,17 @@ use kumo_server_common::config_handle::ConfigHandle; use kumo_server_lifecycle::{is_shutting_down, Activity, ShutdownSubcription}; use kumo_server_runtime::{get_main_runtime, spawn, spawn_blocking_on}; use kumo_template::{context, TemplateEngine}; -use message::message::{MessageList, WeakMessage}; use message::queue_name::QueueNameComponents; use message::Message; use parking_lot::FairMutex; use prometheus::IntGauge; use rfc5321::{EnhancedStatusCode, Response}; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, LazyLock, OnceLock}; use std::time::{Duration, Instant}; use throttle::ThrottleResult; -use timeq::{PopResult, TimerEntryWithDelay}; +use timeq::TimerEntryWithDelay; use tokio::sync::Notify; use tracing::instrument; @@ -1451,201 +1446,6 @@ impl Queue { *self.last_change.lock() } - pub async fn run_singleton_wheel_v2() -> anyhow::Result<()> { - let mut shutdown = ShutdownSubcription::get(); - - tracing::debug!("singleton_wheel_v2: starting up"); - - async fn reinsert_ready( - msg: Message, - to_shrink: &mut HashMap, - ) -> anyhow::Result<()> { - if !msg.is_meta_loaded() { - msg.load_meta().await?; - } - let queue_name = msg.get_queue_name()?; - // Use get_opt rather than resolve here. If the queue is not currently - // tracked in the QueueManager then this message cannot possibly belong - // to it. Using resolve would have the side effect of creating an empty - // queue for it, which will then age out later. It's a waste to do that, - // so we just check and skip. - let queue = QueueManager::get_opt(&queue_name) - .ok_or_else(|| anyhow::anyhow!("no scheduled queue"))?; - - if let Some(b) = AdminBounceEntry::get_for_queue_name(&queue.name) { - // Note that this will cause the msg to be removed from the - // queue so the remove() check below will return false - queue.bounce_all(&b).await; - } - - // Verify that the message is still in the queue - match &queue.queue { - QueueStructure::SingletonTimerWheelV2(q) => { - fn remove(q: &FairMutex>, msg: &Message) -> bool { - q.lock().remove(msg) - } - - if remove(q, &msg) { - queue.metrics().sub(1); - queue - .insert_ready(msg, InsertReason::DueTimeWasReached.into(), None) - .await?; - if !to_shrink.contains_key(&queue_name) { - to_shrink.insert(queue_name, queue); - } - } - } - _ => { - anyhow::bail!("impossible queue strategy"); - } - } - - Ok(()) - } - - let mut to_shrink = HashMap::new(); - - loop { - tokio::select! { - _ = tokio::time::sleep(Duration::from_secs(3)) => { - TOTAL_QMAINT_RUNS.inc(); - - fn pop() -> (MessageList, usize) { - let mut wheel = SINGLETON_WHEEL_V2.lock(); - - let ready = wheel.pop(); - if !ready.is_empty() { - tracing::debug!("singleton_wheel_v2: popped {} messages", ready.len()); - } - - (ready, wheel.len()) - } - - let mut reinserted = 0; - let (msgs, len) = pop(); - for msg in msgs { - reinserted += 1; - if let Err(err) = reinsert_ready(msg, &mut to_shrink).await { - tracing::error!("singleton_wheel_v2: reinsert_ready: {err:#}"); - } - } - tracing::debug!("singleton_wheel_v2: done reinserting {reinserted}. total scheduled={len}"); - - for (_queue_name, queue) in to_shrink.drain() { - queue.queue.shrink(); - } - to_shrink.shrink_to_fit(); - } - _ = shutdown.shutting_down() => { - tracing::info!("singleton_wheel: stopping"); - return Ok(()); - } - } - } - } - - pub async fn run_singleton_wheel() -> anyhow::Result<()> { - let mut shutdown = ShutdownSubcription::get(); - - tracing::debug!("singleton_wheel: starting up"); - - async fn reinsert_ready( - msg: Message, - to_shrink: &mut HashMap, - ) -> anyhow::Result<()> { - if !msg.is_meta_loaded() { - msg.load_meta().await?; - } - let queue_name = msg.get_queue_name()?; - // Use get_opt rather than resolve here. If the queue is not currently - // tracked in the QueueManager then this message cannot possibly belong - // to it. Using resolve would have the side effect of creating an empty - // queue for it, which will then age out later. It's a waste to do that, - // so we just check and skip. - let queue = QueueManager::get_opt(&queue_name) - .ok_or_else(|| anyhow::anyhow!("no scheduled queue"))?; - - if let Some(b) = AdminBounceEntry::get_for_queue_name(&queue.name) { - // Note that this will cause the msg to be removed from the - // queue so the remove() check below will return false - queue.bounce_all(&b).await; - } - - // Verify that the message is still in the queue - match &queue.queue { - QueueStructure::SingletonTimerWheel(q) => { - fn remove(q: &FairMutex>, msg: &Message) -> bool { - q.lock().remove(msg) - } - - if remove(q, &msg) { - queue.metrics().sub(1); - queue - .insert_ready(msg, InsertReason::DueTimeWasReached.into(), None) - .await?; - if !to_shrink.contains_key(&queue_name) { - to_shrink.insert(queue_name, queue); - } - } - } - _ => { - anyhow::bail!("impossible queue strategy"); - } - } - - Ok(()) - } - - let mut to_shrink = HashMap::new(); - - loop { - tokio::select! { - _ = tokio::time::sleep(Duration::from_secs(3)) => { - TOTAL_QMAINT_RUNS.inc(); - - fn pop() -> (Vec, usize) { - let mut wheel = SINGLETON_WHEEL.lock(); - - let msgs = if let PopResult::Items(weak_messages) = wheel.pop() { - tracing::debug!("singleton_wheel: popped {} messages", weak_messages.len()); - weak_messages - } else { - vec![] - }; - - (msgs, wheel.len()) - } - - let mut reinserted = 0; - let (msgs, len) = pop(); - let mut messages = vec![]; - for weak_message in msgs { - if let Some(msg) = weak_message.upgrade() { - messages.push(msg); - } - } - wait_for_message_batch(&messages).await; - for msg in messages { - reinserted += 1; - if let Err(err) = reinsert_ready(msg, &mut to_shrink).await { - tracing::error!("singleton_wheel: reinsert_ready: {err:#}"); - } - } - tracing::debug!("singleton_wheel: done reinserting {reinserted}. total scheduled={len}"); - - for (_queue_name, queue) in to_shrink.drain() { - queue.queue.shrink(); - } - to_shrink.shrink_to_fit(); - } - _ = shutdown.shutting_down() => { - tracing::info!("singleton_wheel: stopping"); - return Ok(()); - } - } - } - } - pub fn is_empty(&self) -> bool { self.queue.is_empty() } diff --git a/crates/kumod/src/queue/strategy.rs b/crates/kumod/src/queue/strategy.rs index ea5702d7..c36abf75 100644 --- a/crates/kumod/src/queue/strategy.rs +++ b/crates/kumod/src/queue/strategy.rs @@ -1,5 +1,4 @@ -use crate::queue::maintainer::QMAINT_RUNTIME; -use crate::queue::queue::Queue; +use crate::queue::maintainer::{run_singleton_wheel_v1, run_singleton_wheel_v2, QMAINT_RUNTIME}; use chrono::{DateTime, Utc}; use crossbeam_skiplist::SkipSet; use message::message::WeakMessage; @@ -209,8 +208,8 @@ impl QueueStructure { STARTED_SINGLETON_WHEEL.call_once(|| { QMAINT_RUNTIME .spawn("singleton_wheel".to_string(), async move { - if let Err(err) = Queue::run_singleton_wheel().await { - tracing::error!("run_singleton_wheel: {err:#}"); + if let Err(err) = run_singleton_wheel_v1().await { + tracing::error!("run_singleton_wheel_v1: {err:#}"); } }) .expect("failed to spawn singleton_wheel"); @@ -241,7 +240,7 @@ impl QueueStructure { STARTED_SINGLETON_WHEEL_V2.call_once(|| { QMAINT_RUNTIME .spawn("singleton_wheel_v2".to_string(), async move { - if let Err(err) = Queue::run_singleton_wheel_v2().await { + if let Err(err) = run_singleton_wheel_v2().await { tracing::error!("run_singleton_wheel_v2: {err:#}"); } })