diff --git a/crates/kumod/src/queue/maintainer.rs b/crates/kumod/src/queue/maintainer.rs index 70ab575a..e142f5d2 100644 --- a/crates/kumod/src/queue/maintainer.rs +++ b/crates/kumod/src/queue/maintainer.rs @@ -159,97 +159,51 @@ pub async fn maintain_named_queue(q: &QueueHandle) -> anyhow::Result<()> { } } -pub async fn run_singleton_wheel_v2() -> anyhow::Result<()> { - let mut shutdown = ShutdownSubcription::get(); +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"))?; - 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(()) + 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; } - let mut to_shrink = HashMap::new(); + fn remove(q: &FairMutex>, msg: &Message) -> bool { + q.lock().remove(msg) + } - 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()) + // 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) { + to_shrink.insert(queue_name, queue); } - - 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(()); } } + _ => { + anyhow::bail!("impossible queue strategy"); + } } + + Ok(()) } pub async fn run_singleton_wheel_v1() -> anyhow::Result<()> { @@ -257,53 +211,6 @@ pub async fn run_singleton_wheel_v1() -> anyhow::Result<()> { 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 { @@ -353,3 +260,49 @@ pub async fn run_singleton_wheel_v1() -> anyhow::Result<()> { } } } + +pub async fn run_singleton_wheel_v2() -> anyhow::Result<()> { + let mut shutdown = ShutdownSubcription::get(); + + tracing::debug!("singleton_wheel_v2: starting up"); + + 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(()); + } + } + } +}