diff --git a/crates/kumod/src/queue/maintainer.rs b/crates/kumod/src/queue/maintainer.rs index 320af98b..ca2876d9 100644 --- a/crates/kumod/src/queue/maintainer.rs +++ b/crates/kumod/src/queue/maintainer.rs @@ -2,13 +2,13 @@ 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::strategy::{QueueStructure, WheelV1Entry, SINGLETON_WHEEL, SINGLETON_WHEEL_V2}; use crate::queue::wait_for_message_batch; use crate::ready_queue::ReadyQueueManager; use anyhow::Context; use kumo_server_lifecycle::{Activity, ShutdownSubcription}; use kumo_server_runtime::Runtime; -use message::message::{MessageList, WeakMessage}; +use message::message::MessageList; use message::Message; use parking_lot::FairMutex; use prometheus::{Histogram, IntCounter}; @@ -177,11 +177,94 @@ pub async fn maintain_named_queue(q: &QueueHandle) -> anyhow::Result<()> { async fn reinsert_ready( msg: Message, + queue: QueueHandle, to_shrink: &mut HashMap, ) -> anyhow::Result<()> { - if !msg.is_meta_loaded() { - msg.load_meta().await?; + 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; } + + fn remove(q: &FairMutex>, msg: &Message) -> bool { + q.lock().remove(msg) + } + + // Verify that the message is still in the queue + match &queue.queue { + QueueStructure::SingletonTimerWheel(q) | QueueStructure::SingletonTimerWheelV2(q) => { + if remove(q, &msg) { + queue.metrics().sub(1); + queue + .insert_ready(msg, InsertReason::DueTimeWasReached.into(), None) + .await?; + if !to_shrink.contains_key(queue.name.as_str()) { + to_shrink.insert(queue.name.to_string(), queue); + } + } + } + _ => { + anyhow::bail!("impossible queue strategy"); + } + } + + Ok(()) +} + +async fn reinsert_batch(messages: Vec<(Message, QueueHandle)>, total_scheduled: usize) { + let mut to_shrink = HashMap::new(); + let mut reinserted = 0; + + let messages_only: Vec = messages.iter().map(|(msg, _q)| msg.clone()).collect(); + wait_for_message_batch(&messages_only).await; + for (msg, queue) in messages { + reinserted += 1; + if let Err(err) = reinsert_ready(msg, queue, &mut to_shrink).await { + tracing::error!("singleton_wheel: reinsert_ready: {err:#}"); + } + } + tracing::debug!( + "singleton_wheel: done reinserting {reinserted}. total scheduled={total_scheduled}" + ); + + for (_queue_name, queue) in to_shrink.drain() { + queue.queue.shrink(); + } + to_shrink.shrink_to_fit(); +} + +async fn process_batch(messages: Vec<(Message, QueueHandle)>, total_scheduled: usize) { + if messages.is_empty() { + return; + } + + if !SPAWN_REINSERTION.load(Ordering::Relaxed) { + reinsert_batch(messages, total_scheduled).await; + return; + } + + if let Err(err) = + QMAINT_RUNTIME.spawn("reinsert_batch", reinsert_batch(messages, total_scheduled)) + { + tracing::error!("run_singleton_wheel_v1: failed to spawn reinsert_batch: {err:#}"); + } +} + +async fn reinsert_ready_v2( + msg: Message, + to_shrink: &mut HashMap, +) -> anyhow::Result<()> { + // Note that there is a potential race here that we cannot detect + // until we try to call the v2 cancel method. + // If we are no longer responsible for the message (detected later), + // we might load the metadata here while a concurrent actor is requeing + // the message and releasing the metadata. + // That can cause the get_queue_name method to fail. + // We address this in the v1 flavor of this function by keeping + // the associate Queue handle so that we don't need to muck with the + // metadata at all. + // Here, we don't and can't do that. + msg.load_meta_if_needed().await?; let queue_name = msg.get_queue_name().context("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 @@ -209,8 +292,8 @@ async fn reinsert_ready( queue .insert_ready(msg, InsertReason::DueTimeWasReached.into(), None) .await?; - if !to_shrink.contains_key(&queue_name) { - to_shrink.insert(queue_name, queue); + if !to_shrink.contains_key(queue.name.as_str()) { + to_shrink.insert(queue.name.to_string(), queue); } } } @@ -222,14 +305,14 @@ async fn reinsert_ready( Ok(()) } -async fn reinsert_batch(messages: Vec, total_scheduled: usize) { +async fn reinsert_batch_v2(messages: Vec, total_scheduled: usize) { let mut to_shrink = HashMap::new(); let mut reinserted = 0; wait_for_message_batch(&messages).await; for msg in messages { reinserted += 1; - if let Err(err) = reinsert_ready(msg, &mut to_shrink).await { + if let Err(err) = reinsert_ready_v2(msg, &mut to_shrink).await { tracing::error!("singleton_wheel: reinsert_ready: {err:#}"); } } @@ -243,19 +326,20 @@ async fn reinsert_batch(messages: Vec, total_scheduled: usize) { to_shrink.shrink_to_fit(); } -async fn process_batch(messages: Vec, total_scheduled: usize) { +async fn process_batch_v2(messages: Vec, total_scheduled: usize) { if messages.is_empty() { return; } if !SPAWN_REINSERTION.load(Ordering::Relaxed) { - reinsert_batch(messages, total_scheduled).await; + reinsert_batch_v2(messages, total_scheduled).await; return; } - if let Err(err) = - QMAINT_RUNTIME.spawn("reinsert_batch", reinsert_batch(messages, total_scheduled)) - { + if let Err(err) = QMAINT_RUNTIME.spawn( + "reinsert_batch", + reinsert_batch_v2(messages, total_scheduled), + ) { tracing::error!("run_singleton_wheel_v1: failed to spawn reinsert_batch: {err:#}"); } } @@ -284,7 +368,7 @@ async fn run_singleton_wheel_v1() -> anyhow::Result<()> { tracing::trace!("singleton_wheel_v1 ticking"); TOTAL_QMAINT_RUNS.inc(); - fn pop() -> (Vec, usize) { + fn pop() -> (Vec, usize) { let _timer = POP_LATENCY.start_timer(); let mut wheel = SINGLETON_WHEEL.lock(); @@ -302,8 +386,8 @@ async fn run_singleton_wheel_v1() -> anyhow::Result<()> { let mut messages = vec![]; for weak_message in msgs { - if let Some(msg) = weak_message.upgrade() { - messages.push(msg); + if let Some((msg, queue)) = weak_message.upgrade() { + messages.push((msg, queue)); } } process_batch(messages, total_scheduled).await; @@ -339,7 +423,7 @@ async fn run_singleton_wheel_v2() -> anyhow::Result<()> { } let (messages, total_scheduled) = pop(); - process_batch(messages.into_iter().collect(), total_scheduled).await; + process_batch_v2(messages.into_iter().collect(), total_scheduled).await; } } diff --git a/crates/kumod/src/queue/queue.rs b/crates/kumod/src/queue/queue.rs index f900421e..f63b7e2c 100644 --- a/crates/kumod/src/queue/queue.rs +++ b/crates/kumod/src/queue/queue.rs @@ -346,9 +346,9 @@ impl Queue { } /// Insert into the timeq, and updates the counters. - fn timeq_insert(&self, msg: Message) -> Result<(), Message> { + fn timeq_insert(self: &Arc, msg: Message) -> Result<(), Message> { tracing::trace!("timeq_insert {} due={:?}", self.name, msg.get_due()); - match self.queue.insert(msg) { + match self.queue.insert(msg, self) { QueueInsertResult::Inserted { should_notify } => { self.metrics().inc(); if should_notify { @@ -373,7 +373,7 @@ impl Queue { } async fn do_rebind( - &self, + self: &Arc, msg: Message, rebind: &Arc, context: InsertContext, @@ -441,7 +441,7 @@ impl Queue { } Ok(queue) => { queue_holder = queue; - &*queue_holder + &queue_holder } }; @@ -494,7 +494,7 @@ impl Queue { } #[instrument(skip(self))] - pub async fn rebind_all(&self, rebind: &Arc) { + pub async fn rebind_all(self: &Arc, rebind: &Arc) { let msgs = self.drain_timeq(); let count = msgs.len(); if count > 0 { @@ -668,7 +668,7 @@ impl Queue { /// The requeue_message event is NOT called by this function. #[instrument(skip(self, msg))] pub async fn requeue_message_internal( - &self, + self: &Arc, msg: Message, increment_attempts: IncrementAttempts, delay: Option, @@ -734,7 +734,7 @@ impl Queue { #[instrument(skip(self, msg))] async fn insert_delayed( - &self, + self: &Arc, msg: Message, context: InsertContext, ) -> anyhow::Result { @@ -764,7 +764,11 @@ impl Queue { } #[instrument(skip(self, msg))] - async fn force_into_delayed(&self, msg: Message, context: InsertContext) -> anyhow::Result<()> { + async fn force_into_delayed( + self: &Arc, + msg: Message, + context: InsertContext, + ) -> anyhow::Result<()> { tracing::trace!("force_into_delayed {}", msg.id()); loop { match self.insert_delayed(msg.clone(), context.clone()).await? { @@ -909,7 +913,7 @@ impl Queue { #[instrument(skip(self, msg))] pub async fn insert_ready( - &self, + self: &Arc, msg: Message, mut context: InsertContext, deadline: Option, @@ -1397,7 +1401,7 @@ impl Queue { /// into this queue #[instrument(fields(self.name), skip(self, msg))] pub async fn insert( - &self, + self: &Arc, msg: Message, context: InsertContext, deadline: Option, diff --git a/crates/kumod/src/queue/strategy.rs b/crates/kumod/src/queue/strategy.rs index 262cea4d..6378ec11 100644 --- a/crates/kumod/src/queue/strategy.rs +++ b/crates/kumod/src/queue/strategy.rs @@ -1,4 +1,6 @@ use crate::queue::maintainer::{start_singleton_wheel_v1, start_singleton_wheel_v2}; +use crate::queue::queue::QueueHandle; +use crate::queue::Queue; use chrono::{DateTime, Utc}; use crossbeam_skiplist::SkipSet; use message::message::WeakMessage; @@ -8,11 +10,11 @@ use mlua::prelude::*; use parking_lot::FairMutex; use serde::{Deserialize, Serialize}; use std::collections::HashSet; -use std::sync::{Arc, LazyLock}; +use std::sync::{Arc, LazyLock, Weak}; use std::time::Duration; -use timeq::{PopResult, TimeQ, TimerError}; +use timeq::{PopResult, TimeQ, TimerEntryWithDelay, TimerError}; -pub static SINGLETON_WHEEL: LazyLock>>> = +pub static SINGLETON_WHEEL: LazyLock>>> = LazyLock::new(|| Arc::new(FairMutex::new(TimeQ::new()))); pub static SINGLETON_WHEEL_V2: LazyLock>> = @@ -35,6 +37,32 @@ pub enum QueueInsertResult { Full(Message), } +#[derive(Debug)] +pub struct WheelV1Entry { + weak: WeakMessage, + queue: Weak, +} + +impl WheelV1Entry { + pub fn upgrade(self) -> Option<(Message, QueueHandle)> { + let message = self.weak.upgrade()?; + let queue = self.queue.upgrade()?; + Some((message, queue)) + } +} + +impl TimerEntryWithDelay for WheelV1Entry { + fn delay(&self) -> Duration { + match self.weak.upgrade() { + None => { + // Dangling/Cancelled. Make it appear due immediately + Duration::from_millis(0) + } + Some(msg) => msg.delay(), + } + } +} + pub enum QueueStructure { TimerWheel(FairMutex>), SkipList(SkipSet), @@ -165,7 +193,7 @@ impl QueueStructure { } } - pub fn insert(&self, msg: Message) -> QueueInsertResult { + pub fn insert(&self, msg: Message, queue: &Arc) -> QueueInsertResult { match self { Self::TimerWheel(q) => match q.lock().insert(msg) { Ok(()) => QueueInsertResult::Inserted { @@ -198,7 +226,10 @@ impl QueueStructure { } Self::SingletonTimerWheel(q) => { let mut wheel = SINGLETON_WHEEL.lock(); - match wheel.insert(msg.weak()) { + match wheel.insert(WheelV1Entry { + weak: msg.weak(), + queue: Arc::downgrade(queue), + }) { Ok(()) => { q.lock().insert(msg); drop(wheel); @@ -316,6 +347,7 @@ impl Ord for DelayedEntry { #[cfg(test)] mod test { use super::*; + use kumo_server_lifecycle::LifeCycle; use message::EnvelopeAddress; use spool::SpoolId; @@ -339,7 +371,22 @@ mod test { .await .unwrap(); eprintln!("due {due:?}"); - let result = qs.insert(msg); + + // This is a bit inelegant; the queue object that we need + // to pass to the insert method needs to be able to construct + // an Activity instance, which will fail with "shutting down" + // if no lifecycle has been started. + // Let's start one up in this test context. + // This is bad because future tests that might have this dependency + // might now intermittently start to pass depending on their + // runtime ordering wrt. tests that call this function. + static TEST_LIFE_CYCLE: LazyLock = LazyLock::new(|| LifeCycle::new()); + LazyLock::force(&TEST_LIFE_CYCLE); + let queue = Queue::new(format!("dummy-{:?}", qs.strategy())) + .await + .unwrap(); + + let result = qs.insert(msg, &queue); eprintln!("result: {result:?}"); assert!(matches!(result, QueueInsertResult::Full(_))); }