rmqtt 0.22.0

//! MQTT Retained Message Storage Implementation
//!
//! Provides hierarchical storage and efficient retrieval of MQTT retained messages with:
//! 1. **Topic Pattern Matching**:
//!    - Multi-level wildcard (#) and single-level wildcard (+) support
//!    - Metadata-aware filtering for system topics (starting with $)
//!
//! 2. **Message Management**:
//!    - Time-based expiration with automatic cleanup
//!    - Atomic counters for message statistics
//!    - Thread-safe operations using async RwLock
//!
//! 3. **Core Components**:
//!    - `RetainTree`: Trie-like structure for O(log n) topic operations
//!    - `TimedValue`: Wrapper for message expiration tracking
//!    - `DefaultRetainStorage`: Production implementation with plugin integration
//!
//! ## Design Highlights
//! - **Hierarchical Storage**:
//!   ```text
//!   Root
//!   ├── iot
//!   │   └── b
//!   │       ├── x (value=1)
//!   │       ├── y (value=2)
//!   │       └── z (value=3)
//!   └── x
//!       └── y
//!           └── z (value=4)
//!   ```
//!   Enables efficient wildcard pattern matching through trie traversal
//!
//! - **Expiration Mechanism**:
//!   ```rust,ignore
//!   TimedValue::new(retain, timeout) // Wraps message with TTL
//!   remove_expired_messages() // Scheduled cleanup task
//!   ```
//!   Automatically evicts stale messages using duration-based tracking
//!
//! - **Plugin Architecture**:
//!   ```rust,ignore
//!   #[async_trait]
//!   impl RetainStorage for DefaultRetainStorage {
//!       async fn set(...) { /* delegates to rmqtt-retainer plugin */ }
//!   }
//!   ```
//!   Enables extension through external plugins while maintaining core logic
//!
//! ## Key Operations
//! | Method                | Complexity | Description                     |
//! |-----------------------|------------|---------------------------------|
//! | `insert()`            | O(k)       | k = topic depth levels          |
//! | `matches()`           | O(k+m)     | m = matching branches           |
//! | `remove_expired()`    | O(n)       | n = total stored messages       |
//! | `retain()`            | O(n)       | Conditional bulk removal        |
//!
//! ## Usage Note
//! The base implementation intentionally delegates to `rmqtt-retainer` plugin for:
//! - Distributed storage support
//! - Enhanced persistence mechanisms
//! - Cluster-wide message synchronization
//!
//! See `rmqtt-retainer` documentation for production deployment recommendations

use std::str::FromStr;
use std::time::Duration;

use async_trait::async_trait;
use tokio::sync::RwLock;

use crate::topic::{Level, Topic};
use crate::types::{HashMap, Retain, TimedValue, TopicFilter, TopicName};
use crate::utils::{Counter, StatsMergeMode};
use crate::Result;

/// Abstraction for retained message storage backends.
///
/// Implementations can provide in-memory, plugin-backed, or
/// distributed storage for retained messages.
///
/// # Default Behavior
///
/// The default in-memory implementation ([`DefaultRetainStorage`]) is
/// functional but logs a warning recommending the `rmqtt-retainer`
/// What type of retain synchronization a storage backend requires.
///
/// - `Full`: send the full retain message (`SetRetain`) — for local-only
///   backends like in-memory (Ram) or per-node Sled.
/// - `TopicOnly`: send only the topic name (`SetRetainTopic`) — for shared
///   backends like Redis where the retain data is already visible to all
///   nodes, but the in-memory topic index must be kept in sync.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RetainSyncMode {
    Full,
    TopicOnly,
}

/// plugin for production use.
#[async_trait]
pub trait RetainStorage: Sync + Send {
    /// Whether retained message storage is enabled.
    #[inline]
    fn enable(&self) -> bool {
        false
    }

    /// Whether merging retentions from remote cluster nodes is needed during retrieval.
    #[inline]
    fn merge_on_read(&self) -> bool {
        false
    }

    /// Whether retain storage needs cluster-wide synchronization at startup.
    ///
    /// Returns `true` (default) if the storage is local (in-memory, per-node Sled)
    /// and needs proactive sync of retains from peer nodes when a new node joins.
    /// Returns `false` if the storage is shared (e.g., Redis) so all nodes already
    /// see the same data and no startup sync is needed.
    #[inline]
    fn need_sync(&self) -> bool {
        true
    }

    /// Store a retained message for the given topic.
    ///
    /// If the payload is empty, the retained message is cleared.
    async fn set(&self, topic: &TopicName, retain: Retain, expiry_interval: Option<Duration>) -> Result<()>;

    /// Retrieve all retained messages matching the given topic filter.
    async fn get(&self, topic_filter: &TopicFilter) -> Result<Vec<(TopicName, Retain)>>;

    /// Retrieve a paginated snapshot of all non-expired retained messages.
    ///
    /// Returns `(items, has_more)`. Used by cluster plugins to synchronize
    /// the retain store between nodes. Default implementation returns empty.
    async fn get_all_paginated(
        &self,
        offset: usize,
        limit: usize,
    ) -> Result<(Vec<(TopicName, Retain, Option<Duration>)>, bool)> {
        let _ = (offset, limit);
        Ok((Vec::new(), false))
    }

    /// Current count of retained messages.
    async fn count(&self) -> isize;

    /// Maximum number of retained messages ever stored.
    async fn max(&self) -> isize;

    /// How stats from this storage should be merged with other sources.
    #[inline]
    fn stats_merge_mode(&self) -> StatsMergeMode {
        StatsMergeMode::None
    }

    /// What type of retain synchronization this storage requires from the
    /// cluster plugin when [`retain_set_broadcast`] is called.
    ///
    /// [`retain_set_broadcast`]: crate::shared::ClusterShared::retain_set_broadcast
    #[inline]
    fn retain_sync_mode(&self) -> RetainSyncMode {
        RetainSyncMode::Full
    }

    /// Handle a topic-only retain sync notification from a cluster peer.
    ///
    /// Called when the cluster plugin receives `SetRetainTopicAdd` or
    /// `SetRetainTopicRemove` — only used when [`retain_sync_mode`] returns
    /// `TopicOnly`. The storage backend should unconditionally insert or
    /// remove the topic from its in-memory index **without** querying the
    /// shared storage, since the actual retain data is already in Redis.
    ///
    /// - `is_set = true` : retain exists, insert into index
    /// - `is_set = false`: retain deleted, remove from index
    ///
    /// [`retain_sync_mode`]: Self::retain_sync_mode
    async fn sync_retain_topic(
        &self,
        _topic: &TopicName,
        _expiry_interval: Option<Duration>,
        _is_set: bool,
    ) -> Result<()> {
        Ok(())
    }
}

/// Default in-memory retained message storage.
///
/// Uses a [`RetainTree`] (trie-like structure) for O(log n) topic
/// operations and a [`Counter`] for atomic statistics tracking.
///
/// # Production Note
///
/// This is a simple in-memory implementation suitable for testing.
/// Production deployments should use the `rmqtt-retainer` plugin
/// for persistent and distributed storage.
pub struct DefaultRetainStorage {
    pub messages: RwLock<RetainTree<TimedValue<Retain>>>,
    retaineds: Counter,
}

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

impl DefaultRetainStorage {
    #[inline]
    pub fn new() -> DefaultRetainStorage {
        Self { messages: RwLock::new(RetainTree::default()), retaineds: Counter::new() }
    }

    #[inline]
    pub async fn remove_expired_messages(&self) -> usize {
        let mut messages = self.messages.write().await;
        messages.retain(usize::MAX, |tv| {
            if tv.is_expired() {
                self.retaineds.dec();
                false
            } else {
                true
            }
        })
    }

    #[inline]
    pub async fn set_with_timeout(
        &self,
        topic: &TopicName,
        retain: Retain,
        timeout: Option<Duration>,
    ) -> Result<()> {
        let topic = Topic::from_str(topic)?;
        let mut messages = self.messages.write().await;
        let old = messages.remove(&topic);
        if !retain.publish.payload.is_empty() {
            messages.insert(&topic, TimedValue::new(retain, timeout));
            if old.is_none() {
                self.retaineds.inc();
            }
        } else if old.is_some() {
            self.retaineds.dec();
        }
        Ok(())
    }

    #[inline]
    pub async fn get_message(&self, topic_filter: &TopicFilter) -> Result<Vec<(TopicName, Retain)>> {
        let topic = Topic::from_str(topic_filter)?;
        let retains = self
            .messages
            .read()
            .await
            .matches(&topic)
            .drain(..)
            .filter_map(|(t, r)| {
                if r.is_expired() {
                    None
                } else {
                    Some((TopicName::from(t.to_string()), r.into_value()))
                }
            })
            .collect::<Vec<(TopicName, Retain)>>();
        Ok(retains)
    }

    /// Return a paginated snapshot of all stored (non-expired) messages.
    ///
    /// Uses the wildcard `#` filter to match all topics.
    #[inline]
    pub async fn get_all_paginated(
        &self,
        offset: usize,
        limit: usize,
    ) -> Result<(Vec<(TopicName, Retain, Option<Duration>)>, bool)> {
        let topic = match Topic::from_str("#") {
            Err(e) => return Err(anyhow::anyhow!(e)),
            Ok(t) => t,
        };
        let messages = self.messages.read().await;
        let all: Vec<(TopicName, Retain, Option<Duration>)> = messages
            .matches(&topic)
            .into_iter()
            .filter_map(|(t, tv)| {
                if tv.is_expired() {
                    return None;
                }
                let topic = TopicName::from(t.to_string());
                let remaining = tv.remaining();
                let retain = tv.into_value();
                Some((topic, retain, remaining))
            })
            .collect();
        let total = all.len();
        let has_more = offset + limit < total;
        let items = all.into_iter().skip(offset).take(limit).collect();
        Ok((items, has_more))
    }
}

#[async_trait]
impl RetainStorage for DefaultRetainStorage {
    #[inline]
    async fn set(
        &self,
        _topic: &TopicName,
        _retain: Retain,
        _expiry_interval: Option<Duration>,
    ) -> Result<()> {
        log::warn!("Please use the \"rmqtt-retainer\" plugin as the main program no longer supports retain messages.");
        Ok(())
    }

    #[inline]
    async fn get(&self, _topic_filter: &TopicFilter) -> Result<Vec<(TopicName, Retain)>> {
        log::warn!("Please use the \"rmqtt-retainer\" plugin as the main program no longer supports retain messages.");
        Ok(Vec::new())
    }

    #[inline]
    async fn count(&self) -> isize {
        self.retaineds.count()
    }

    #[inline]
    async fn max(&self) -> isize {
        self.retaineds.max()
    }
}

/// Trie-like tree structure for retained message storage.
///
/// Supports efficient topic-level insertion, removal, and
/// wildcard pattern matching. The tree hierarchy mirrors
/// the MQTT topic hierarchy.
///
/// # Performance
///
/// | Operation    | Complexity |
/// |--------------|------------|
/// | `insert()`   | O(k)       |
/// | `remove()`   | O(k)       |
/// | `matches()`  | O(k+m)     |
/// | `retain()`   | O(n)       |
///
/// where k = topic depth, m = matching branches, n = total values.
pub type RetainTree<V> = Node<V>;

/// A single node in the retained message trie.
///
/// Each node optionally holds a value and maps topic levels
/// to child nodes, forming a tree that represents the topic hierarchy.
pub struct Node<V> {
    value: Option<V>,
    branches: HashMap<Level, Node<V>>,
}

impl<V> Default for Node<V> {
    #[inline]
    fn default() -> Node<V> {
        Self { value: None, branches: HashMap::default() }
    }
}

impl<V> Node<V>
where
    V: std::fmt::Debug + Clone,
{
    /// Insert a value at the given topic path, creating intermediate nodes as needed.
    #[inline]
    pub fn insert(&mut self, topic: &Topic, value: V) {
        let mut path = topic.levels().clone();
        path.reverse();
        self._insert(path, value);
    }

    #[inline]
    fn _insert(&mut self, mut path: Vec<Level>, value: V) {
        if let Some(first) = path.pop() {
            self.branches.entry(first).or_default()._insert(path, value)
        } else {
            self.value.replace(value);
        }
    }

    /// Remove the stored value at the given topic path.
    ///
    /// Prunes empty intermediate nodes after removal.
    /// Returns the removed value, if any.
    #[inline]
    pub fn remove(&mut self, topic: &Topic) -> Option<V> {
        self._remove(topic.levels().as_ref())
    }

    #[inline]
    fn _remove(&mut self, path: &[Level]) -> Option<V> {
        if path.is_empty() {
            self.value.take()
        } else {
            let t = &path[0];
            if let Some(x) = self.branches.get_mut(t) {
                let res = x._remove(&path[1..]);
                if x.value.is_none() && x.branches.is_empty() {
                    self.branches.remove(t);
                }
                res
            } else {
                None
            }
        }
    }

    /// Remove values for which the predicate returns `false`.
    ///
    /// Respects `max_limit` to cap the number of removals per call.
    /// Returns the count of removed entries.
    #[inline]
    pub fn retain<F>(&mut self, max_limit: usize, mut f: F) -> usize
    where
        F: FnMut(&mut V) -> bool,
    {
        let mut removeds = 0;
        self._retain(&mut f, &mut removeds, max_limit);
        removeds
    }

    #[inline]
    fn _retain<F>(&mut self, f: &mut F, removeds: &mut usize, max_limit: usize)
    where
        F: FnMut(&mut V) -> bool,
    {
        if *removeds >= max_limit {
            return;
        }
        self.branches.retain(|_, child_node| {
            child_node._retain(f, removeds, max_limit);
            if let Some(v) = child_node.value_mut() {
                if !f(v) {
                    let _ = child_node.value.take();
                    *removeds += 1;
                }
            }
            !(child_node.value.is_none() && child_node.branches.is_empty())
        });
    }

    #[inline]
    pub fn matches(&self, topic: &Topic) -> Vec<(Topic, V)> {
        let mut out = Vec::new();
        self._matches(topic.levels(), Vec::new(), &mut out);
        out
    }

    #[inline]
    fn _matches(&self, path: &[Level], mut sub_path: Vec<Level>, out: &mut Vec<(Topic, V)>) {
        let add_to_out = |levels: Vec<Level>, v: V, out: &mut Vec<(Topic, V)>| {
            out.push((Topic::from(levels), v));
        };

        //let node_map = &self.branches;

        if self.branches.is_empty() || path.is_empty() {
            if path.is_empty() {
                //Precise matching
                if let Some(v) = self.value.as_ref() {
                    add_to_out(sub_path, v.clone(), out);
                }
            }
        } else if !path.is_empty() {
            if let Some(r) = self.branches.get(&path[0]) {
                //Precise matching
                sub_path.push(path[0].clone());

                if path.len() > 1 && path[1] == Level::MultiWildcard {
                    //# Match parent, subscription ending with #
                    if let Some(v) = r.value.as_ref() {
                        add_to_out(sub_path.clone(), v.clone(), out);
                    }
                }
                r._matches(&path[1..], sub_path, out);
            } else if matches!(path[0], Level::SingleWildcard) {
                //Single layer matching
                for (k, v) in self.branches.iter() {
                    if sub_path.is_empty() && !matches!(k, Level::Blank) && k.is_metadata() {
                        //TopicName names starting with the $character cannot be matched with topic
                        //filters starting with wildcards (# or +)
                        continue;
                    }
                    let mut sub_path = sub_path.clone();
                    sub_path.push(k.clone());

                    if path.len() > 1 && path[1] == Level::MultiWildcard {
                        //# Match parent, subscription ending with #
                        if let Some(v) = v.value.as_ref() {
                            add_to_out(sub_path.clone(), v.clone(), out);
                        }
                    }
                    v._matches(&path[1..], sub_path, out);
                }
            } else if path[0] == Level::MultiWildcard {
                //Multilayer matching
                for (k, v) in self.branches.iter() {
                    if sub_path.is_empty() && !matches!(k, Level::Blank) && k.is_metadata() {
                        //TopicName names starting with the $character cannot be matched with topic
                        //filters starting with wildcards (# or +)
                        continue;
                    }
                    let mut sub_path = sub_path.clone();
                    sub_path.push(k.clone());

                    if v.branches.is_empty() {
                        if let Some(v) = v.value.as_ref() {
                            add_to_out(sub_path, v.clone(), out);
                        }
                    } else {
                        if let Some(v) = v.value.as_ref() {
                            add_to_out(sub_path.clone(), v.clone(), out);
                        }
                        v._matches(path, sub_path, out);
                    }
                }
            }
        }
    }

    #[inline]
    pub fn value(&self) -> Option<&V> {
        self.value.as_ref()
    }

    #[inline]
    pub fn value_mut(&mut self) -> Option<&mut V> {
        self.value.as_mut()
    }

    #[inline]
    pub fn children(&self) -> &HashMap<Level, Node<V>> {
        &self.branches
    }

    #[inline]
    pub fn child(&self, l: &Level) -> Option<&Node<V>> {
        self.branches.get(l)
    }

    #[inline]
    pub fn values_size(&self) -> usize {
        let len: usize = self.branches.values().map(|n| n.values_size()).sum();
        if self.value.is_some() {
            len + 1
        } else {
            len
        }
    }

    #[inline]
    pub fn nodes_size(&self) -> usize {
        let len: usize = self.branches.values().map(|n| n.nodes_size()).sum();
        self.branches.len() + len
    }

    #[inline]
    pub fn list(&self, mut top: usize) -> Vec<String> {
        let mut out = Vec::new();
        let parent = Level::Blank;
        self._list(&mut out, &parent, &mut top, 0);
        out
    }

    #[inline]
    fn _list(&self, out: &mut Vec<String>, _parent: &Level, top: &mut usize, depth: usize) {
        if *top == 0 {
            return;
        }
        for (l, n) in self.branches.iter() {
            out.push(format!("{} {:?}", " ".repeat(depth * 3), l));
            *top -= 1;
            n._list(out, l, top, depth + 1);
            if *top == 0 {
                return;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;

    use super::{RetainTree, Topic};

    fn match_one(tree: &RetainTree<i32>, topic_filter: &str, vs: &[i32]) -> bool {
        let mut matcheds = 0;
        let t = Topic::from_str(topic_filter).unwrap();
        //println!("[retain] {} ===> {:?}", topic_filter, tree.matches(&t));
        for (topic, v) in tree.matches(&t).iter() {
            println!("[retain] {topic_filter}({topic}) => {v:?}, {vs:?}");
            if !vs.contains(v) {
                return false;
            }
            matcheds += 1;
        }
        matcheds == vs.len()
    }

    #[test]
    fn retain() {
        let mut tree: RetainTree<i32> = RetainTree::default();
        tree.insert(&Topic::from_str("/iot/b/x").unwrap(), 1);
        tree.insert(&Topic::from_str("/iot/b/y").unwrap(), 2);
        tree.insert(&Topic::from_str("/iot/b/z").unwrap(), 3);
        tree.insert(&Topic::from_str("/iot/b").unwrap(), 123);
        tree.insert(&Topic::from_str("/x/y/z").unwrap(), 4);

        assert!(match_one(&tree, "/iot/b/y", &[2]));
        assert!(match_one(&tree, "/iot/b/+", &[1, 2, 3]));
        assert!(match_one(&tree, "/x/y/z", &[4]));
        assert!(!match_one(&tree, "/x/y/z", &[1]));

        tree.insert(&Topic::from_str("/xx/yy").unwrap(), -1);
        tree.insert(&Topic::from_str("/xx/yy/").unwrap(), 0);
        tree.insert(&Topic::from_str("/xx/yy/1").unwrap(), 1);
        tree.insert(&Topic::from_str("/xx/yy/2").unwrap(), 2);
        tree.insert(&Topic::from_str("/xx/yy/3").unwrap(), 3);

        tree.insert(&Topic::from_str("/xx/yy/3/4").unwrap(), 4);
        tree.insert(&Topic::from_str("/xx/yy/3/4/5").unwrap(), 5);

        assert!(match_one(&tree, "/xx/yy/+", &[0, 1, 2, 3]));
        assert!(match_one(&tree, "/xx/yy/3/+", &[4]));
        assert!(match_one(&tree, "/xx/yy/3/4/+", &[5]));
        assert!(match_one(&tree, "/xx/yy/1/+", &[]));

        println!("1 tree.values_size: {}", tree.values_size());
        println!("1 tree.nodes_size: {}", tree.nodes_size());
        tree.retain(usize::MAX, |_| false);
        println!("2 tree.values_size: {}", tree.values_size());
        println!("2 tree.nodes_size: {}", tree.nodes_size());
    }
}