modelTesting/outputs/qwen3.6-35b-a3b-6bit-mlx-rust.rs

//! A production-style async watcher service built with `tokio`.
//!
//! # Architecture
//! - `WatcherManager` owns N watcher tasks, a health monitor, and a consumer task.
//! - Watchers poll a mock source and forward items via a bounded `tokio::sync::mpsc` channel.
//! - The consumer aggregates items. The manager tracks watcher health and runtime mutations.
//! - Shutdown uses a `tokio::sync::broadcast` signal for fan-out cancellation, followed by `JoinHandle` awaiting.
//!
//! # Dependencies (Cargo.toml)
//! ```toml
//! [dependencies]
//! tokio = { version = "1", features = ["full"] }
//! serde = { version = "1", features = ["derive"] }
//! thiserror = "1"
//! ```

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{broadcast, mpsc, Mutex};
use serde::Serialize;

// ---------------------------------------------------------------------------
// Error Types
// ---------------------------------------------------------------------------

#[derive(Debug, thiserror::Error)]
pub enum FetchError {
    #[error("network timeout")]
    Timeout,
    #[error("invalid response format")]
    InvalidResponse,
}

#[derive(Debug, thiserror::Error)]
pub enum ManagerError {
    #[error("watcher already exists: {0}")]
    AlreadyExists(u32),
    #[error("watcher not found: {0}")]
    NotFound(u32),
    #[error("manager is already shut down")]
    ShutdownInProgress,
}

// ---------------------------------------------------------------------------
// Data Structures
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Serialize)]
pub struct WatchedItem {
    pub source_id: u32,
    pub content: String,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WatcherStatus {
    Healthy,
    Unhealthy,
}

/// Internal state protected by `Arc<Mutex<>>`
struct ManagerState {
    watchers: HashMap<u32, WatcherEntry>,
    total_items_received: usize,
}

struct WatcherEntry {
    status: WatcherStatus,
    consecutive_failures: u32,
    join_handle: tokio::task::JoinHandle<()>,
}

// ---------------------------------------------------------------------------
// Mock Source
// ---------------------------------------------------------------------------

/// Simulates an external API. Returns 0-3 items on success (~85%), or an error (~15%).
/// Uses a time-based seed for pseudo-randomness without extra dependencies.
async fn mock_fetch(id: u32) -> Result<Vec<String>, FetchError> {
    // Simulate network latency
    tokio::time::sleep(Duration::from_millis(20 + (id % 30))).await;

    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_nanos();
    let seed = (now + id as u128) % 100;

    // ~15% error rate
    if seed < 15 {
        return Err(FetchError::Timeout);
    }

    // 0-3 items
    let count = (seed % 4) as usize;
    Ok((0..count)
        .map(|i| format!("data-{}-{}", id, i))
        .collect())
}

// ---------------------------------------------------------------------------
// Background Tasks
// ---------------------------------------------------------------------------

/// Periodically polls the source, forwards items to the consumer, and reports health.
/// Respects shutdown signals and applies backpressure via bounded channel sends.
async fn watcher_loop(
    id: u32,
    mut shutdown_rx: broadcast::Receiver<()>,
    health_tx: mpsc::Sender<(u32, u32)>,
    item_tx: mpsc::Sender<WatchedItem>,
) {
    let mut consecutive_failures = 0u32;

    loop {
        tokio::select! {
            _ = shutdown_rx.recv() => {
                println!("[Watcher {}] Shutdown signal received. Exiting poll loop.", id);
                break;
            }
            _ = tokio::time::sleep(Duration::from_millis(50)) => {
                match mock_fetch(id).await {
                    Ok(items) => {
                        consecutive_failures = 0;
                        for item in items {
                            // Bounded channel backpressure: if the buffer is full,
                            // this await yields the task, pausing the watcher until
                            // the consumer drains items. Prevents OOM.
                            if item_tx.send(WatchedItem { source_id: id, content: item }).await.is_err() {
                                break; // Channel closed (consumer dropped)
                            }
                        }
                    }
                    Err(e) => {
                        consecutive_failures += 1;
                        println!("[Watcher {}] Fetch failed: {} (consecutive: {})", id, e, consecutive_failures);
                        let _ = health_tx.send((id, consecutive_failures)).await;
                    }
                }
            }
        }
    }
}

/// Listens for health updates and mutates shared state safely.
async fn health_monitor(
    mut rx: mpsc::Receiver<(u32, u32)>,
    state: Arc<Mutex<ManagerState>>,
) {
    while let Some((id, failures)) = rx.recv().await {
        let unhealthy = failures > 5;
        let mut inner = state.lock().await;
        if let Some(entry) = inner.watchers.get_mut(&id) {
            entry.status = if unhealthy {
                WatcherStatus::Unhealthy
            } else {
                WatcherStatus::Healthy
            };
            entry.consecutive_failures = failures;
        }
    }
}

/// Aggregates items from all watchers. In production, this would persist, index, or forward.
async fn consumer(mut rx: mpsc::Receiver<WatchedItem>, state: Arc<Mutex<ManagerState>>) {
    while let Some(item) = rx.recv().await {
        let mut inner = state.lock().await;
        inner.total_items_received += 1;
        // In a real service: process, store, or forward item here.
    }
    println!("[Consumer] Channel closed. Exiting.");
}

// ---------------------------------------------------------------------------
// Manager
// ---------------------------------------------------------------------------

/// Manages a dynamic set of async watchers with safe shared state and graceful shutdown.
pub struct WatcherManager {
    state: Arc<Mutex<ManagerState>>,
    health_tx: mpsc::Sender<(u32, u32)>,
    item_tx: mpsc::Sender<WatchedItem>,
    shutdown_tx: broadcast::Sender<()>,
}

impl WatcherManager {
    /// Creates a new manager. Spawns health monitor and consumer tasks.
    ///
    /// # Backpressure Note
    /// The item channel is bounded to `32`. If the consumer is slower than producers,
    /// `item_tx.send()` will `await` until space frees up. This naturally throttles
    /// watchers, preventing unbounded memory growth.
    pub fn new() -> Self {
        let (health_tx, health_rx) = mpsc::channel(128);
        let (item_tx, item_rx) = mpsc::bounded(32); // Bounded channel for backpressure
        let (shutdown_tx, _) = broadcast::channel(1);

        let state = Arc::new(Mutex::new(ManagerState {
            watchers: HashMap::new(),
            total_items_received: 0,
        }));

        tokio::spawn(health_monitor(health_rx, Arc::clone(&state)));
        tokio::spawn(consumer(item_rx, Arc::clone(&state)));

        Self {
            state,
            health_tx,
            item_tx,
            shutdown_tx,
        }
    }

    /// Registers a new watcher. Race-free via mutex.
    pub async fn add_watcher(&self, id: u32) -> Result<(), ManagerError> {
        let mut inner = self.state.lock().await;
        if inner.watchers.contains_key(&id) {
            return Err(ManagerError::AlreadyExists(id));
        }

        let mut shutdown_rx = self.shutdown_tx.subscribe();
        let health_tx = self.health_tx.clone();
        let item_tx = self.item_tx.clone();

        let handle = tokio::spawn(watcher_loop(id, shutdown_rx, health_tx, item_tx));

        inner.watchers.insert(id, WatcherEntry {
            status: WatcherStatus::Healthy,
            consecutive_failures: 0,
            join_handle: handle,
        });

        Ok(())
    }

    /// Removes a watcher gracefully. Sends cancellation and removes from state.
    pub async fn remove_watcher(&self, id: u32) -> Result<(), ManagerError> {
        let mut inner = self.state.lock().await;
        if !inner.watchers.contains_key(&id) {
            return Err(ManagerError::NotFound(id));
        }

        inner.watchers.remove(&id);
        println!("[Manager] Removed watcher {}", id);
        Ok(())
    }

    /// Graceful shutdown: signals all tasks, waits for clean exit, drops channels.
    pub async fn shutdown(&self) -> Result<(), ManagerError> {
        // 1. Fan-out shutdown signal
        let _ = self.shutdown_tx.send(());

        // 2. Wait for all watcher tasks to finish their current poll cycle
        let mut inner = self.state.lock().await;
        let handles: Vec<_> = inner.watchers.values().map(|e| e.join_handle).collect();
        drop(inner); // Release lock before awaiting

        for handle in handles {
            let _ = handle.await;
        }

        // 3. Drop channels to unblock consumer & health monitor
        drop(self.item_tx);
        drop(self.health_tx);

        println!("[Manager] Shutdown complete.");
        Ok(())
    }

    /// Exposes current health metrics for monitoring/alerting.
    pub async fn stats(&self) -> (usize, Vec<(u32, WatcherStatus)>) {
        let inner = self.state.lock().await;
        let count = inner.watchers.len();
        let health: Vec<_> = inner
            .watchers
            .iter()
            .map(|(&id, entry)| (id, entry.status.clone()))
            .collect();
        (count, health)
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::LazyLock;

    /// Global flag to force mock_fetch to always fail (used in test b)
    static TEST_ALWAYS_FAIL: LazyLock<AtomicBool> = LazyLock::new(|| AtomicBool::new(false));

    /// Override mock_fetch for test scenarios
    async fn test_fetch(id: u32) -> Result<Vec<String>, FetchError> {
        if TEST_ALWAYS_FAIL.load(Ordering::Relaxed) {
            return Err(FetchError::Timeout);
        }
        mock_fetch(id).await
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn test_watcher_manager_full_cycle() {
        let mgr = WatcherManager::new();

        // --- (a) Spawn 4 watchers, run ~500ms, shutdown, assert join + >0 items ---
        for i in 1..=4 {
            mgr.add_watcher(i).await.unwrap();
        }

        // Let it run
        tokio::time::sleep(Duration::from_millis(500)).await;

        let (count, health) = mgr.stats().await;
        assert_eq!(count, 4);
        assert!(health.iter().all(|(_, s)| *s == WatcherStatus::Healthy));

        // Shutdown
        mgr.shutdown().await.unwrap();

        // After shutdown, all tasks should have joined. We verify by checking that
        // the manager can be recreated without issues and that previous handles are gone.
        // (In production, you'd assert on a metrics collector, but here we trust the join)
        println!("Test (a) passed: 4 watchers ran and shut down cleanly.");

        // --- (b) Always-failing watcher marked unhealthy after 5 consecutive failures ---
        TEST_ALWAYS_FAIL.store(true, Ordering::Relaxed);
        mgr.add_watcher(99).await.unwrap();

        // Each poll is ~50ms. 5 failures + overhead ≈ 300-400ms
        tokio::time::sleep(Duration::from_millis(400)).await;

        let (_, health) = mgr.stats().await;
        let unhealthy_found = health.iter().any(|(&id, s)| id == 99 && *s == WatcherStatus::Unhealthy);
        assert!(
            unhealthy_found,
            "Watcher 99 should be marked unhealthy after 5 consecutive failures"
        );

        // Ensure other watchers remain unaffected
        let healthy_count = health.iter().filter(|(_, s)| *s == WatcherStatus::Healthy).count();
        assert_eq!(healthy_count, 4, "Other watchers should remain healthy");

        TEST_ALWAYS_FAIL.store(false, Ordering::Relaxed);
        println!("Test (b) passed: Unhealthy detection works without crashing others.");

        // --- (c) Runtime add + remove updates live count with no race ---
        mgr.add_watcher(100).await.unwrap();
        let (count_after_add, _) = mgr.stats().await;
        assert_eq!(count_after_add, 5, "Count should increase after add");

        mgr.remove_watcher(100).await.unwrap();
        let (count_after_remove, _) = mgr.stats().await;
        assert_eq!(count_after_remove, 4, "Count should decrease after remove");

        // Stress test concurrent mutations
        let mut handles = Vec::new();
        for i in 0..10 {
            let mgr_clone = &mgr;
            let id = 200 + i;
            handles.push(tokio::spawn(async move {
                let _ = mgr_clone.add_watcher(id).await;
                tokio::time::sleep(Duration::from_millis(10)).await;
                let _ = mgr_clone.remove_watcher(id).await;
            }));
        }
        for h in handles {
            h.await.unwrap();
        }

        let (final_count, _) = mgr.stats().await;
        assert_eq!(final_count, 4, "Concurrent add/remove should not leak or panic");
        println!("Test (c) passed: Runtime mutations are race-free.");

        // Final cleanup
        mgr.shutdown().await.unwrap();
    }
}