queue: avoid potential "singleton_wheel: reinsert_ready: metadata must be loaded first"

There is a potential race between the timerwheel tick and bulk queue
operations (flushing, bouncing).  The logic in the tick case has
some accommodations for this, but there is another outstanding
that can result in racing to manage the metadata load state
on a message.

In the tick case, we load it if needed, then resolve the queue
name, so that we can resolve the Queue handle.

We need the queue handle to decide whether we are responsible
for the message, or whether we might be racing with a concurrent
action.

In the race case, the other actor may have decided to release
the message metadata, which will cause the queue name resolution
to fail.

While we could just shrug and silently ignore the error in that
case, that makes me uneasy.

What this commit does is adjust the v1 wheel entry so that we
capture weak references to both the message and its containing
queue.

Then when we tick, we can simply upgrade both of those and
reconcile without needing to manipulate any message metadata.

These changes unfortunately result in re-duplication of some
of the logic that I recently refactored to be shared with the
v2 tick implementation.

The v2 tick implementation has the same edge case around
metadata, but cannot be easily adjusted to use the same
pattern.

So this commit sticks a comment in the code; nobody is using
v2 AFAIK, and there have been reports of some wonky behavior
with it.  My recommendation at this time is to avoid using
the v2 wheel and we can fix this all up for it later.

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<String, QueueHandle>,
 ) -> 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<HashSet<Message>>, 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<Message> = 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<String, QueueHandle>,
+) -> 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<Message>, total_scheduled: usize) {
+async fn reinsert_batch_v2(messages: Vec<Message>, 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<Message>, total_scheduled: usize) {
     to_shrink.shrink_to_fit();
 }
 
-async fn process_batch(messages: Vec<Message>, total_scheduled: usize) {
+async fn process_batch_v2(messages: Vec<Message>, 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<WeakMessage>, usize) {
+        fn pop() -> (Vec<WheelV1Entry>, 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;
     }
 }
 

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<Self>, 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<Self>,
         msg: Message,
         rebind: &Arc<AdminRebindEntry>,
         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<AdminRebindEntry>) {
+    pub async fn rebind_all(self: &Arc<Self>, rebind: &Arc<AdminRebindEntry>) {
         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<Self>,
         msg: Message,
         increment_attempts: IncrementAttempts,
         delay: Option<chrono::Duration>,
@@ -734,7 +734,7 @@ impl Queue {
 
     #[instrument(skip(self, msg))]
     async fn insert_delayed(
-        &self,
+        self: &Arc<Self>,
         msg: Message,
         context: InsertContext,
     ) -> anyhow::Result<InsertResult> {
@@ -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<Self>,
+        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<Self>,
         msg: Message,
         mut context: InsertContext,
         deadline: Option<Instant>,
@@ -1397,7 +1401,7 @@ impl Queue {
     /// into this queue
     #[instrument(fields(self.name), skip(self, msg))]
     pub async fn insert(
-        &self,
+        self: &Arc<Self>,
         msg: Message,
         context: InsertContext,
         deadline: Option<Instant>,

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<Arc<FairMutex<TimeQ<WeakMessage>>>> =
+pub static SINGLETON_WHEEL: LazyLock<Arc<FairMutex<TimeQ<WheelV1Entry>>>> =
     LazyLock::new(|| Arc::new(FairMutex::new(TimeQ::new())));
 
 pub static SINGLETON_WHEEL_V2: LazyLock<Arc<FairMutex<TriTimeQ>>> =
@@ -35,6 +37,32 @@ pub enum QueueInsertResult {
     Full(Message),
 }
 
+#[derive(Debug)]
+pub struct WheelV1Entry {
+    weak: WeakMessage,
+    queue: Weak<Queue>,
+}
+
+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<TimeQ<Message>>),
     SkipList(SkipSet<DelayedEntry>),
@@ -165,7 +193,7 @@ impl QueueStructure {
         }
     }
 
-    pub fn insert(&self, msg: Message) -> QueueInsertResult {
+    pub fn insert(&self, msg: Message, queue: &Arc<Queue>) -> 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<LifeCycle> = 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(_)));
     }