peat_mesh/storage/

flow_control.rs

//! Flow control for Automerge sync
//!
//! This module provides production-grade flow control for sync operations:
//! - Rate limiting per peer (token bucket algorithm)
//! - Memory-bounded message queues
//! - Sync storm prevention (cooldown after rapid syncs)
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────────┐
//! │  SyncCoordinator    │
//! └──────────┬──────────┘
//!            │
//!            ▼
//! ┌─────────────────────┐
//! │   FlowController    │
//! ├─────────────────────┤
//! │ - TokenBucket/peer  │
//! │ - MessageQueue/peer │
//! │ - SyncCooldown/doc  │
//! └─────────────────────┘
//! ```

use iroh::EndpointId;
use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use thiserror::Error;

/// Flow control errors
#[derive(Error, Debug, Clone, PartialEq, Eq)]
pub enum FlowControlError {
    /// Rate limit exceeded for peer
    #[error("Rate limit exceeded for peer")]
    RateLimitExceeded,
    /// Message queue full for peer
    #[error("Message queue full for peer (max {max_size} messages)")]
    QueueFull { max_size: usize },
    /// Sync cooldown active (storm prevention)
    #[error("Sync cooldown active, {remaining_ms}ms remaining")]
    CooldownActive { remaining_ms: u64 },
}

/// Configuration for flow control
#[derive(Debug, Clone)]
pub struct FlowControlConfig {
    /// Maximum messages per second per peer (token bucket capacity)
    pub max_messages_per_second: u32,
    /// Token refill rate (tokens per refill interval)
    pub tokens_per_refill: u32,
    /// Token refill interval
    pub refill_interval: Duration,
    /// Maximum queue size per peer
    pub max_queue_size: usize,
    /// Sync cooldown period (minimum time between syncs for same doc)
    pub sync_cooldown: Duration,
    /// Maximum memory per peer for sync state (bytes)
    pub max_memory_per_peer: usize,
}

impl Default for FlowControlConfig {
    fn default() -> Self {
        Self {
            max_messages_per_second: 100,
            tokens_per_refill: 10,
            refill_interval: Duration::from_millis(100), // 10 refills/sec * 10 tokens = 100/sec
            max_queue_size: 1000,
            sync_cooldown: Duration::from_millis(100), // 100ms minimum between syncs
            max_memory_per_peer: 10 * 1024 * 1024,     // 10MB per peer
        }
    }
}

/// Token bucket rate limiter
///
/// Implements the token bucket algorithm for rate limiting:
/// - Bucket starts full (capacity tokens)
/// - Each operation consumes one token
/// - Tokens refill at a fixed rate
/// - Operations blocked when bucket is empty
#[derive(Debug)]
pub struct TokenBucket {
    /// Maximum tokens (bucket capacity)
    capacity: u32,
    /// Current available tokens
    tokens: AtomicU32,
    /// Tokens to add per refill
    tokens_per_refill: u32,
    /// Refill interval
    refill_interval: Duration,
    /// Last refill timestamp
    last_refill: RwLock<Instant>,
}

impl TokenBucket {
    /// Create a new token bucket
    pub fn new(capacity: u32, tokens_per_refill: u32, refill_interval: Duration) -> Self {
        Self {
            capacity,
            tokens: AtomicU32::new(capacity),
            tokens_per_refill,
            refill_interval,
            last_refill: RwLock::new(Instant::now()),
        }
    }

    /// Try to consume a token
    ///
    /// Returns true if a token was available and consumed, false otherwise.
    pub fn try_acquire(&self) -> bool {
        // First, refill tokens if needed
        self.refill();

        // Try to consume a token using CAS loop
        loop {
            let current = self.tokens.load(Ordering::Acquire);
            if current == 0 {
                return false;
            }
            if self
                .tokens
                .compare_exchange(current, current - 1, Ordering::AcqRel, Ordering::Acquire)
                .is_ok()
            {
                return true;
            }
            // CAS failed, retry
        }
    }

    /// Get current available tokens
    pub fn available_tokens(&self) -> u32 {
        self.refill();
        self.tokens.load(Ordering::Acquire)
    }

    /// Refill tokens based on elapsed time
    fn refill(&self) {
        let now = Instant::now();
        let mut last = self.last_refill.write().unwrap_or_else(|e| e.into_inner());

        let elapsed = now.duration_since(*last);
        if elapsed < self.refill_interval {
            return;
        }

        // Calculate how many refill periods have passed
        let periods = (elapsed.as_millis() / self.refill_interval.as_millis()) as u32;
        if periods == 0 {
            return;
        }

        // Add tokens (capped at capacity)
        let tokens_to_add = periods.saturating_mul(self.tokens_per_refill);
        loop {
            let current = self.tokens.load(Ordering::Acquire);
            let new_tokens = (current + tokens_to_add).min(self.capacity);
            if self
                .tokens
                .compare_exchange(current, new_tokens, Ordering::AcqRel, Ordering::Acquire)
                .is_ok()
            {
                break;
            }
        }

        // Update last refill time
        *last = now;
    }
}

/// Bounded message queue for a peer
///
/// Provides a FIFO queue with configurable maximum size.
/// When full, oldest messages are dropped (or operations fail).
#[derive(Debug)]
pub struct BoundedQueue<T> {
    /// Queue contents
    queue: VecDeque<T>,
    /// Maximum queue size
    max_size: usize,
    /// Total messages enqueued (all time)
    total_enqueued: u64,
    /// Total messages dropped due to overflow
    total_dropped: u64,
}

impl<T> BoundedQueue<T> {
    /// Create a new bounded queue
    pub fn new(max_size: usize) -> Self {
        Self {
            queue: VecDeque::with_capacity(max_size.min(1000)), // Pre-allocate up to 1000
            max_size,
            total_enqueued: 0,
            total_dropped: 0,
        }
    }

    /// Enqueue a message, dropping oldest if full
    ///
    /// Returns the dropped message if one was evicted, None otherwise.
    pub fn enqueue(&mut self, item: T) -> Option<T> {
        self.total_enqueued += 1;

        let dropped = if self.queue.len() >= self.max_size {
            self.total_dropped += 1;
            self.queue.pop_front()
        } else {
            None
        };

        self.queue.push_back(item);
        dropped
    }

    /// Try to enqueue a message, failing if full
    ///
    /// Returns Ok(()) if enqueued, Err if queue is full.
    pub fn try_enqueue(&mut self, item: T) -> Result<(), T> {
        if self.queue.len() >= self.max_size {
            return Err(item);
        }
        self.total_enqueued += 1;
        self.queue.push_back(item);
        Ok(())
    }

    /// Dequeue the next message
    pub fn dequeue(&mut self) -> Option<T> {
        self.queue.pop_front()
    }

    /// Peek at the next message without removing
    pub fn peek(&self) -> Option<&T> {
        self.queue.front()
    }

    /// Get current queue length
    pub fn len(&self) -> usize {
        self.queue.len()
    }

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

    /// Get total messages enqueued (all time)
    pub fn total_enqueued(&self) -> u64 {
        self.total_enqueued
    }

    /// Get total messages dropped due to overflow
    pub fn total_dropped(&self) -> u64 {
        self.total_dropped
    }

    /// Clear the queue
    pub fn clear(&mut self) {
        self.queue.clear();
    }
}

/// Sync cooldown tracker for storm prevention
///
/// Tracks last sync time for each (peer, document) pair to prevent
/// rapid repeated syncs that could indicate a sync storm.
#[derive(Debug)]
pub struct SyncCooldownTracker {
    /// Last sync time for (peer_id, doc_key) pairs
    last_sync: HashMap<(EndpointId, String), Instant>,
    /// Cooldown duration
    cooldown: Duration,
    /// Total syncs blocked by cooldown
    blocked_count: u64,
}

impl SyncCooldownTracker {
    /// Create a new cooldown tracker
    pub fn new(cooldown: Duration) -> Self {
        Self {
            last_sync: HashMap::new(),
            cooldown,
            blocked_count: 0,
        }
    }

    /// Check if sync is allowed (not in cooldown)
    ///
    /// Returns Ok(()) if sync is allowed, Err with remaining time if in cooldown.
    pub fn check_cooldown(
        &mut self,
        peer_id: &EndpointId,
        doc_key: &str,
    ) -> Result<(), FlowControlError> {
        let key = (*peer_id, doc_key.to_string());
        let now = Instant::now();

        if let Some(last) = self.last_sync.get(&key) {
            let elapsed = now.duration_since(*last);
            if elapsed < self.cooldown {
                self.blocked_count += 1;
                let remaining = self.cooldown - elapsed;
                return Err(FlowControlError::CooldownActive {
                    remaining_ms: remaining.as_millis() as u64,
                });
            }
        }

        Ok(())
    }

    /// Record a sync operation (updates last sync time)
    pub fn record_sync(&mut self, peer_id: &EndpointId, doc_key: &str) {
        let key = (*peer_id, doc_key.to_string());
        self.last_sync.insert(key, Instant::now());
    }

    /// Get count of syncs blocked by cooldown
    pub fn blocked_count(&self) -> u64 {
        self.blocked_count
    }

    /// Clean up old entries (entries older than 10x cooldown)
    pub fn cleanup(&mut self) {
        let now = Instant::now();
        let threshold = self.cooldown * 10;
        self.last_sync
            .retain(|_, last| now.duration_since(*last) < threshold);
    }
}

/// Per-peer resource tracking
#[derive(Debug)]
pub struct PeerResourceTracker {
    /// Estimated memory usage (bytes)
    memory_usage: AtomicU64,
    /// Maximum allowed memory (bytes)
    max_memory: u64,
    /// Messages sent
    messages_sent: AtomicU64,
    /// Messages received
    messages_received: AtomicU64,
    /// Messages dropped (rate limited or queue overflow)
    messages_dropped: AtomicU64,
}

impl PeerResourceTracker {
    /// Create a new resource tracker
    pub fn new(max_memory: u64) -> Self {
        Self {
            memory_usage: AtomicU64::new(0),
            max_memory,
            messages_sent: AtomicU64::new(0),
            messages_received: AtomicU64::new(0),
            messages_dropped: AtomicU64::new(0),
        }
    }

    /// Add to memory usage, returns false if would exceed limit
    pub fn try_allocate(&self, bytes: u64) -> bool {
        loop {
            let current = self.memory_usage.load(Ordering::Acquire);
            let new_usage = current + bytes;
            if new_usage > self.max_memory {
                return false;
            }
            if self
                .memory_usage
                .compare_exchange(current, new_usage, Ordering::AcqRel, Ordering::Acquire)
                .is_ok()
            {
                return true;
            }
        }
    }

    /// Free memory
    pub fn free(&self, bytes: u64) {
        self.memory_usage.fetch_sub(bytes, Ordering::Release);
    }

    /// Get current memory usage
    pub fn memory_usage(&self) -> u64 {
        self.memory_usage.load(Ordering::Acquire)
    }

    /// Record message sent
    pub fn record_sent(&self) {
        self.messages_sent.fetch_add(1, Ordering::Relaxed);
    }

    /// Record message received
    pub fn record_received(&self) {
        self.messages_received.fetch_add(1, Ordering::Relaxed);
    }

    /// Record message dropped
    pub fn record_dropped(&self) {
        self.messages_dropped.fetch_add(1, Ordering::Relaxed);
    }

    /// Get messages sent count
    pub fn messages_sent(&self) -> u64 {
        self.messages_sent.load(Ordering::Relaxed)
    }

    /// Get messages received count
    pub fn messages_received(&self) -> u64 {
        self.messages_received.load(Ordering::Relaxed)
    }

    /// Get messages dropped count
    pub fn messages_dropped(&self) -> u64 {
        self.messages_dropped.load(Ordering::Relaxed)
    }
}

/// Flow controller statistics
#[derive(Debug, Clone, Default)]
pub struct FlowControlStats {
    /// Total messages rate limited
    pub rate_limited: u64,
    /// Total messages queue dropped
    pub queue_dropped: u64,
    /// Total syncs blocked by cooldown
    pub cooldown_blocked: u64,
    /// Total memory usage across all peers (bytes)
    pub total_memory_usage: u64,
    /// Number of active peers
    pub active_peers: usize,
}

/// Main flow controller
///
/// Coordinates rate limiting, queue management, and resource tracking
/// for all peers.
pub struct FlowController {
    /// Configuration
    config: FlowControlConfig,
    /// Rate limiters per peer
    rate_limiters: Arc<RwLock<HashMap<EndpointId, TokenBucket>>>,
    /// Sync cooldown tracker
    cooldowns: Arc<RwLock<SyncCooldownTracker>>,
    /// Resource trackers per peer
    resources: Arc<RwLock<HashMap<EndpointId, PeerResourceTracker>>>,
    /// Global rate limit counter
    rate_limited_count: AtomicU64,
}

impl FlowController {
    /// Create a new flow controller with default config
    pub fn new() -> Self {
        Self::with_config(FlowControlConfig::default())
    }

    /// Create a new flow controller with custom config
    pub fn with_config(config: FlowControlConfig) -> Self {
        Self {
            cooldowns: Arc::new(RwLock::new(SyncCooldownTracker::new(config.sync_cooldown))),
            config,
            rate_limiters: Arc::new(RwLock::new(HashMap::new())),
            resources: Arc::new(RwLock::new(HashMap::new())),
            rate_limited_count: AtomicU64::new(0),
        }
    }

    /// Check if a sync operation is allowed
    ///
    /// Performs all flow control checks:
    /// 1. Rate limiting (token bucket)
    /// 2. Sync cooldown (storm prevention)
    /// 3. Resource limits (memory)
    ///
    /// Returns Ok(()) if allowed, Err with specific reason if blocked.
    pub fn check_sync_allowed(
        &self,
        peer_id: &EndpointId,
        doc_key: &str,
    ) -> Result<(), FlowControlError> {
        // 1. Check rate limit
        {
            let mut limiters = self
                .rate_limiters
                .write()
                .unwrap_or_else(|e| e.into_inner());
            let limiter = limiters.entry(*peer_id).or_insert_with(|| {
                TokenBucket::new(
                    self.config.max_messages_per_second,
                    self.config.tokens_per_refill,
                    self.config.refill_interval,
                )
            });

            if !limiter.try_acquire() {
                self.rate_limited_count.fetch_add(1, Ordering::Relaxed);
                return Err(FlowControlError::RateLimitExceeded);
            }
        }

        // 2. Check cooldown
        {
            let mut cooldowns = self.cooldowns.write().unwrap_or_else(|e| e.into_inner());
            cooldowns.check_cooldown(peer_id, doc_key)?;
        }

        Ok(())
    }

    /// Record a successful sync operation
    ///
    /// Updates cooldown tracker to prevent sync storms.
    pub fn record_sync(&self, peer_id: &EndpointId, doc_key: &str) {
        let mut cooldowns = self.cooldowns.write().unwrap_or_else(|e| e.into_inner());
        cooldowns.record_sync(peer_id, doc_key);
    }

    /// Get or create resource tracker for peer
    pub fn get_resource_tracker(&self, peer_id: &EndpointId) -> Arc<PeerResourceTracker> {
        let mut resources = self.resources.write().unwrap_or_else(|e| e.into_inner());
        if !resources.contains_key(peer_id) {
            resources.insert(
                *peer_id,
                PeerResourceTracker::new(self.config.max_memory_per_peer as u64),
            );
        }
        // Return a clone since PeerResourceTracker contains atomics that are already thread-safe
        // This is a simplification - in production you'd use Arc<PeerResourceTracker>
        // For now, create a snapshot
        let tracker = resources.get(peer_id).unwrap();
        Arc::new(PeerResourceTracker {
            memory_usage: AtomicU64::new(tracker.memory_usage.load(Ordering::Acquire)),
            max_memory: tracker.max_memory,
            messages_sent: AtomicU64::new(tracker.messages_sent.load(Ordering::Relaxed)),
            messages_received: AtomicU64::new(tracker.messages_received.load(Ordering::Relaxed)),
            messages_dropped: AtomicU64::new(tracker.messages_dropped.load(Ordering::Relaxed)),
        })
    }

    /// Get current statistics
    pub fn stats(&self) -> FlowControlStats {
        let cooldowns = self.cooldowns.read().unwrap_or_else(|e| e.into_inner());
        let resources = self.resources.read().unwrap_or_else(|e| e.into_inner());

        let total_memory: u64 = resources
            .values()
            .map(|r| r.memory_usage.load(Ordering::Relaxed))
            .sum();

        let queue_dropped: u64 = resources
            .values()
            .map(|r| r.messages_dropped.load(Ordering::Relaxed))
            .sum();

        FlowControlStats {
            rate_limited: self.rate_limited_count.load(Ordering::Relaxed),
            queue_dropped,
            cooldown_blocked: cooldowns.blocked_count(),
            total_memory_usage: total_memory,
            active_peers: resources.len(),
        }
    }

    /// Clean up stale data
    pub fn cleanup(&self) {
        let mut cooldowns = self.cooldowns.write().unwrap_or_else(|e| e.into_inner());
        cooldowns.cleanup();
    }

    /// Get current config
    pub fn config(&self) -> &FlowControlConfig {
        &self.config
    }

    /// Get available tokens for a peer
    pub fn available_tokens(&self, peer_id: &EndpointId) -> u32 {
        let limiters = self.rate_limiters.read().unwrap_or_else(|e| e.into_inner());
        limiters
            .get(peer_id)
            .map(|l| l.available_tokens())
            .unwrap_or(self.config.max_messages_per_second)
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    fn create_test_peer_id() -> EndpointId {
        use iroh::SecretKey;
        let mut rng = rand::rng();
        SecretKey::generate(&mut rng).public()
    }

    #[test]
    fn test_token_bucket_basic() {
        let bucket = TokenBucket::new(10, 1, Duration::from_millis(100));

        // Should have 10 tokens initially
        assert_eq!(bucket.available_tokens(), 10);

        // Consume all tokens
        for _ in 0..10 {
            assert!(bucket.try_acquire());
        }

        // Should be empty now
        assert!(!bucket.try_acquire());
        assert_eq!(bucket.available_tokens(), 0);
    }

    #[test]
    fn test_token_bucket_refill() {
        let bucket = TokenBucket::new(10, 5, Duration::from_millis(10));

        // Consume all tokens
        for _ in 0..10 {
            bucket.try_acquire();
        }
        assert_eq!(bucket.available_tokens(), 0);

        // Wait for refill
        std::thread::sleep(Duration::from_millis(25));

        // Should have refilled some tokens (at least 5)
        let available = bucket.available_tokens();
        assert!(
            available >= 5,
            "Expected at least 5 tokens, got {}",
            available
        );
    }

    #[test]
    fn test_bounded_queue_basic() {
        let mut queue: BoundedQueue<i32> = BoundedQueue::new(3);

        assert!(queue.is_empty());
        assert_eq!(queue.len(), 0);

        // Enqueue items
        queue.enqueue(1);
        queue.enqueue(2);
        queue.enqueue(3);

        assert_eq!(queue.len(), 3);
        assert_eq!(queue.total_enqueued(), 3);
        assert_eq!(queue.total_dropped(), 0);

        // Dequeue
        assert_eq!(queue.dequeue(), Some(1));
        assert_eq!(queue.dequeue(), Some(2));
        assert_eq!(queue.dequeue(), Some(3));
        assert_eq!(queue.dequeue(), None);
    }

    #[test]
    fn test_bounded_queue_overflow() {
        let mut queue: BoundedQueue<i32> = BoundedQueue::new(3);

        queue.enqueue(1);
        queue.enqueue(2);
        queue.enqueue(3);

        // This should drop item 1
        let dropped = queue.enqueue(4);
        assert_eq!(dropped, Some(1));
        assert_eq!(queue.total_dropped(), 1);

        // Queue should now be [2, 3, 4]
        assert_eq!(queue.dequeue(), Some(2));
        assert_eq!(queue.dequeue(), Some(3));
        assert_eq!(queue.dequeue(), Some(4));
    }

    #[test]
    fn test_bounded_queue_try_enqueue() {
        let mut queue: BoundedQueue<i32> = BoundedQueue::new(2);

        assert!(queue.try_enqueue(1).is_ok());
        assert!(queue.try_enqueue(2).is_ok());
        assert!(queue.try_enqueue(3).is_err()); // Should fail
    }

    #[test]
    fn test_sync_cooldown_tracker() {
        let peer_id = create_test_peer_id();
        let mut tracker = SyncCooldownTracker::new(Duration::from_millis(50));

        // First sync should be allowed
        assert!(tracker.check_cooldown(&peer_id, "doc1").is_ok());
        tracker.record_sync(&peer_id, "doc1");

        // Immediate second sync should be blocked
        let result = tracker.check_cooldown(&peer_id, "doc1");
        assert!(matches!(
            result,
            Err(FlowControlError::CooldownActive { .. })
        ));

        // Different doc should be allowed
        assert!(tracker.check_cooldown(&peer_id, "doc2").is_ok());

        // Wait for cooldown
        std::thread::sleep(Duration::from_millis(60));

        // Should be allowed now
        assert!(tracker.check_cooldown(&peer_id, "doc1").is_ok());
    }

    #[test]
    fn test_peer_resource_tracker() {
        let tracker = PeerResourceTracker::new(1000);

        // Should start empty
        assert_eq!(tracker.memory_usage(), 0);

        // Allocate some memory
        assert!(tracker.try_allocate(500));
        assert_eq!(tracker.memory_usage(), 500);

        // Allocate more
        assert!(tracker.try_allocate(400));
        assert_eq!(tracker.memory_usage(), 900);

        // This should fail (would exceed 1000)
        assert!(!tracker.try_allocate(200));
        assert_eq!(tracker.memory_usage(), 900);

        // Free some
        tracker.free(300);
        assert_eq!(tracker.memory_usage(), 600);
    }

    #[test]
    fn test_flow_controller_rate_limiting() {
        let config = FlowControlConfig {
            max_messages_per_second: 5,
            tokens_per_refill: 1,
            refill_interval: Duration::from_millis(200),
            sync_cooldown: Duration::ZERO, // Disable cooldown for rate limit testing
            ..Default::default()
        };
        let controller = FlowController::with_config(config);
        let peer_id = create_test_peer_id();

        // Should allow first 5 syncs
        for i in 0..5 {
            assert!(
                controller.check_sync_allowed(&peer_id, "doc1").is_ok(),
                "Sync {} should be allowed",
                i
            );
            controller.record_sync(&peer_id, "doc1");
        }

        // 6th should be rate limited
        let result = controller.check_sync_allowed(&peer_id, "doc1");
        assert!(
            matches!(result, Err(FlowControlError::RateLimitExceeded)),
            "Expected rate limit, got {:?}",
            result
        );
    }

    #[test]
    fn test_flow_controller_stats() {
        let controller = FlowController::new();
        let peer_id = create_test_peer_id();

        // Do some operations
        controller.check_sync_allowed(&peer_id, "doc1").ok();
        controller.record_sync(&peer_id, "doc1");

        let stats = controller.stats();
        assert_eq!(stats.active_peers, 0); // No resource tracker created yet
        assert_eq!(stats.rate_limited, 0);
    }

    #[test]
    fn test_flow_controller_cleanup() {
        let config = FlowControlConfig {
            sync_cooldown: Duration::from_millis(10),
            ..Default::default()
        };
        let controller = FlowController::with_config(config);
        let peer_id = create_test_peer_id();

        // Record a sync
        controller.record_sync(&peer_id, "doc1");

        // Wait for cooldown to expire and beyond
        std::thread::sleep(Duration::from_millis(150));

        // Cleanup should not panic
        controller.cleanup();
    }
}