qrusty 0.20.9

// src/memory_storage.rs
// Implements: SYS-0013, SYS-0018, PER-0011

//! # In-Memory Storage Backend
//!
//! Provides a volatile, in-memory implementation of `StorageApi` using a
//! `BTreeMap` for sorted key-value storage. This backend preserves the
//! same key-ordering semantics as the RocksDB backend, ensuring identical
//! functional behavior for priority ordering, locking, ACK/NACK, DLQ,
//! and duplicate detection.
//!
//! ## When to Use
//!
//! - Hardware with failing or absent persistent storage
//! - Development and testing environments
//! - Scenarios where message durability is not required
//!
//! ## Limitations
//!
//! - All data is lost on process restart
//! - No crash recovery or WAL
//! - Write operations hold an exclusive lock on the data store

use anyhow::Result;
use chrono::{DateTime, Utc};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::sync::{Arc, Mutex};
use tokio::sync::RwLock;

use crate::message::{
    BatchAckResult, BatchNackResult, Message, Priority, PriorityOrdering, QueueConfig, QueueStats,
};
use crate::storage::{LockedIdIndex, QueueCounts, RenameQueueError};

/// Computes a 128-bit xxh3 hash of a payload string for duplicate detection.
/// See `storage::hash_payload` for rationale.
// Implements: PER-0014
fn hash_payload(payload: &str) -> [u8; 16] {
    xxhash_rust::xxh3::xxh3_128(payload.as_bytes()).to_ne_bytes()
}

/// In-memory storage backend using a sorted `BTreeMap` to replicate
/// RocksDB's lexicographic key ordering.
pub struct MemoryStorage {
    /// Sorted key-value store. Key format matches RocksDB Storage:
    /// `{queue}/{priority_key}/{timestamp}/{uuid}` for messages,
    /// `_queue_config/{queue}` for configs,
    /// `_dlq/{queue}/{id}` for dead-letter entries.
    data: Arc<RwLock<BTreeMap<String, Vec<u8>>>>,

    /// In-memory lock tracking (identical to Storage).
    locked_index: Arc<RwLock<HashMap<String, DateTime<Utc>>>>,

    /// Secondary index mapping `(queue, message_id) -> storage_key` for
    /// O(1) lookup of locked messages in `ack`, `nack`, `renew`,
    /// `batch_ack`, and `batch_nack`.  Mirrors the persistent backend's
    /// `locked_id_index`.
    ///
    /// MemoryStorage does NOT need the persistent backend's cap or
    /// sticky `untracked_locks_possible` flag: the index is the sole
    /// source of truth for lock state, it has no persistence to drift
    /// from, and there is no post-restart recovery to guard against.
    /// A miss on `(queue, message_id)` therefore provably means the
    /// message is not locked, and the fast paths return `Ok(false)`
    /// without any fallback scan.
    ///
    /// Implements: DLV-0014
    locked_id_index: Arc<RwLock<LockedIdIndex>>,

    /// Queue configurations (identical to Storage).
    queue_configs: Arc<RwLock<HashMap<String, QueueConfig>>>,

    /// Per-queue payload hash sets for duplicate detection (identical to Storage).
    /// Stores 128-bit xxh3 hashes instead of full payload strings (PER-0014).
    payload_sets: Arc<RwLock<HashMap<String, HashSet<[u8; 16]>>>>,

    /// Per-queue available/locked counts for O(n_queues) stats.
    ///
    /// Implements: SYS-0018
    ///
    /// Maintained incrementally by every state-changing operation so that
    /// `get_queue_stats`, `get_all_queue_stats`, and `list_queues` can
    /// read directly from this cache instead of scanning `data`.
    queue_counters: Arc<Mutex<HashMap<String, QueueCounts>>>,
}

impl Default for MemoryStorage {
    fn default() -> Self {
        Self::new()
    }
}

/// Percent-encodes characters that conflict with the storage key separator.
/// See `storage::encode_text_priority` for full documentation.
fn encode_text_priority(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for b in s.bytes() {
        match b {
            b'%' => out.push_str("%25"),
            b'/' => out.push_str("%2F"),
            0 => out.push_str("%00"),
            _ => out.push(b as char),
        }
    }
    out
}

impl MemoryStorage {
    /// Creates a new in-memory storage instance.
    ///
    /// Logs a warning that all data will be lost on restart.
    pub fn new() -> Self {
        tracing::warn!("Running with in-memory storage. All data will be lost on restart.");
        Self {
            data: Arc::new(RwLock::new(BTreeMap::new())),
            locked_index: Arc::new(RwLock::new(HashMap::new())),
            locked_id_index: Arc::new(RwLock::new(HashMap::new())),
            queue_configs: Arc::new(RwLock::new(HashMap::new())),
            payload_sets: Arc::new(RwLock::new(HashMap::new())),
            queue_counters: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// O(1) lookup of the storage key for a currently-locked
    /// `(queue, message_id)` pair.  Returns `None` if the entry is
    /// not in the secondary index — which, on the MemoryStorage
    /// backend, is conclusive (the message is not locked).
    async fn locked_key_for_inner(&self, queue: &str, message_id: &str) -> Option<String> {
        self.locked_id_index
            .read()
            .await
            .get(queue)
            .and_then(|m| m.get(message_id).cloned())
    }

    /// Removes the (queue, message_id) entry from the secondary index
    /// and cleans up the queue's sub-map if it becomes empty.
    async fn remove_locked_id_index_entry(&self, queue: &str, message_id: &str) {
        let mut idx = self.locked_id_index.write().await;
        if let Some(map) = idx.get_mut(queue) {
            map.remove(message_id);
            if map.is_empty() {
                idx.remove(queue);
            }
        }
    }

    /// Test-only: inject a secondary-index entry pointing at a storage
    /// key that may or may not exist.  Lets tests exercise stale-entry
    /// cleanup on the fast paths.
    #[cfg(any(test, feature = "test-helpers"))]
    #[doc(hidden)]
    pub async fn __test_insert_locked_id_index_entry(
        &self,
        queue: &str,
        message_id: &str,
        key: &str,
    ) {
        self.locked_id_index
            .write()
            .await
            .entry(queue.to_string())
            .or_default()
            .insert(message_id.to_string(), key.to_string());
    }

    /// Test-only: check whether the secondary index has an entry for a
    /// given `(queue, message_id)`.
    #[cfg(any(test, feature = "test-helpers"))]
    #[doc(hidden)]
    pub async fn __test_locked_id_index_has(&self, queue: &str, message_id: &str) -> bool {
        self.locked_id_index
            .read()
            .await
            .get(queue)
            .is_some_and(|m| m.contains_key(message_id))
    }

    /// Extract the queue name from a storage key like `queue/priority/ts/uuid`.
    /// Returns None for keys outside the message namespace (e.g. `_dlq/…`,
    /// `_queue_config/…`) or malformed keys.
    fn queue_from_key(key: &str) -> Option<&str> {
        if key.starts_with('_') {
            return None;
        }
        let mut parts = key.splitn(2, '/');
        parts.next()
    }

    /// Generate a storage key for a message, identical to Storage::generate_message_key.
    fn generate_message_key(&self, msg: &Message, config: &QueueConfig) -> String {
        let ts = msg.created_at.timestamp_nanos_opt().unwrap_or(0);
        match &msg.priority {
            Priority::Numeric(n) => {
                let priority_key = match config.ordering {
                    PriorityOrdering::MaxFirst => u64::MAX - n,
                    PriorityOrdering::MinFirst => *n,
                    PriorityOrdering::Fifo => 0,
                };
                format!("{}/{:020}/{:016}/{}", msg.queue, priority_key, ts, msg.id)
            }
            Priority::Text(s) => {
                let key_segment = match config.ordering {
                    PriorityOrdering::MaxFirst => {
                        // Pad to fixed width with 0x00 before complementing so
                        // that variable-length strings reverse correctly.
                        let bytes = s.as_bytes();
                        let pad_len = 512usize.max(bytes.len());
                        let mut buf = vec![0u8; pad_len];
                        buf[..bytes.len()].copy_from_slice(bytes);
                        format!(
                            "s:{}",
                            buf.iter()
                                .map(|b| format!("{:02x}", b ^ 0xFF))
                                .collect::<String>()
                        )
                    }
                    PriorityOrdering::MinFirst => format!("s:{}", encode_text_priority(s)),
                    PriorityOrdering::Fifo => "s:".to_string(),
                };
                format!("{}/{}/{:016}/{}", msg.queue, key_segment, ts, msg.id)
            }
        }
    }

    /// Returns the queue config, defaulting to MaxFirst if not configured.
    async fn get_queue_config(&self, queue_name: &str) -> QueueConfig {
        self.queue_configs
            .read()
            .await
            .get(queue_name)
            .cloned()
            .unwrap_or_default()
    }

    /// Checks if a queue has any messages or a stored configuration.
    async fn queue_exists_inner(&self, queue_name: &str) -> Result<bool> {
        // Check in-memory configs first
        if self.queue_configs.read().await.contains_key(queue_name) {
            return Ok(true);
        }
        // Check for any message keys
        let data = self.data.read().await;
        let prefix = format!("{}/", queue_name);
        Ok(data
            .range(prefix.clone()..)
            .next()
            .is_some_and(|(k, _)| k.starts_with(&prefix)))
    }

    /// Creates or updates queue configuration.
    async fn configure_queue_inner(&self, queue_name: &str, config: QueueConfig) {
        let old_config = self.get_queue_config(queue_name).await;
        let existed = self.queue_exists_inner(queue_name).await.unwrap_or(false);

        let mut effective_config = config;
        if existed {
            effective_config.ordering = old_config.ordering;
        }

        // Store in memory configs
        self.queue_configs
            .write()
            .await
            .insert(queue_name.to_string(), effective_config.clone());

        // Persist config in the BTreeMap (mirrors RocksDB _queue_config/ key)
        let config_key = format!("_queue_config/{}", queue_name);
        if let Ok(config_json) = serde_json::to_vec(&effective_config) {
            self.data.write().await.insert(config_key, config_json);
        }

        // SYS-0018: ensure a counter entry exists so list_queues and stats
        // include the queue even when it has no messages yet.
        self.queue_counters
            .lock()
            .unwrap()
            .entry(queue_name.to_string())
            .or_default();

        // If we just disabled duplicates, de-dupe and build payload set
        if old_config.allow_duplicates && !effective_config.allow_duplicates {
            let removed = self.dedupe_unlocked_messages(queue_name).await;
            if removed > 0 {
                tracing::info!(
                    "De-duplicated {} unlocked message(s) in queue '{}' after disabling duplicates",
                    removed,
                    queue_name
                );
            }

            // Rebuild payload set
            let data = self.data.read().await;
            let prefix = format!("{}/", queue_name);
            let mut set = HashSet::new();
            for (k, v) in data.range(prefix.clone()..) {
                if !k.starts_with(&prefix) {
                    break;
                }
                if let Ok(msg) = serde_json::from_slice::<Message>(v) {
                    set.insert(hash_payload(&msg.payload));
                }
            }
            self.payload_sets
                .write()
                .await
                .insert(queue_name.to_string(), set);
        }

        // If we just enabled duplicates, drop the payload set
        if !old_config.allow_duplicates && effective_config.allow_duplicates {
            self.payload_sets.write().await.remove(queue_name);
        }
    }

    /// Remove duplicate unlocked messages by payload, keeping the first occurrence.
    async fn dedupe_unlocked_messages(&self, queue_name: &str) -> usize {
        let mut data = self.data.write().await;
        let prefix = format!("{}/", queue_name);
        let now = Utc::now();

        let mut seen: HashSet<String> = HashSet::new();
        let mut keys_to_remove: Vec<String> = Vec::new();

        for (k, v) in data.range(prefix.clone()..) {
            if !k.starts_with(&prefix) {
                break;
            }
            if let Ok(msg) = serde_json::from_slice::<Message>(v) {
                let is_locked = msg.locked_until.is_some_and(|lu| lu > now);
                if is_locked {
                    // Locked messages always count towards "seen" but aren't removed
                    seen.insert(msg.payload);
                    continue;
                }
                if !seen.insert(msg.payload.clone()) {
                    keys_to_remove.push(k.clone());
                }
            }
        }

        let count = keys_to_remove.len();
        for k in keys_to_remove {
            data.remove(&k);
        }

        // SYS-0018: these were all unlocked messages, so available -= count.
        if count > 0 {
            let mut counters = self.queue_counters.lock().unwrap();
            if let Some(entry) = counters.get_mut(queue_name) {
                entry.available = entry.available.saturating_sub(count as u64);
            }
        }

        count
    }

    /// Unlocks a single message by storage key. Returns true if actually
    /// unlocked.  Also removes the secondary-index entry so the fast
    /// paths (DLV-0014) don't serve a now-stale key.
    async fn unlock_message_by_key(&self, key: &str) -> Result<bool> {
        let unlocked_info: Option<(String, String)> = {
            let mut data = self.data.write().await;
            match data.get(key) {
                Some(value) => {
                    let mut msg: Message = serde_json::from_slice(value)?;
                    if msg.locked_until.is_some() {
                        msg.locked_until = None;
                        msg.locked_by = None;
                        let new_value = serde_json::to_vec(&msg)?;
                        data.insert(key.to_string(), new_value);
                        Some((msg.queue, msg.id))
                    } else {
                        None
                    }
                }
                None => None,
            }
        };

        if let Some((queue, id)) = unlocked_info {
            // DLV-0014: keep the secondary index consistent with
            // locked_index; expired locks removed here must not serve
            // stale keys to later fast-path lookups.
            self.remove_locked_id_index_entry(&queue, &id).await;
            Ok(true)
        } else {
            Ok(false)
        }
    }
}

#[async_trait::async_trait]
impl crate::api::StorageApi for MemoryStorage {
    async fn queue_exists(&self, queue_name: &str) -> Result<bool> {
        self.queue_exists_inner(queue_name).await
    }

    async fn create_queue(&self, queue_name: &str, config: QueueConfig) {
        self.configure_queue_inner(queue_name, config).await;
    }

    async fn rename_queue(
        &self,
        from: &str,
        to: &str,
    ) -> std::result::Result<(), RenameQueueError> {
        if from.trim().is_empty() || to.trim().is_empty() {
            return Err(RenameQueueError::Storage(anyhow::anyhow!(
                "queue names cannot be empty"
            )));
        }
        if from == to {
            return Ok(());
        }
        if !self
            .queue_exists_inner(from)
            .await
            .map_err(RenameQueueError::Storage)?
        {
            return Err(RenameQueueError::NotFound);
        }
        if self
            .queue_exists_inner(to)
            .await
            .map_err(RenameQueueError::Storage)?
        {
            return Err(RenameQueueError::AlreadyExists);
        }

        let existing_config = self.get_queue_config(from).await;
        let now = Utc::now();

        let mut data = self.data.write().await;
        // (old_key, new_key, locked_until) for still-locked messages so
        // the primary locked_index update below can remove exact keys
        // instead of retain()-scanning the whole cross-queue map.
        let mut moved_locked: Vec<(String, String, DateTime<Utc>)> = Vec::new();
        // DLV-0014: (message_id, new_storage_key) to rebuild the
        // secondary index under the new queue name.
        let mut moved_locked_ids: Vec<(String, String)> = Vec::new();

        // Collect and move message keys
        let prefix_from = format!("{}/", from);
        let entries: Vec<(String, Vec<u8>)> = data
            .range(prefix_from.clone()..)
            .take_while(|(k, _)| k.starts_with(&prefix_from))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();

        for (key, value) in entries {
            let mut msg: Message = serde_json::from_slice(&value).map_err(|e| {
                RenameQueueError::Storage(anyhow::anyhow!("deserialize message failed: {e}"))
            })?;
            msg.queue = to.to_string();

            let new_key = self.generate_message_key(&msg, &existing_config);
            let new_value = serde_json::to_vec(&msg).map_err(|e| {
                RenameQueueError::Storage(anyhow::anyhow!("serialize message failed: {e}"))
            })?;

            if let Some(locked_until) = msg.locked_until {
                if locked_until > now {
                    moved_locked.push((key.clone(), new_key.clone(), locked_until));
                    moved_locked_ids.push((msg.id.clone(), new_key.clone()));
                }
            }

            data.insert(new_key, new_value);
            data.remove(&key);
        }

        // Move DLQ entries
        let dlq_prefix_from = format!("_dlq/{}/", from);
        let dlq_entries: Vec<(String, Vec<u8>)> = data
            .range(dlq_prefix_from.clone()..)
            .take_while(|(k, _)| k.starts_with(&dlq_prefix_from))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();

        for (key, value) in dlq_entries {
            let suffix = key
                .strip_prefix(&dlq_prefix_from)
                .unwrap_or_default()
                .to_string();
            let mut msg: Message = serde_json::from_slice(&value).map_err(|e| {
                RenameQueueError::Storage(anyhow::anyhow!("deserialize DLQ message failed: {e}"))
            })?;
            msg.queue = to.to_string();
            let new_value = serde_json::to_vec(&msg).map_err(|e| {
                RenameQueueError::Storage(anyhow::anyhow!("serialize DLQ message failed: {e}"))
            })?;

            let new_key = format!("_dlq/{}/{}", to, suffix);
            data.insert(new_key, new_value);
            data.remove(&key);
        }

        // Update config keys
        let config_key_from = format!("_queue_config/{}", from);
        let config_key_to = format!("_queue_config/{}", to);
        let config_json = serde_json::to_vec(&existing_config).map_err(|e| {
            RenameQueueError::Storage(anyhow::anyhow!("serialize config failed: {e}"))
        })?;
        data.remove(&config_key_from);
        data.insert(config_key_to, config_json);

        // Drop data lock before acquiring other locks
        drop(data);

        // Update in-memory configs
        {
            let mut configs = self.queue_configs.write().await;
            configs.remove(from);
            configs.insert(to.to_string(), existing_config);
        }

        // Update locked_index: remove old keys by exact match and insert
        // new keys.  Precisely targeted instead of a retain() pass so
        // rename cost stays proportional to the source queue's locked
        // set, not the total locked set across all queues.
        {
            let mut index = self.locked_index.write().await;
            for (old_key, new_key, until) in moved_locked {
                index.remove(&old_key);
                index.insert(new_key, until);
            }
        }

        // DLV-0014: rebuild the secondary index under the new queue
        // name.  Drop the entire `from` sub-map and insert moved
        // (msg_id, new_storage_key) pairs into the `to` sub-map.
        {
            let mut id_index = self.locked_id_index.write().await;
            id_index.remove(from);
            if !moved_locked_ids.is_empty() {
                let entry = id_index.entry(to.to_string()).or_default();
                for (msg_id, new_key) in moved_locked_ids {
                    entry.insert(msg_id, new_key);
                }
            }
        }

        // Move payload set
        {
            let mut sets = self.payload_sets.write().await;
            if let Some(set) = sets.remove(from) {
                sets.insert(to.to_string(), set);
            }
        }

        // SYS-0018: move counter entry to the new name.
        {
            let mut counters = self.queue_counters.lock().unwrap();
            if let Some(entry) = counters.remove(from) {
                counters.insert(to.to_string(), entry);
            }
        }

        Ok(())
    }

    async fn push(&self, msg: Message) -> Result<String> {
        // Implements: DLV-0012 (lock-order: data → payload_sets)
        let config = self.get_queue_config(&msg.queue).await;

        // Enforce priority kind matches queue configuration
        if msg.priority.kind() != config.priority_kind {
            return Err(anyhow::anyhow!(
                "Priority kind mismatch: queue expects {:?} but got {:?}",
                config.priority_kind,
                msg.priority.kind()
            ));
        }

        // Validate text priority constraints
        if let Priority::Text(ref s) = msg.priority {
            if s.is_empty() {
                return Err(anyhow::anyhow!("Text priority must not be empty"));
            }
        }

        if !config.allow_duplicates {
            // Acquire data first, then payload_sets — consistent with ack/nack/batch
            // (DLV-0012).  Reversed order (payload_sets → data) caused lock-order
            // inversions and deadlocks under concurrent producer+consumer load.
            let payload_hash = hash_payload(&msg.payload);
            let mut data = self.data.write().await;
            let mut sets = self.payload_sets.write().await;
            let set = sets.entry(msg.queue.clone()).or_default();
            if set.contains(&payload_hash) {
                return Err(anyhow::anyhow!("Duplicate payload rejected"));
            }
            let key = self.generate_message_key(&msg, &config);
            let value = serde_json::to_vec(&msg)?;
            data.insert(key, value);
            set.insert(payload_hash);

            // SYS-0018: available++
            self.queue_counters
                .lock()
                .unwrap()
                .entry(msg.queue.clone())
                .or_default()
                .available += 1;

            return Ok(msg.id);
        }

        let key = self.generate_message_key(&msg, &config);
        let value = serde_json::to_vec(&msg)?;
        self.data.write().await.insert(key, value);

        // SYS-0018: available++
        self.queue_counters
            .lock()
            .unwrap()
            .entry(msg.queue.clone())
            .or_default()
            .available += 1;

        Ok(msg.id)
    }

    async fn pop(
        &self,
        queue: &str,
        consumer_id: &str,
        timeout_secs: u64,
    ) -> Result<Option<Message>> {
        let prefix = format!("{}/", queue);
        let now = Utc::now();
        let lock_until = now + chrono::Duration::seconds(timeout_secs as i64);

        let mut data = self.data.write().await;

        // Find the first available message.  Capture whether the
        // selected message was already in the locked state (expired
        // lock) so we can avoid double-counting in the counter update
        // below — see SYS-0018 comment in pop's counter block.
        let found = {
            let mut result = None;
            for (k, v) in data.range(prefix.clone()..) {
                if !k.starts_with(&prefix) {
                    break;
                }
                let msg: Message = serde_json::from_slice(v)?;
                let was_already_locked = if let Some(locked_until) = msg.locked_until {
                    if locked_until > now {
                        continue;
                    }
                    true
                } else {
                    false
                };
                result = Some((k.clone(), msg, was_already_locked));
                break;
            }
            result
        };

        if let Some((key, mut msg, was_already_locked)) = found {
            msg.locked_until = Some(lock_until);
            msg.locked_by = Some(consumer_id.to_string());
            msg.retry_count += 1;

            let new_value = serde_json::to_vec(&msg)?;
            data.insert(key.clone(), new_value);

            self.locked_index
                .write()
                .await
                .insert(key.clone(), lock_until);

            // DLV-0014: populate the secondary index for O(1) lookup
            // by (queue, message_id) in ack/nack/renew/batch ops.
            self.locked_id_index
                .write()
                .await
                .entry(queue.to_string())
                .or_default()
                .insert(msg.id.clone(), key);

            // SYS-0018: counter update.
            //
            // Fresh pop (was_already_locked=false): available -> locked.
            //
            // Re-pop of expired lock (was_already_locked=true): the
            // message is already on the locked count from the prior
            // pop; do not double-count.  Without this guard the counter
            // drifts by +1 locked per uncleared expiry, which manifests
            // as "queues with a few messages that never get processed."
            if !was_already_locked {
                let mut counters = self.queue_counters.lock().unwrap();
                let entry = counters.entry(queue.to_string()).or_default();
                entry.available = entry.available.saturating_sub(1);
                entry.locked += 1;
            }

            Ok(Some(msg))
        } else {
            Ok(None)
        }
    }

    async fn ack(&self, queue: &str, message_id: &str, consumer_id: &str) -> Result<bool> {
        // Fast path (DLV-0014): O(1) lookup via the secondary index.
        let key = match self.locked_key_for_inner(queue, message_id).await {
            Some(k) => k,
            None => return Ok(false),
        };

        let msg = {
            let mut data = self.data.write().await;
            let value = match data.get(&key) {
                Some(v) => v.clone(),
                None => {
                    // Stale secondary-index entry — clean up and bail.
                    drop(data);
                    self.remove_locked_id_index_entry(queue, message_id).await;
                    return Ok(false);
                }
            };
            let msg: Message = serde_json::from_slice(&value)?;
            if msg.locked_by.as_deref() != Some(consumer_id) {
                return Ok(false);
            }
            data.remove(&key);
            msg
        };

        self.locked_index.write().await.remove(&key);
        self.remove_locked_id_index_entry(queue, message_id).await;

        {
            let mut sets = self.payload_sets.write().await;
            if let Some(set) = sets.get_mut(queue) {
                set.remove(&hash_payload(&msg.payload));
            }
        }

        // SYS-0018: locked--
        {
            let mut counters = self.queue_counters.lock().unwrap();
            if let Some(entry) = counters.get_mut(queue) {
                entry.locked = entry.locked.saturating_sub(1);
            }
        }

        Ok(true)
    }

    async fn nack(&self, queue: &str, message_id: &str, consumer_id: &str) -> Result<bool> {
        // Fast path (DLV-0014): O(1) lookup via the secondary index.
        let key = match self.locked_key_for_inner(queue, message_id).await {
            Some(k) => k,
            None => return Ok(false),
        };

        let (msg, moved_to_dlq) = {
            let mut data = self.data.write().await;
            let value = match data.get(&key) {
                Some(v) => v.clone(),
                None => {
                    drop(data);
                    self.remove_locked_id_index_entry(queue, message_id).await;
                    return Ok(false);
                }
            };
            let mut msg: Message = serde_json::from_slice(&value)?;
            if msg.locked_by.as_deref() != Some(consumer_id) {
                return Ok(false);
            }

            let moved_to_dlq = msg.retry_count >= msg.max_retries;
            if moved_to_dlq {
                let dlq_key = format!("_dlq/{}/{}", queue, msg.id);
                data.insert(dlq_key, value);
                data.remove(&key);
            } else {
                msg.locked_until = None;
                msg.locked_by = None;
                let new_value = serde_json::to_vec(&msg)?;
                data.insert(key.clone(), new_value);
            }
            (msg, moved_to_dlq)
        };

        self.locked_index.write().await.remove(&key);
        self.remove_locked_id_index_entry(queue, message_id).await;

        if moved_to_dlq {
            let mut sets = self.payload_sets.write().await;
            if let Some(set) = sets.get_mut(queue) {
                set.remove(&hash_payload(&msg.payload));
            }
        }

        // SYS-0018: locked--; if retry, available++
        {
            let mut counters = self.queue_counters.lock().unwrap();
            if let Some(entry) = counters.get_mut(queue) {
                entry.locked = entry.locked.saturating_sub(1);
                if !moved_to_dlq {
                    entry.available += 1;
                }
            }
        }

        Ok(true)
    }

    async fn renew(
        &self,
        queue: &str,
        message_id: &str,
        consumer_id: &str,
        timeout_secs: u64,
    ) -> Result<bool> {
        let now = Utc::now();
        let new_expiry = now + chrono::Duration::seconds(timeout_secs as i64);

        // Fast path (DLV-0014): O(1) lookup via the secondary index.
        let key = match self.locked_key_for_inner(queue, message_id).await {
            Some(k) => k,
            None => return Ok(false),
        };

        {
            let mut data = self.data.write().await;
            let value = match data.get(&key) {
                Some(v) => v.clone(),
                None => {
                    drop(data);
                    self.remove_locked_id_index_entry(queue, message_id).await;
                    return Ok(false);
                }
            };
            let mut msg: Message = serde_json::from_slice(&value)?;
            if msg.locked_by.as_deref() != Some(consumer_id) {
                return Ok(false);
            }
            msg.locked_until = Some(new_expiry);
            data.insert(key.clone(), serde_json::to_vec(&msg)?);
        }

        self.locked_index.write().await.insert(key, new_expiry);

        Ok(true)
    }

    async fn batch_ack(
        &self,
        queue: &str,
        consumer_id: &str,
        message_ids: &[String],
    ) -> Result<BatchAckResult> {
        if message_ids.is_empty() {
            return Ok(BatchAckResult::default());
        }

        // Fast path (DLV-0014): resolve each id to its storage key via
        // the secondary index in O(batch_size) instead of scanning the
        // whole queue.  MemoryStorage has no cap or sticky flag, so a
        // missing entry provably means the message is not locked.
        let snapshot: Vec<(String, String)> = {
            let idx = self.locked_id_index.read().await;
            match idx.get(queue) {
                Some(map) => message_ids
                    .iter()
                    .filter_map(|id| map.get(id).map(|key| (id.clone(), key.clone())))
                    .collect(),
                None => Vec::new(),
            }
        };

        let mut acked: Vec<String> = Vec::new();
        let mut acked_payloads: Vec<[u8; 16]> = Vec::new();
        let mut keys_to_remove: Vec<String> = Vec::new();
        let mut stale: Vec<String> = Vec::new();

        {
            let mut data = self.data.write().await;
            for (id, key) in snapshot {
                let value = match data.get(&key) {
                    Some(v) => v.clone(),
                    None => {
                        stale.push(id);
                        continue;
                    }
                };
                let msg: Message = serde_json::from_slice(&value)?;
                if msg.locked_by.as_deref() == Some(consumer_id) {
                    data.remove(&key);
                    acked.push(id);
                    acked_payloads.push(hash_payload(&msg.payload));
                    keys_to_remove.push(key);
                }
            }
        }

        // Clean up stale secondary-index entries.
        for id in stale {
            self.remove_locked_id_index_entry(queue, &id).await;
        }

        if !keys_to_remove.is_empty() {
            let mut locked_index = self.locked_index.write().await;
            for key in &keys_to_remove {
                locked_index.remove(key);
            }
        }

        // DLV-0014: drop acked entries from the secondary index.
        if !acked.is_empty() {
            let mut idx = self.locked_id_index.write().await;
            if let Some(map) = idx.get_mut(queue) {
                for id in &acked {
                    map.remove(id);
                }
                if map.is_empty() {
                    idx.remove(queue);
                }
            }
        }

        if !acked_payloads.is_empty() {
            let mut sets = self.payload_sets.write().await;
            if let Some(set) = sets.get_mut(queue) {
                for payload_hash in &acked_payloads {
                    set.remove(payload_hash);
                }
            }
        }

        // SYS-0018: locked -= acked_count
        if !acked.is_empty() {
            let mut counters = self.queue_counters.lock().unwrap();
            if let Some(entry) = counters.get_mut(queue) {
                entry.locked = entry.locked.saturating_sub(acked.len() as u64);
            }
        }

        let acked_set: HashSet<&str> = acked.iter().map(String::as_str).collect();
        let not_found = message_ids
            .iter()
            .filter(|id| !acked_set.contains(id.as_str()))
            .cloned()
            .collect();

        Ok(BatchAckResult { acked, not_found })
    }

    async fn batch_nack(
        &self,
        queue: &str,
        consumer_id: &str,
        message_ids: &[String],
    ) -> Result<BatchNackResult> {
        if message_ids.is_empty() {
            return Ok(BatchNackResult::default());
        }

        // Fast path (DLV-0014): resolve each id via the secondary index.
        let snapshot: Vec<(String, String)> = {
            let idx = self.locked_id_index.read().await;
            match idx.get(queue) {
                Some(map) => message_ids
                    .iter()
                    .filter_map(|id| map.get(id).map(|key| (id.clone(), key.clone())))
                    .collect(),
                None => Vec::new(),
            }
        };

        struct Found {
            id: String,
            key: String,
            msg: Message,
        }

        let mut targets: Vec<Found> = Vec::new();
        let mut stale: Vec<String> = Vec::new();

        {
            let data = self.data.read().await;
            for (id, key) in snapshot {
                let value = match data.get(&key) {
                    Some(v) => v.clone(),
                    None => {
                        stale.push(id);
                        continue;
                    }
                };
                let msg: Message = serde_json::from_slice(&value)?;
                if msg.locked_by.as_deref() == Some(consumer_id) {
                    targets.push(Found { id, key, msg });
                }
            }
        }

        for id in stale {
            self.remove_locked_id_index_entry(queue, &id).await;
        }

        let mut result = BatchNackResult::default();
        let mut index_keys: Vec<String> = Vec::new();
        let mut processed_ids: Vec<String> = Vec::new();
        let mut dlq_payloads: Vec<[u8; 16]> = Vec::new();

        {
            let mut data = self.data.write().await;
            for Found { id, key, mut msg } in targets {
                if msg.retry_count >= msg.max_retries {
                    let dlq_key = format!("_dlq/{}/{}", queue, msg.id);
                    let value = data.get(&key).cloned().unwrap_or_default();
                    data.insert(dlq_key, value);
                    data.remove(&key);
                    result.dead_lettered.push(msg.id.clone());
                    dlq_payloads.push(hash_payload(&msg.payload));
                } else {
                    msg.locked_until = None;
                    msg.locked_by = None;
                    let new_value = serde_json::to_vec(&msg)?;
                    data.insert(key.clone(), new_value);
                    result.unlocked.push(msg.id.clone());
                }
                index_keys.push(key);
                processed_ids.push(id);
            }
        }

        let processed: HashSet<&str> = result
            .unlocked
            .iter()
            .chain(result.dead_lettered.iter())
            .chain(result.dropped.iter())
            .map(String::as_str)
            .collect();
        result.not_found = message_ids
            .iter()
            .filter(|id| !processed.contains(id.as_str()))
            .cloned()
            .collect();

        if !index_keys.is_empty() {
            let mut locked_index = self.locked_index.write().await;
            for key in index_keys {
                locked_index.remove(&key);
            }
        }

        // DLV-0014: drop every processed entry from the secondary index.
        if !processed_ids.is_empty() {
            let mut idx = self.locked_id_index.write().await;
            if let Some(map) = idx.get_mut(queue) {
                for id in &processed_ids {
                    map.remove(id);
                }
                if map.is_empty() {
                    idx.remove(queue);
                }
            }
        }

        if !dlq_payloads.is_empty() {
            let mut sets = self.payload_sets.write().await;
            if let Some(set) = sets.get_mut(queue) {
                for payload_hash in dlq_payloads {
                    set.remove(&payload_hash);
                }
            }
        }

        // SYS-0018: locked -= (unlocked + dead_lettered); available += unlocked
        {
            let total_processed = result.unlocked.len() + result.dead_lettered.len();
            if total_processed > 0 {
                let mut counters = self.queue_counters.lock().unwrap();
                if let Some(entry) = counters.get_mut(queue) {
                    entry.locked = entry.locked.saturating_sub(total_processed as u64);
                    entry.available += result.unlocked.len() as u64;
                }
            }
        }

        Ok(result)
    }

    async fn delete_queue(&self, queue_name: &str) -> Result<usize> {
        let mut data = self.data.write().await;
        let prefix = format!("{}/", queue_name);

        // Collect message keys to delete
        let keys: Vec<String> = data
            .range(prefix.clone()..)
            .take_while(|(k, _)| k.starts_with(&prefix))
            .map(|(k, _)| k.clone())
            .collect();

        let deleted_count = keys.len();
        for key in &keys {
            data.remove(key);
        }

        // Delete dead letter queue entries for this queue
        let dlq_prefix = format!("_dlq/{}/", queue_name);
        let dlq_keys: Vec<String> = data
            .range(dlq_prefix.clone()..)
            .take_while(|(k, _)| k.starts_with(&dlq_prefix))
            .map(|(k, _)| k.clone())
            .collect();
        for key in &dlq_keys {
            data.remove(key);
        }

        // Remove config key
        let config_key = format!("_queue_config/{}", queue_name);
        data.remove(&config_key);

        drop(data);

        // Clean up locked index
        if deleted_count > 0 {
            let mut locked_index = self.locked_index.write().await;
            for key in &keys {
                locked_index.remove(key);
            }
        }

        // Remove in-memory config and payload set
        self.queue_configs.write().await.remove(queue_name);
        self.payload_sets.write().await.remove(queue_name);

        // DLV-0014: drop the queue's sub-map from the secondary index.
        self.locked_id_index.write().await.remove(queue_name);

        // SYS-0018: drop the counter entry.
        self.queue_counters.lock().unwrap().remove(queue_name);

        Ok(deleted_count)
    }

    async fn purge_queue(&self, queue_name: &str) -> Result<usize> {
        let mut data = self.data.write().await;
        let prefix = format!("{}/", queue_name);

        // Collect message keys to delete (not config keys)
        let keys: Vec<String> = data
            .range(prefix.clone()..)
            .take_while(|(k, _)| k.starts_with(&prefix))
            .map(|(k, _)| k.clone())
            .collect();

        let purged_count = keys.len();
        for key in &keys {
            data.remove(key);
        }

        drop(data);

        // Clean up locked index
        if purged_count > 0 {
            let mut locked_index = self.locked_index.write().await;
            for key in &keys {
                locked_index.remove(key);
            }
        }

        // Clear payload set (queue still exists, just empty)
        let mut sets = self.payload_sets.write().await;
        if let Some(set) = sets.get_mut(queue_name) {
            set.clear();
        }

        // DLV-0014: drop the queue's sub-map from the secondary index.
        self.locked_id_index.write().await.remove(queue_name);

        // SYS-0018: queue still exists — zero the counter instead of removing.
        if let Some(entry) = self.queue_counters.lock().unwrap().get_mut(queue_name) {
            entry.available = 0;
            entry.locked = 0;
        }

        Ok(purged_count)
    }

    // Implements: SYS-0018
    async fn get_all_queue_stats(&self) -> Result<Vec<QueueStats>> {
        // Snapshot counters under a short-lived lock, release before the
        // async config read.  SYS-0018 mandates this path does not scan
        // `data`.
        let snapshot: Vec<(String, QueueCounts)> = {
            let counters = self.queue_counters.lock().unwrap();
            counters
                .iter()
                .map(|(k, v)| (k.clone(), v.clone()))
                .collect()
        };

        let configs = self.queue_configs.read().await;

        let mut stats: Vec<QueueStats> = snapshot
            .into_iter()
            .map(|(queue_name, counts)| {
                let config = configs.get(&queue_name).cloned().unwrap_or_default();
                QueueStats {
                    name: queue_name,
                    available: counts.available as usize,
                    locked: counts.locked as usize,
                    total: (counts.available + counts.locked) as usize,
                    config,
                }
            })
            .collect();

        stats.sort_by(|a, b| a.name.cmp(&b.name));
        Ok(stats)
    }

    // Implements: SYS-0018
    async fn get_queue_stats(&self, queue_name: &str) -> Result<QueueStats> {
        let counts = self
            .queue_counters
            .lock()
            .unwrap()
            .get(queue_name)
            .cloned()
            .unwrap_or_default();

        Ok(QueueStats {
            name: queue_name.to_string(),
            available: counts.available as usize,
            locked: counts.locked as usize,
            total: (counts.available + counts.locked) as usize,
            config: self.get_queue_config(queue_name).await,
        })
    }

    // Implements: SYS-0018
    async fn list_queues(&self) -> Result<Vec<String>> {
        let mut queues: Vec<String> = self
            .queue_counters
            .lock()
            .unwrap()
            .keys()
            .cloned()
            .collect();
        queues.sort();
        Ok(queues)
    }

    async fn unlock_expired_messages(&self) -> Result<usize> {
        let now = Utc::now();
        let mut expired_keys = Vec::new();

        {
            let locked_index = self.locked_index.read().await;
            for (key, lock_until) in locked_index.iter() {
                if *lock_until <= now {
                    expired_keys.push(key.clone());
                }
            }
        }

        let mut unlocked_count = 0;
        // SYS-0018: track per-queue transitions so we can update counters
        // after the unlock loop without holding multiple locks.
        let mut per_queue_unlocks: HashMap<String, u64> = HashMap::new();

        for key in &expired_keys {
            match self.unlock_message_by_key(key).await {
                Ok(true) => {
                    unlocked_count += 1;
                    if let Some(q) = Self::queue_from_key(key) {
                        *per_queue_unlocks.entry(q.to_string()).or_insert(0) += 1;
                    }
                }
                Ok(false) => {}
                Err(e) => {
                    tracing::error!("Failed to unlock message {}: {}", key, e);
                }
            }
        }

        if !expired_keys.is_empty() {
            let mut locked_index = self.locked_index.write().await;
            for key in expired_keys {
                locked_index.remove(&key);
            }
        }

        // SYS-0018: locked -= N, available += N per queue.
        if !per_queue_unlocks.is_empty() {
            let mut counters = self.queue_counters.lock().unwrap();
            for (queue, n) in per_queue_unlocks {
                if let Some(entry) = counters.get_mut(&queue) {
                    entry.locked = entry.locked.saturating_sub(n);
                    entry.available += n;
                }
            }
        }

        Ok(unlocked_count)
    }

    // Implements: API-0014
    async fn force_unlock_queue(&self, queue_name: &str) -> Result<usize> {
        let prefix = format!("{}/", queue_name);

        // Collect keys of locked messages on this queue.
        let keys_to_unlock: Vec<String> = {
            let data = self.data.read().await;
            data.range(prefix.clone()..)
                .take_while(|(k, _)| k.starts_with(&prefix))
                .filter_map(|(k, v)| {
                    serde_json::from_slice::<Message>(v)
                        .ok()
                        .filter(|m| m.locked_until.is_some())
                        .map(|_| k.clone())
                })
                .collect()
        };

        // Unlock each one via the existing helper (clears locked_until/locked_by
        // and writes back to the BTreeMap).
        let mut unlocked = 0usize;
        for key in &keys_to_unlock {
            if self.unlock_message_by_key(key).await.unwrap_or(false) {
                unlocked += 1;
            }
        }

        // Remove from locked_index.
        if !keys_to_unlock.is_empty() {
            let mut locked_index = self.locked_index.write().await;
            for key in &keys_to_unlock {
                locked_index.remove(key);
            }
        }

        // DLV-0014: defensively drop the queue's entire sub-map from
        // the secondary index.  `unlock_message_by_key` already removes
        // each unlocked message individually, but this guarantees
        // consistency for any entry that might have been skipped (e.g.
        // a message whose record failed to deserialize and therefore
        // never reached unlock_message_by_key).  Matches the
        // persistent backend's force_unlock_queue_sync behavior.
        self.locked_id_index.write().await.remove(queue_name);

        // SYS-0018: locked -= unlocked, available += unlocked.
        if unlocked > 0 {
            let mut counters = self.queue_counters.lock().unwrap();
            if let Some(entry) = counters.get_mut(queue_name) {
                entry.locked = entry.locked.saturating_sub(unlocked as u64);
                entry.available += unlocked as u64;
            }
        }

        Ok(unlocked)
    }

    async fn diagnose_queue(
        &self,
        queue_name: &str,
    ) -> Result<Option<crate::message::QueueDiagnostic>> {
        let prefix = format!("{}/", queue_name);
        let dlq_prefix = format!("_dlq/{}/", queue_name);
        let now = Utc::now();

        let mut scanned_total: u64 = 0;
        let mut scanned_available: u64 = 0;
        let mut scanned_locked_active: u64 = 0;
        let mut scanned_locked_expired: u64 = 0;
        let mut scanned_dlq: u64 = 0;

        {
            let data = self.data.read().await;
            for (k, v) in data.range(prefix.clone()..) {
                if !k.starts_with(&prefix) {
                    break;
                }
                if k.starts_with("_queue_config/") || k.starts_with("_dlq/") {
                    continue;
                }
                let msg: Message = match serde_json::from_slice(v) {
                    Ok(m) => m,
                    Err(_) => continue,
                };
                scanned_total += 1;
                match msg.locked_until {
                    Some(t) if t > now => scanned_locked_active += 1,
                    Some(_) => scanned_locked_expired += 1,
                    None => scanned_available += 1,
                }
            }
            for (k, _) in data.range(dlq_prefix.clone()..) {
                if !k.starts_with(&dlq_prefix) {
                    break;
                }
                scanned_dlq += 1;
            }
        }

        let (counter_available, counter_locked) = self
            .queue_counters
            .lock()
            .unwrap()
            .get(queue_name)
            .map(|c| (c.available, c.locked))
            .unwrap_or((0, 0));

        let locked_index_for_queue = self
            .locked_index
            .read()
            .await
            .keys()
            .filter(|k| k.starts_with(&prefix))
            .count() as u64;

        let locked_id_index_for_queue = self
            .locked_id_index
            .read()
            .await
            .get(queue_name)
            .map(|m| m.len() as u64)
            .unwrap_or(0);

        let expected_counter_available = scanned_available + scanned_locked_expired;
        let mut discrepancies = Vec::new();
        if counter_available != expected_counter_available {
            discrepancies.push(format!(
                "counter.available mismatch: counter={}, scan_available+scan_locked_expired={} (drift {:+})",
                counter_available,
                expected_counter_available,
                counter_available as i64 - expected_counter_available as i64
            ));
        }
        if counter_locked != scanned_locked_active {
            discrepancies.push(format!(
                "counter.locked mismatch: counter={}, scan_locked_active={} (drift {:+})",
                counter_locked,
                scanned_locked_active,
                counter_locked as i64 - scanned_locked_active as i64
            ));
        }
        if scanned_locked_expired > 0 {
            discrepancies.push(format!(
                "{} message(s) have expired locks not yet swept by unlock_expired_messages",
                scanned_locked_expired
            ));
        }

        Ok(Some(crate::message::QueueDiagnostic {
            queue: queue_name.to_string(),
            counter_available,
            counter_locked,
            scanned_total,
            scanned_available,
            scanned_locked_active,
            scanned_locked_expired,
            scanned_dlq,
            locked_index_for_queue,
            locked_id_index_for_queue,
            // MemoryStorage has no hot tier — there is no separate
            // cache layer between the BTreeMap and the consumer.
            hot_tier_size: 0,
            discrepancies,
        }))
    }

    async fn reseed_queue_counters(&self, queue_name: &str) -> Result<Option<(u64, u64)>> {
        let diag = match self.diagnose_queue(queue_name).await? {
            Some(d) => d,
            None => return Ok(None),
        };
        let new_available = diag.scanned_available + diag.scanned_locked_expired;
        let new_locked = diag.scanned_locked_active;
        let mut counters = self.queue_counters.lock().unwrap();
        let entry = counters.entry(queue_name.to_string()).or_default();
        entry.available = new_available;
        entry.locked = new_locked;
        tracing::info!(
            queue = queue_name,
            available = new_available,
            locked = new_locked,
            "reseed_queue_counters: counters reset from scan"
        );
        Ok(Some((new_available, new_locked)))
    }
}