path: root/timing/src/mixing_pool.rs

use std::collections::VecDeque;
use std::time::{Duration, Instant};
use rand::seq::SliceRandom;
use serde::{Deserialize, Serialize};

/// A pool that collects items and releases them in shuffled batches
/// Used for temporal obfuscation of transaction broadcasts
#[derive(Debug)]
pub struct MixingPool<T> {
    /// Items waiting to be released
    pool: VecDeque<PoolItem<T>>,

    /// Minimum time to collect items before releasing
    min_batch_duration: Duration,

    /// Maximum time to collect items before forcing release
    max_batch_duration: Duration,

    /// Maximum number of items in pool before forcing release
    max_pool_size: usize,

    /// Time when current batch started collecting
    batch_start: Option<Instant>,
}

#[derive(Debug, Clone)]
struct PoolItem<T> {
    item: T,
    #[allow(dead_code)]
    added_at: Instant,
}

impl<T> MixingPool<T> {
    /// Create a new mixing pool with custom parameters
    pub fn new(
        min_batch_duration: Duration,
        max_batch_duration: Duration,
        max_pool_size: usize,
    ) -> Self {
        Self {
            pool: VecDeque::new(),
            min_batch_duration,
            max_batch_duration,
            max_pool_size,
            batch_start: None,
        }
    }

    /// Create a mixing pool with default parameters (10-30 min batches)
    pub fn default() -> Self {
        Self::new(
            Duration::from_secs(10 * 60), // 10 minutes min
            Duration::from_secs(30 * 60), // 30 minutes max
            1000,                         // max 1000 items
        )
    }

    /// Add an item to the mixing pool
    pub fn add(&mut self, item: T) {
        if self.batch_start.is_none() {
            self.batch_start = Some(Instant::now());
        }

        self.pool.push_back(PoolItem {
            item,
            added_at: Instant::now(),
        });
    }

    /// Check if pool should release a batch
    pub fn should_release(&self) -> bool {
        if self.pool.is_empty() {
            return false;
        }

        if let Some(start) = self.batch_start {
            let elapsed = start.elapsed();

            // Force release if max duration reached
            if elapsed >= self.max_batch_duration {
                return true;
            }

            // Force release if pool is full
            if self.pool.len() >= self.max_pool_size {
                return true;
            }

            // Release if min duration passed and we have items
            if elapsed >= self.min_batch_duration && !self.pool.is_empty() {
                return true;
            }
        }

        false
    }

    /// Release a shuffled batch of all pooled items
    pub fn release_batch(&mut self) -> Vec<T> {
        let mut items: Vec<T> = self.pool.drain(..).map(|pi| pi.item).collect();

        // Shuffle to remove temporal ordering
        let mut rng = rand::thread_rng();
        items.shuffle(&mut rng);

        // Reset batch timer
        self.batch_start = None;

        items
    }

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

    pub fn is_empty(&self) -> bool {
        self.pool.is_empty()
    }

    /// Get time until next potential release (min batch duration)
    pub fn time_until_release(&self) -> Option<Duration> {
        self.batch_start.map(|start| {
            let elapsed = start.elapsed();
            if elapsed >= self.min_batch_duration {
                Duration::from_secs(0)
            } else {
                self.min_batch_duration - elapsed
            }
        })
    }
}

/// Stats about mixing pool performance
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MixingPoolStats {
    pub total_items_processed: u64,
    pub total_batches_released: u64,
    pub avg_batch_size: f64,
    pub avg_wait_time_secs: f64,
}

/// Mixing pool with statistics tracking
#[derive(Debug)]
pub struct MixingPoolWithStats<T> {
    pool: MixingPool<T>,
    stats: MixingPoolStats,
    #[allow(dead_code)]
    item_wait_times: Vec<Duration>,
}

impl<T> MixingPoolWithStats<T> {
    pub fn new(pool: MixingPool<T>) -> Self {
        Self {
            pool,
            stats: MixingPoolStats {
                total_items_processed: 0,
                total_batches_released: 0,
                avg_batch_size: 0.0,
                avg_wait_time_secs: 0.0,
            },
            item_wait_times: Vec::new(),
        }
    }

    pub fn add(&mut self, item: T) {
        self.pool.add(item);
    }

    pub fn should_release(&self) -> bool {
        self.pool.should_release()
    }

    pub fn release_batch(&mut self) -> Vec<T> {
        let batch = self.pool.release_batch();
        let batch_size = batch.len();

        // Update stats
        self.stats.total_items_processed += batch_size as u64;
        self.stats.total_batches_released += 1;

        // Update average batch size
        let total_batches = self.stats.total_batches_released as f64;
        self.stats.avg_batch_size =
            (self.stats.avg_batch_size * (total_batches - 1.0) + batch_size as f64) / total_batches;

        batch
    }

    pub fn get_stats(&self) -> &MixingPoolStats {
        &self.stats
    }

    pub fn len(&self) -> usize {
        self.pool.len()
    }

    pub fn is_empty(&self) -> bool {
        self.pool.is_empty()
    }
}

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

    #[test]
    fn test_mixing_pool_add() {
        let mut pool = MixingPool::<u32>::default();
        pool.add(1);
        pool.add(2);
        pool.add(3);

        assert_eq!(pool.len(), 3);
    }

    #[test]
    fn test_mixing_pool_release() {
        let mut pool = MixingPool::<u32>::new(
            Duration::from_millis(10),
            Duration::from_secs(1),
            100,
        );

        pool.add(1);
        pool.add(2);
        pool.add(3);

        // Should not release immediately
        assert!(!pool.should_release());

        // Wait for min duration
        std::thread::sleep(Duration::from_millis(15));

        assert!(pool.should_release());

        let batch = pool.release_batch();
        assert_eq!(batch.len(), 3);
        assert!(pool.is_empty());
    }

    #[test]
    fn test_mixing_pool_shuffle() {
        let mut pool = MixingPool::<u32>::new(
            Duration::from_millis(10),
            Duration::from_secs(1),
            100,
        );

        // Add items in order
        for i in 0..100 {
            pool.add(i);
        }

        std::thread::sleep(Duration::from_millis(15));

        let batch = pool.release_batch();

        // Check all items are present
        let mut sorted = batch.clone();
        sorted.sort();
        assert_eq!(sorted, (0..100).collect::<Vec<_>>());

        // Check they were shuffled (very unlikely to be in original order)
        assert_ne!(batch, sorted);
    }

    #[test]
    fn test_mixing_pool_max_size() {
        let mut pool = MixingPool::<u32>::new(
            Duration::from_secs(100), // Long min duration
            Duration::from_secs(200),
            10, // Small max size
        );

        // Add items up to max
        for i in 0..10 {
            pool.add(i);
        }

        // Should force release due to max size
        assert!(pool.should_release());
    }

    #[test]
    fn test_mixing_pool_with_stats() {
        let pool = MixingPool::<u32>::new(
            Duration::from_millis(10),
            Duration::from_secs(1),
            100,
        );
        let mut pool_with_stats = MixingPoolWithStats::new(pool);

        for i in 0..5 {
            pool_with_stats.add(i);
        }

        std::thread::sleep(Duration::from_millis(15));

        pool_with_stats.release_batch();

        let stats = pool_with_stats.get_stats();
        assert_eq!(stats.total_items_processed, 5);
        assert_eq!(stats.total_batches_released, 1);
        assert_eq!(stats.avg_batch_size, 5.0);
    }
}