modelTesting/outputs/muse-glimmer-28b-gguf-rust.rs

// Cargo.toml
// tokio = { version = "1", features = ["full"] }
// tokio-util = "0.7"
// serde = { version = "1", features = ["derive"] }
// thiserror = "1"
// rand = "0.9"

use rand::Rng;
use serde::Serialize;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use thiserror::Error;
use tokio::sync::{mpsc, Mutex};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;

#[derive(Error, Debug)]
enum FetchError {
    #[error("transient")]
    Transient,
    #[error("always fail")]
    AlwaysFail,
}

async fn mock_fetch(id: u32) -> Result<Vec<String>, FetchError> {
    tokio::time::sleep(Duration::from_millis(10)).await;
    if id == 999 {
        return Err(FetchError::AlwaysFail);
    }
    let mut rng = rand::rng();
    if rng.random_bool(0.15) {
        Err(FetchError::Transient)
    } else {
        Ok((0..rng.random_range(0..4))
            .map(|i| format!("item-{}-{}", id, i))
            .collect())
    }
}

#[derive(Debug, Clone, Serialize)]
struct Item {
    watcher_id: u32,
    items: Vec<String>,
    ts: u64,
}

#[derive(Debug, Serialize, Default, Clone)]
struct Aggregated {
    items: Vec<Item>,
    total_count: usize,
}

struct WatcherMeta {
    token: CancellationToken,
    handle: JoinHandle<()>,
    consecutive_failures: Arc<Mutex<u32>>,
    healthy: Arc<Mutex<bool>>,
}

pub struct WatcherManager {
    pub(crate) tx: mpsc::Sender<Item>,
    pub(crate) aggregated: Arc<Mutex<Aggregated>>,
    pub(crate) watchers: Arc<Mutex<HashMap<u32, WatcherMeta>>>,
    pub(crate) next_id: Arc<Mutex<u32>>,
    pub(crate) shutdown_token: CancellationToken,
    pub(crate) consumer_handle: JoinHandle<Aggregated>,
}

impl WatcherManager {
    pub async fn new(num_watchers: usize) -> Self {
        let (tx, rx) = mpsc::channel(32); // bounded -> backpressure
        let aggregated = Arc::new(Mutex::new(Aggregated::default()));
        let shutdown_token = CancellationToken::new();
        let watchers = Arc::new(Mutex::new(HashMap::new()));
        let next_id = Arc::new(Mutex::new(1));
        let consumer_handle = tokio::spawn(consumer_task(rx, shutdown_token.clone(), aggregated.clone()));

        let manager = WatcherManager {
            tx,
            aggregated,
            watchers,
            next_id,
            shutdown_token,
            consumer_handle,
        };

        for _ in 0..num_watchers {
            manager.add_watcher().await;
        }
        manager
    }

    pub async fn add_watcher(&self) -> u32 {
        let mut id_guard = self.next_id.lock().await;
        let id = *id_guard;
        *id_guard += 1;
        drop(id_guard);
        self.add_watcher_with_id(id).await
    }

    pub async fn add_watcher_with_id(&self, id: u32) -> u32 {
        let token = CancellationToken::new();
        let tx = self.tx.clone();
        let shutdown_token = self.shutdown_token.clone();
        let consecutive_failures = Arc::new(Mutex::new(0u32));
        let healthy = Arc::new(Mutex::new(true));

        let handle = tokio::spawn(watcher_task(
            id,
            token.clone(),
            shutdown_token,
            tx,
            consecutive_failures.clone(),
            healthy.clone(),
        ));

        let meta = WatcherMeta {
            token,
            handle,
            consecutive_failures,
            healthy,
        };
        self.watchers.lock().await.insert(id, meta);
        id
    }

    pub async fn remove_watcher(&self, id: u32) -> Result<(), String> {
        let mut watchers = self.watchers.lock().await;
        if let Some(meta) = watchers.remove(&id) {
            meta.token.cancel();
            drop(watchers);
            let _ = meta.handle.await;
            Ok(())
        } else {
            Err("not found".into())
        }
    }

    pub async fn shutdown(&mut self) -> Aggregated {
        self.shutdown_token.cancel();
        let mut watchers = self.watchers.lock().await;
        let mut handles = Vec::new();
        for (_, meta) in watchers.drain() {
            meta.token.cancel();
            handles.push(meta.handle);
        }
        drop(watchers);
        for h in handles {
            let _ = h.await;
        }
        self.consumer_handle.await.unwrap()
    }
}

async fn watcher_task(
    id: u32,
    token: CancellationToken,
    shutdown_token: CancellationToken,
    tx: mpsc::Sender<Item>,
    consecutive_failures: Arc<Mutex<u32>>,
    healthy: Arc<Mutex<bool>>,
) {
    let mut interval = tokio::time::interval(Duration::from_millis(100));
    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
    loop {
        tokio::select! {
            _ = token.cancelled() => break,
            _ = shutdown_token.cancelled() => break,
            _ = interval.tick() => {}
        }
        match mock_fetch(id).await {
            Ok(items) => {
                *consecutive_failures.lock().await = 0;
                *healthy.lock().await = true;
                if !items.is_empty() {
                    let item = Item {
                        watcher_id: id,
                        items,
                        ts: tokio::time::Instant::now().elapsed().as_millis() as u64,
                    };
                    // backpressure: await if channel full
                    if tx.send(item).await.is_err() {
                        break;
                    }
                }
            }
            Err(e) => {
                let mut fails = consecutive_failures.lock().await;
                *fails += 1;
                if *fails > 5 {
                    *healthy.lock().await = false;
                }
                eprintln!("watcher {} error: {:?}", id, e);
            }
        }
    }
}

async fn consumer_task(
    mut rx: mpsc::Receiver<Item>,
    shutdown_token: CancellationToken,
    aggregated: Arc<Mutex<Aggregated>>,
) -> Aggregated {
    let mut agg = Aggregated::default();
    loop {
        tokio::select! {
            _ = shutdown_token.cancelled() => break,
            maybe_item = rx.recv() => {
                match maybe_item {
                    Some(item) => {
                        agg.items.push(item);
                        agg.total_count += 1;
                        *aggregated.lock().await = agg.clone();
                    }
                    None => break,
                }
            }
        }
    }
    agg
}

#[tokio::test]
async fn test_manager_runs_and_shuts_down() {
    let mut manager = WatcherManager::new(4).await;
    tokio::time::sleep(Duration::from_millis(500)).await;
    let agg = manager.shutdown().await;
    assert!(agg.total_count > 0, "consumer should have received items");
}

#[tokio::test]
async fn test_unhealthy_watcher() {
    let mut manager = WatcherManager::new(0).await;
    let id = manager.add_watcher_with_id(999).await;
    tokio::time::sleep(Duration::from_millis(800)).await;
    let watchers = manager.watchers.lock().await;
    let meta = watchers.get(&id).expect("watcher exists");
    let healthy = *meta.healthy.lock().await;
    assert!(!healthy, "watcher should be marked unhealthy after >5 failures");
    drop(watchers);
    let _ = manager.shutdown().await;
}

#[tokio::test]
async fn test_add_remove_concurrent() {
    use std::sync::Arc;
    let manager = Arc::new(Mutex::new(WatcherManager::new(0).await));
    let mut handles = Vec::new();
    for _ in 0..20 {
        let m = manager.clone();
        handles.push(tokio::spawn(async move {
            let mut mgr = m.lock().await;
            let id = mgr.add_watcher().await;
            drop(mgr);
            tokio::time::sleep(Duration::from_millis(5)).await;
            let mut mgr = m.lock().await;
            let _ = mgr.remove_watcher(id).await;
        }));
    }
    for h in handles {
        h.await.unwrap();
    }
    let mgr = manager.lock().await;
    assert_eq!(mgr.watchers.lock().await.len(), 0);
}