offline_protocol_reliability/

retry_queue.rs

//! Retry queue with exponential backoff for message retries.

use crate::constants::{
    DEFAULT_BACKOFF_MULTIPLIER, DEFAULT_INITIAL_DELAY_MS, DEFAULT_MAX_DELAY_MS,
    DEFAULT_MAX_RETRIES, DEFAULT_OUTBOX_LIFETIME_MS, DEFAULT_PENDING_MESSAGE_LIFETIME_MS,
};
use chrono::{DateTime, Utc};
use offline_protocol_core::{Message, MessagePriority};
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};

/// Configuration for retry behavior.
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Maximum number of retries per message.
    pub max_retries: u32,

    /// Initial retry delay in milliseconds.
    pub initial_delay_ms: u64,

    /// Maximum retry delay in milliseconds.
    pub max_delay_ms: u64,

    /// Backoff multiplier for exponential backoff.
    pub backoff_multiplier: f32,

    /// Maximum lifetime for messages in the outbox (milliseconds).
    pub outbox_max_lifetime_ms: u64,

    /// Maximum lifetime for messages waiting on MLS session establishment.
    pub pending_message_max_lifetime_ms: u64,

    /// Opt-in: treat a durably-unreachable direct message as EDGE-DRIVEN once
    /// its escalating reachability probe has run a few times, instead of the
    /// default of probing it forever at the 15s->600s cap on every carrier.
    ///
    /// Default `false` preserves the documented contract (a parked message
    /// "never goes fully quiet"; see docs/message-delivery.md) and every
    /// existing native/third-party app's behaviour is unchanged. When `true`,
    /// after a bounded number of probes the message stops being timed-probed
    /// and rests in the outbox, re-driven only when the peer next proves
    /// reachable (an inbound frame or presence-online edge flushes it). It also
    /// lets a restart skip re-driving a durably-failing backlog.
    ///
    /// This trades the SDK's "self-recovers even for a silent returning peer
    /// with no presence" guarantee for zero steady-state relay traffic to gone
    /// peers. It is correct for deployments whose peers always interact or
    /// advertise presence on return (e.g. a machine-to-machine capability
    /// exchange), and MUST NOT be enabled for consumers that rely on the timed
    /// probe as their only recovery path.
    pub edge_driven_unreachable_dm: bool,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_retries: DEFAULT_MAX_RETRIES,
            initial_delay_ms: DEFAULT_INITIAL_DELAY_MS,
            max_delay_ms: DEFAULT_MAX_DELAY_MS,
            backoff_multiplier: DEFAULT_BACKOFF_MULTIPLIER,
            outbox_max_lifetime_ms: DEFAULT_OUTBOX_LIFETIME_MS,
            pending_message_max_lifetime_ms: DEFAULT_PENDING_MESSAGE_LIFETIME_MS,
            // Default off: preserve the documented perpetual-probe contract.
            edge_driven_unreachable_dm: false,
        }
    }
}

impl RetryConfig {
    /// Calculates the backoff delay for the given retry count:
    /// `initial_delay_ms * backoff_multiplier^retry_count`, clamped to
    /// `max_delay_ms`.
    ///
    /// Computed in f64 and clamped before the integer cast, so it is total
    /// over all configs: `f64::min` returns the other operand when one side
    /// is NaN (NaN/infinity collapse to `max_delay_ms`), and float-to-int
    /// `as` casts saturate.
    pub fn delay_for_retry(&self, retry_count: u32) -> u64 {
        let exponent = i32::try_from(retry_count).unwrap_or(i32::MAX);
        let factor = f64::from(self.backoff_multiplier).powi(exponent);
        let delay_ms = self.initial_delay_ms as f64 * factor;
        delay_ms.min(self.max_delay_ms as f64) as u64
    }
}

/// A message scheduled for retry.
#[derive(Debug, Clone)]
pub struct RetryEntry {
    /// The message to retry.
    pub message: Message,

    /// Number of times this message has been retried.
    pub retry_count: u32,

    /// When this message was first added to the queue.
    pub added_at: DateTime<Utc>,

    /// When this message should be retried next.
    pub retry_at: DateTime<Utc>,
}

impl RetryEntry {
    /// Checks if this entry is ready for retry.
    pub fn is_ready(&self) -> bool {
        Utc::now() >= self.retry_at
    }

    /// Checks if this entry has expired (exceeded max lifetime).
    pub fn is_expired(&self, max_lifetime_ms: u64) -> bool {
        let elapsed = Utc::now().signed_duration_since(self.added_at);
        elapsed.num_milliseconds() >= max_lifetime_ms as i64
    }
}

// Implement Ord for priority queue (earlier retry_at = higher priority)
impl Ord for RetryEntry {
    fn cmp(&self, other: &Self) -> Ordering {
        // Reverse comparison for min-heap (earliest retry_at first)
        other.retry_at.cmp(&self.retry_at).then_with(|| {
            // Tie-breaker: higher priority messages first (non-reversed,
            // so High > Low makes High pop first from the max-heap)
            self.message.priority.cmp(&other.message.priority)
        })
    }
}

impl PartialOrd for RetryEntry {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Eq for RetryEntry {}

impl PartialEq for RetryEntry {
    fn eq(&self, other: &Self) -> bool {
        self.retry_at == other.retry_at && self.message.priority == other.message.priority
    }
}

/// Retry queue for managing message retries with exponential backoff.
pub struct RetryQueue {
    config: RetryConfig,
    queue: BinaryHeap<RetryEntry>,
    index: HashMap<String, ()>, // For fast lookup
}

impl RetryQueue {
    /// Creates a new retry queue with default configuration.
    pub fn new() -> Self {
        Self::with_config(RetryConfig::default())
    }

    /// Creates a new retry queue with custom configuration.
    pub fn with_config(config: RetryConfig) -> Self {
        Self {
            config,
            queue: BinaryHeap::new(),
            index: HashMap::new(),
        }
    }

    /// Adds a message to the retry queue.
    ///
    /// The retry queue is a pure scheduling mechanism — it does not enforce a
    /// retry limit.  ACK-level retry limits are enforced by the ACK manager in
    /// the protocol layer.
    ///
    /// Duplicate messages (same ID already in queue) are silently ignored.
    ///
    /// # Arguments
    ///
    /// * `message` - The message to queue for retry
    /// * `retry_count` - Current retry count (used for backoff calculation)
    ///
    /// # Returns
    ///
    /// The absolute time the retry is scheduled for, or `None` if the
    /// message was already queued (nothing was scheduled).
    pub fn enqueue(&mut self, message: Message, retry_count: u32) -> Option<DateTime<Utc>> {
        let delay_ms = self.config.delay_for_retry(retry_count);
        self.enqueue_with_delay(message, retry_count, delay_ms)
    }

    /// Adds a message to the retry queue with an explicit delay, bypassing
    /// the backoff formula.
    ///
    /// Used for reachability probes toward peers an authoritative relay
    /// verdict reported unreachable, where the interval is driven by
    /// per-peer escalation rather than this message's retry count.
    ///
    /// Duplicate messages (same ID already in queue) are silently ignored;
    /// `None` means nothing was scheduled.
    pub fn enqueue_with_delay(
        &mut self,
        message: Message,
        retry_count: u32,
        delay_ms: u64,
    ) -> Option<DateTime<Utc>> {
        // Prevent duplicate entries for the same message
        if self.index.contains_key(&message.id.as_str()) {
            return None;
        }

        let retry_at =
            Utc::now() + chrono::Duration::milliseconds(delay_ms.min(i64::MAX as u64) as i64);

        let entry = RetryEntry {
            message: message.clone(),
            retry_count,
            added_at: Utc::now(),
            retry_at,
        };

        self.index.insert(message.id.as_str(), ());
        self.queue.push(entry);
        Some(retry_at)
    }

    /// Dequeues the next message that is ready for retry.
    ///
    /// # Returns
    ///
    /// Returns `Some(RetryEntry)` if a message is ready, `None` otherwise.
    pub fn dequeue_ready(&mut self) -> Option<RetryEntry> {
        // Peek at the top entry
        let is_ready = self.queue.peek().is_some_and(|entry| entry.is_ready());
        if is_ready {
            if let Some(entry) = self.queue.pop() {
                self.index.remove(&entry.message.id.as_str());
                return Some(entry);
            }
        }
        None
    }

    /// Gets all messages ready for retry without removing them.
    pub fn peek_ready(&self) -> Vec<&RetryEntry> {
        self.queue.iter().filter(|entry| entry.is_ready()).collect()
    }

    /// Removes a message from the queue (e.g., after successful delivery).
    pub fn remove(&mut self, message_id: &str) -> bool {
        if self.index.remove(message_id).is_some() {
            // Remove from heap (expensive, but necessary)
            self.queue
                .retain(|entry| entry.message.id.as_str() != message_id);
            true
        } else {
            false
        }
    }

    /// Cleans up expired messages that exceeded their max lifetime.
    ///
    /// # Returns
    ///
    /// Returns the number of expired messages removed.
    pub fn cleanup_expired(&mut self) -> usize {
        let max_lifetime = self.config.outbox_max_lifetime_ms;

        let before_count = self.queue.len();

        // Filter out expired entries
        let new_queue: BinaryHeap<RetryEntry> = self
            .queue
            .iter()
            .filter(|entry| !entry.is_expired(max_lifetime))
            .cloned()
            .collect();

        // Update index
        self.index.clear();
        for entry in &new_queue {
            self.index.insert(entry.message.id.as_str(), ());
        }

        self.queue = new_queue;

        before_count - self.queue.len()
    }

    /// Gets the number of messages in the queue.
    pub fn len(&self) -> usize {
        self.queue.len()
    }

    /// Checks if the queue is empty.
    pub fn is_empty(&self) -> bool {
        self.queue.is_empty()
    }

    /// Checks if a message is in the queue.
    pub fn contains(&self, message_id: &str) -> bool {
        self.index.contains_key(message_id)
    }

    /// Drains all entries from the queue, ignoring `retry_at` timing.
    ///
    /// Returns entries in priority order (earliest `retry_at` first, then
    /// highest message priority). Used for immediate flush when transport
    /// becomes available.
    pub fn drain_all(&mut self) -> Vec<RetryEntry> {
        self.index.clear();
        let mut entries = Vec::with_capacity(self.queue.len());
        while let Some(entry) = self.queue.pop() {
            entries.push(entry);
        }
        entries
    }

    /// Gets the time until the next retry is ready.
    ///
    /// # Returns
    ///
    /// Returns `Some(Duration)` if there are pending retries, `None` if empty.
    pub fn time_until_next_retry(&self) -> Option<chrono::Duration> {
        self.queue.peek().map(|entry| {
            let now = Utc::now();
            if entry.retry_at > now {
                entry.retry_at - now
            } else {
                chrono::Duration::zero()
            }
        })
    }

    /// Gets statistics about the retry queue.
    pub fn stats(&self) -> RetryQueueStats {
        let ready_count = self.peek_ready().len();
        let priority_counts = self.count_by_priority();

        RetryQueueStats {
            total_count: self.len(),
            ready_count,
            critical_priority_count: priority_counts[&MessagePriority::Critical],
            high_priority_count: priority_counts[&MessagePriority::High],
            medium_priority_count: priority_counts[&MessagePriority::Medium],
            low_priority_count: priority_counts[&MessagePriority::Low],
        }
    }

    fn count_by_priority(&self) -> HashMap<MessagePriority, usize> {
        let mut counts = HashMap::new();
        counts.insert(MessagePriority::Low, 0);
        counts.insert(MessagePriority::Medium, 0);
        counts.insert(MessagePriority::High, 0);
        counts.insert(MessagePriority::Critical, 0);

        for entry in &self.queue {
            *counts.entry(entry.message.priority).or_insert(0) += 1;
        }

        counts
    }
}

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

/// Statistics about the retry queue.
#[derive(Debug, Clone)]
pub struct RetryQueueStats {
    /// Total number of messages in queue.
    pub total_count: usize,
    /// Number of messages ready for retry.
    pub ready_count: usize,
    /// Number of critical priority messages.
    pub critical_priority_count: usize,
    /// Number of high priority messages.
    pub high_priority_count: usize,
    /// Number of medium priority messages.
    pub medium_priority_count: usize,
    /// Number of low priority messages.
    pub low_priority_count: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use offline_protocol_core::{AppId, UserId};
    use std::thread;
    use std::time::Duration;

    fn create_test_message(priority: MessagePriority) -> Message {
        Message::builder(
            UserId::new("alice").unwrap(),
            UserId::new("bob").unwrap(),
            AppId::new("test").unwrap(),
        )
        .content("Test message")
        .priority(priority)
        .build()
    }

    #[test]
    fn test_edge_driven_unreachable_dm_defaults_off() {
        // The opt-in flood-control behaviors (edge-driven parking, restart
        // age-gate, core resend rate-cap) are all gated on this flag. It MUST
        // default false so every existing native/third-party app keeps the
        // documented perpetual-probe behavior unchanged. If a future change
        // flips the default, this fails and flags the silent regression.
        assert!(!RetryConfig::default().edge_driven_unreachable_dm);
    }

    #[test]
    fn test_enqueue_and_dequeue() {
        let config = RetryConfig {
            initial_delay_ms: 50, // Fast for testing
            ..Default::default()
        };
        let mut queue = RetryQueue::with_config(config);

        let msg = create_test_message(MessagePriority::Medium);
        queue.enqueue(msg.clone(), 0);

        assert_eq!(queue.len(), 1);
        assert!(queue.contains(&msg.id.as_str()));

        // Not ready immediately
        assert!(queue.dequeue_ready().is_none());

        // Wait for retry time
        thread::sleep(Duration::from_millis(60));

        // Should be ready now
        let dequeued = queue.dequeue_ready().unwrap();
        assert_eq!(dequeued.message.id, msg.id);
        assert_eq!(queue.len(), 0);
    }

    #[test]
    fn test_exponential_backoff() {
        let config = RetryConfig {
            initial_delay_ms: 100,
            backoff_multiplier: 2.0,
            max_delay_ms: 1000,
            ..Default::default()
        };

        assert_eq!(config.delay_for_retry(0), 100);
        assert_eq!(config.delay_for_retry(1), 200);
        assert_eq!(config.delay_for_retry(2), 400);
        assert_eq!(config.delay_for_retry(3), 800);
        // 100 * 2^4 = 1600ms, capped at 1000ms
        assert_eq!(config.delay_for_retry(4), 1000);
        assert_eq!(config.delay_for_retry(100), 1000);
    }

    #[test]
    fn test_backoff_extreme_multiplier_clamps_not_overflows() {
        // Regression: the old formula multiplied in u64 before clamping, so a
        // large multiplier overflowed (debug panic / release garbage delay).
        for multiplier in [7.0, 1000.0, f32::MAX] {
            let config = RetryConfig {
                initial_delay_ms: 1000,
                backoff_multiplier: multiplier,
                max_delay_ms: 30_000,
                ..Default::default()
            };

            for retry_count in [1, 20, 1000, u32::MAX] {
                let delay = config.delay_for_retry(retry_count);
                assert!(
                    delay <= 30_000,
                    "multiplier {multiplier} retry {retry_count} gave delay {delay}"
                );
            }
        }
    }

    #[test]
    fn test_backoff_non_finite_multiplier_collapses_to_max_delay() {
        // Locks down the totality claim on delay_for_retry: NaN and infinity
        // collapse to exactly max_delay_ms, not merely something <= it.
        for multiplier in [f32::NAN, f32::INFINITY] {
            let config = RetryConfig {
                initial_delay_ms: 1000,
                backoff_multiplier: multiplier,
                max_delay_ms: 30_000,
                ..Default::default()
            };

            for retry_count in [1, 20, u32::MAX] {
                assert_eq!(
                    config.delay_for_retry(retry_count),
                    30_000,
                    "multiplier {multiplier} retry {retry_count}"
                );
            }
        }
    }

    #[test]
    fn test_backoff_retry_zero_clamped_to_max_delay() {
        // retry 0 no longer bypasses the max_delay_ms clamp: a misconfigured
        // initial_delay_ms > max_delay_ms yields max_delay_ms, not the raw
        // initial delay.
        let config = RetryConfig {
            initial_delay_ms: 60_000,
            backoff_multiplier: 2.0,
            max_delay_ms: 30_000,
            ..Default::default()
        };

        assert_eq!(config.delay_for_retry(0), 30_000);
    }

    #[test]
    fn test_enqueue_backoff_clamped_to_max_delay() {
        // Exercise the real enqueue path with a multiplier that overflowed
        // the old inline formula.
        let config = RetryConfig {
            initial_delay_ms: 1000,
            backoff_multiplier: 10.0,
            max_delay_ms: 30_000,
            ..Default::default()
        };
        let mut queue = RetryQueue::with_config(config);

        let msg = create_test_message(MessagePriority::Medium);
        queue.enqueue(msg, 20);

        let until_retry = queue.time_until_next_retry().unwrap();
        assert!(until_retry <= chrono::Duration::milliseconds(30_000));
        assert!(until_retry > chrono::Duration::zero());
    }

    #[test]
    fn test_enqueue_accepts_high_retry_count() {
        let config = RetryConfig {
            max_retries: 2,
            ..Default::default()
        };
        let mut queue = RetryQueue::with_config(config);

        let msg = create_test_message(MessagePriority::Medium);

        // enqueue no longer rejects based on max_retries — it's a pure
        // scheduling mechanism. ACK timeouts govern permanent failure.
        queue.enqueue(msg.clone(), 0);
        queue.dequeue_ready(); // Clear queue

        queue.enqueue(msg.clone(), 1);
        queue.dequeue_ready();

        // Previously this would fail; now it should succeed
        queue.enqueue(msg.clone(), 2);
        queue.dequeue_ready();

        queue.enqueue(msg.clone(), 100);
    }

    #[test]
    fn test_drain_all() {
        let config = RetryConfig {
            initial_delay_ms: 60_000, // Long delay so nothing is "ready"
            ..Default::default()
        };
        let mut queue = RetryQueue::with_config(config);

        let msg1 = create_test_message(MessagePriority::High);
        let msg2 = create_test_message(MessagePriority::Low);
        let msg3 = create_test_message(MessagePriority::Medium);

        queue.enqueue(msg1.clone(), 0);
        queue.enqueue(msg2.clone(), 0);
        queue.enqueue(msg3.clone(), 0);

        // Nothing should be ready (long delay)
        assert!(queue.dequeue_ready().is_none());
        assert_eq!(queue.len(), 3);

        // drain_all returns everything regardless of timing, in heap order
        // (earliest retry_at first, priority as tiebreaker)
        let entries = queue.drain_all();
        assert_eq!(entries.len(), 3);
        assert!(queue.is_empty());
        assert!(!queue.contains(&msg1.id.as_str()));
    }

    #[test]
    fn test_priority_ordering() {
        let config = RetryConfig {
            initial_delay_ms: 50,
            ..Default::default()
        };
        let mut queue = RetryQueue::with_config(config);

        // Add messages with different priorities (same retry time)
        let low = create_test_message(MessagePriority::Low);
        let high = create_test_message(MessagePriority::High);
        let medium = create_test_message(MessagePriority::Medium);

        queue.enqueue(low.clone(), 0);
        queue.enqueue(high.clone(), 0);
        queue.enqueue(medium.clone(), 0);

        thread::sleep(Duration::from_millis(60));

        // All should be ready
        let ready = queue.peek_ready();
        assert_eq!(ready.len(), 3);

        // Check that high priority is among the ready messages
        let has_high = ready
            .iter()
            .any(|e| e.message.priority == MessagePriority::High);
        assert!(has_high);
    }

    #[test]
    fn test_priority_tiebreaker_ordering() {
        // Verify that when retry_at is identical, higher priority pops first.
        // We construct entries directly to ensure identical timestamps.
        let now = Utc::now();
        let make_entry = |priority: MessagePriority| {
            let mut msg = create_test_message(priority);
            // Force same priority into message
            msg.priority = priority;
            RetryEntry {
                message: msg,
                retry_count: 0,
                added_at: now,
                retry_at: now,
            }
        };

        let mut heap = BinaryHeap::new();
        heap.push(make_entry(MessagePriority::Low));
        heap.push(make_entry(MessagePriority::High));
        heap.push(make_entry(MessagePriority::Medium));

        assert_eq!(heap.pop().unwrap().message.priority, MessagePriority::High);
        assert_eq!(
            heap.pop().unwrap().message.priority,
            MessagePriority::Medium
        );
        assert_eq!(heap.pop().unwrap().message.priority, MessagePriority::Low);
    }

    #[test]
    fn test_enqueue_with_delay_uses_explicit_delay() {
        // The backoff formula would schedule retry 0 at initial_delay_ms
        // (1s default); the explicit delay must win.
        let mut queue = RetryQueue::new();
        let msg = create_test_message(MessagePriority::Medium);

        let before = Utc::now();
        let retry_at = queue.enqueue_with_delay(msg, 0, 60_000).unwrap();

        let delay = retry_at - before;
        assert!(delay >= chrono::Duration::milliseconds(59_000));
        assert!(delay <= chrono::Duration::milliseconds(61_000));
    }

    #[test]
    fn test_enqueue_with_delay_ignores_duplicates() {
        let mut queue = RetryQueue::new();
        let msg = create_test_message(MessagePriority::Medium);

        assert!(queue.enqueue_with_delay(msg.clone(), 0, 1000).is_some());
        assert!(queue.enqueue_with_delay(msg.clone(), 0, 2000).is_none());
        assert!(queue.enqueue(msg, 0).is_none());
        assert_eq!(queue.len(), 1);
    }

    #[test]
    fn test_remove_message() {
        let mut queue = RetryQueue::new();
        let msg = create_test_message(MessagePriority::Medium);

        queue.enqueue(msg.clone(), 0);
        assert!(queue.contains(&msg.id.as_str()));

        assert!(queue.remove(&msg.id.as_str()));
        assert!(!queue.contains(&msg.id.as_str()));
        assert_eq!(queue.len(), 0);
    }

    #[test]
    fn test_cleanup_expired() {
        let config = RetryConfig {
            outbox_max_lifetime_ms: 100, // 100ms for testing
            ..Default::default()
        };
        let mut queue = RetryQueue::with_config(config);

        let msg = create_test_message(MessagePriority::Medium);
        queue.enqueue(msg, 0);

        // Wait for expiration
        thread::sleep(Duration::from_millis(150));

        let removed = queue.cleanup_expired();
        assert_eq!(removed, 1);
        assert_eq!(queue.len(), 0);
    }

    #[test]
    fn test_queue_stats() {
        let config = RetryConfig {
            initial_delay_ms: 50,
            ..Default::default()
        };
        let mut queue = RetryQueue::with_config(config);

        queue.enqueue(create_test_message(MessagePriority::High), 0);
        queue.enqueue(create_test_message(MessagePriority::High), 0);
        queue.enqueue(create_test_message(MessagePriority::Medium), 0);
        queue.enqueue(create_test_message(MessagePriority::Low), 0);

        let stats = queue.stats();
        assert_eq!(stats.total_count, 4);
        assert_eq!(stats.high_priority_count, 2);
        assert_eq!(stats.medium_priority_count, 1);
        assert_eq!(stats.low_priority_count, 1);
    }
}