```rust // ───────────────────────────────────────────────────────────────────────────── // Cargo.toml dependencies (for reference): // tokio = { version = "1", features = ["full"] } // serde = { version = "1", features = ["derive"] } // futures = "0.3" // ───────────────────────────────────────────────────────────────────────────── use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; use futures::stream::{self, StreamExt}; use serde::Serialize; use tokio::sync::{broadcast, mpsc, RwLock}; use tokio::time::{interval, sleep}; // ─── Error types ────────────────────────────────────────────────────────────── /// Errors that can occur when fetching data from a source. #[derive(Debug, thiserror::Error)] enum FetchError { #[error("source temporarily unavailable")] TemporarilyUnavailable, #[error("upstream timeout")] Timeout, } /// Errors specific to the watcher service. #[derive(Debug, thiserror::Error)] enum ServiceError { #[error("watcher with id {0} not found")] WatcherNotFound(u32), #[error("channel send failed: {0}")] SendError(String), #[error("shutdown already in progress")] AlreadyShutdown, } // ─── Shared data types ──────────────────────────────────────────────────────── /// A single item produced by a watcher's source poll. #[derive(Debug, Clone, Serialize)] struct WatchedItem { watcher_id: u32, value: String, } /// Events flowing from a watcher to the consumer. #[derive(Debug, Clone)] enum WatchEvent { Item(WatchedItem), /// Health update: `true` = healthy, `false` = unhealthy. HealthUpdate { id: u32, healthy: bool }, } /// Aggregated output produced by the consumer. #[derive(Debug, Clone, Serialize)] struct AggregatedOutput { total_items: u64, items_per_watcher: HashMap, } // ─── Mock source ────────────────────────────────────────────────────────────── /// Simulates an async data source. /// /// Returns 0–3 random strings on success (~85 %), or a `FetchError` (~15 %). async fn mock_fetch(id: u32) -> Result, FetchError> { // Simulate variable latency. sleep(Duration::from_millis(rand_u32() % 50)).await; if rand_u32() % 100 < 15 { // ~15 % failure rate. if rand_u32() % 2 == 0 { return Err(FetchError::TemporarilyUnavailable); } return Err(FetchError::Timeout); } let count = (rand_u32() % 4) as usize; // 0..=3 Ok((0..count) .map(|i| format!("item-{}-{}", id, i)) .collect()) } /// Tiny deterministic PRNG helper so we don't pull in a crate. fn rand_u32() -> u32 { use std::sync::atomic::{AtomicU32, Ordering}; // In a real service you'd use a proper RNG; this is fine for tests. static SEED: AtomicU32 = AtomicU32::new(0xDEADBEEF); let s = SEED.fetch_add(6364136223846793005, Ordering::Relaxed); s.wrapping_mul(6364136223846793005).wrapping_add(1) } // ─── Watcher task ───────────────────────────────────────────────────────────── /// A single watcher polls `mock_fetch` on a schedule and forwards results. /// /// * Bounded channel capacity: **32** events. When full, `send()` awaits /// until the consumer drains space — this is the backpressure mechanism. /// * Repeated failures (>5 consecutive) mark the watcher unhealthy and it /// stops polling, sending a `HealthUpdate` before exiting. async fn watcher_task( id: u32, tx: mpsc::Sender, mut shutdown_rx: broadcast::Receiver<()>, ) { let mut interval = interval(Duration::from_millis(80)); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); let mut consecutive_failures: u32 = 0; loop { tokio::select! { _ = shutdown_rx.recv() => { tracing::info!(watcher_id = id, "shutdown signal received"); break; } _ = interval.tick() => { match mock_fetch(id).await { Ok(values) => { consecutive_failures = 0; for value in values { let item = WatchedItem { watcher_id: id, value }; if tx.send(WatchEvent::Item(item)).await.is_err() { tracing::warn!(watcher_id = id, "channel closed, stopping"); return; } } } Err(e) => { consecutive_failures += 1; tracing::warn!( watcher_id = id, failures = consecutive_failures, error = %e, "poll failed" ); if consecutive_failures > 5 { tracing::error!(watcher_id = id, "watcher marked unhealthy after {} consecutive failures", consecutive_failures); let _ = tx.send(WatchEvent::HealthUpdate { id, healthy: false }).await; return; } } } } } } tracing::info!(watcher_id = id, "watcher task exited cleanly"); } // ─── Consumer task ──────────────────────────────────────────────────────────── /// Receives from all watcher channels (merged via `select_all`) and aggregates. /// /// The merged stream automatically removes closed receivers, so when a watcher /// drops its sender the consumer adapts without explicit coordination. async fn consumer_task( mut receivers: Vec>, output_tx: mpsc::Sender, mut shutdown_rx: broadcast::Receiver<()>, ) -> AggregatedOutput { let mut output = AggregatedOutput { total_items: 0, items_per_watcher: HashMap::new(), }; let mut stream = stream::select_all( receivers .drain(..) .map(|rx| rx.map(Ok::<_, Infallible>)) .collect::>(), ); loop { tokio::select! { _ = shutdown_rx.recv() => { tracing::info!("consumer: shutdown signal received"); break; } next = stream.next() => { match next { Some(Ok(WatchEvent::Item(item))) => { output.total_items += 1; *output.items_per_watcher.entry(item.watcher_id).or_insert(0) += 1; tracing::trace!( watcher_id = item.watcher_id, value = %item.value, "consumed item" ); } Some(Ok(WatchEvent::HealthUpdate { id, healthy })) => { tracing::info!(watcher_id = id, healthy, "health update received"); } Some(Err(_)) => unreachable!("Infallible"), None => { // All senders dropped — this happens on shutdown. tracing::info!("consumer: all watcher channels closed"); break; } } } } } // Flush final output before exiting. if output.total_items > 0 { let _ = output_tx.send(output.clone()).await; } tracing::info!(total_items = output.total_items, "consumer exited"); output } // ─── WatcherManager ─────────────────────────────────────────────────────────── /// Manages a dynamic set of watcher tasks with add/remove and clean shutdown. /// /// Shared state is protected by `Arc>` so that add/remove /// operations are race-free and concurrent with running watchers. #[derive(Debug)] struct WatcherSet { /// Live watcher entries keyed by their ID. watchers: HashMap, /// Current health status per watcher (true = healthy). health_status: HashMap, } #[derive(Debug)] struct WatcherEntry { id: u32, /// Sender half of the watcher→consumer channel. Dropping this closes the /// receiver, causing the consumer's `select_all` stream to drop it. tx: mpsc::Sender, /// Join handle for the watcher task. handle: tokio::task::JoinHandle<()>, } #[derive(Debug)] pub struct WatcherManager { inner: Arc>, /// Broadcast channel used to signal all tasks to shut down. shutdown_tx: broadcast::Sender<()>, /// Sender for the final aggregated output. output_tx: mpsc::Sender, } impl WatcherManager { /// Creates a new manager and spawns the consumer task. pub fn new() -> (Self, mpsc::Receiver) { let (output_tx, output_rx) = mpsc::channel(1); let (shutdown_tx, _) = broadcast::channel(1); let manager = Self { inner: Arc::new(RwLock::new(WatcherSet { watchers: HashMap::new(), health_status: HashMap::new(), })), shutdown_tx: shutdown_tx.clone(), output_tx, }; // Spawn the consumer; it will receive a clone of `shutdown_tx` internally. tokio::spawn(manager.clone_consumer(shutdown_tx)); (manager, output_rx) } /// Returns the number of currently active watchers. pub async fn watcher_count(&self) -> usize { self.inner.read().await.watchers.len() } /// Returns the health status of a watcher, if it exists. pub async fn is_healthy(&self, id: u32) -> Option { self.inner.read().await.health_status.get(&id).copied() } /// Adds a new watcher with the given ID. Returns `Ok(())` or an error if /// a watcher with that ID already exists. pub async fn add_watcher(&self, id: u32) -> Result<(), ServiceError> { let mut set = self.inner.write().await; if set.watchers.contains_key(&id) { return Err(ServiceError::WatcherNotFound(id)); } // Bounded channel: capacity 32. When full, `send()` awaits until the // consumer drains space — this is our backpressure mechanism. If the // consumer is permanently slow, watchers will block on send rather than // buffering unboundedly in memory. let (tx, rx) = mpsc::channel(32); let mut shutdown_rx = self.shutdown_tx.subscribe(); let handle = tokio::spawn(watcher_task(id, tx.clone(), shutdown_rx)); set.watchers.insert( id, WatcherEntry { id, tx, handle }, ); set.health_status.insert(id, true); // We need to give the consumer its receiver. Since the consumer was // spawned before any watchers existed, we store receivers inside the // shared state and the consumer reads them dynamically. // // Actually, the consumer uses `select_all` on a Vec it owns. We need // to push the new receiver into that Vec. We'll store receivers in the // shared state and have the consumer re-build its stream periodically. // // Simpler approach: store receivers in the shared state and have a // separate "receiver manager" that the consumer polls. drop(set); // Push receiver into the consumer's collection via a side channel. // (In production you'd use a dedicated coordination mechanism.) // For this demo, we'll store receivers in the shared state and have // the consumer rebuild its select_all stream when notified. Ok(()) } /// Removes a watcher by ID, dropping its sender so the task exits. pub async fn remove_watcher(&self, id: u32) -> Result<(), ServiceError> { let mut set = self.inner.write().await; if let Some(entry) = set.watchers.remove(&id) { // Dropping `tx` closes the channel; the watcher task will see // `send()` fail and exit. We await the handle to ensure clean join. let id = entry.id; tokio::spawn(async move { if entry.handle.await.is_err() { tracing::warn!(watcher_id = id, "watcher task panicked"); } }); set.health_status.remove(&id); } else { return Err(ServiceError::WatcherNotFound(id)); } Ok(()) } /// Signals all tasks to shut down and waits for them to finish. pub async fn shutdown(&self) -> Result<(), ServiceError> { // Broadcast the shutdown signal. All watchers and the consumer listen // to this channel and will exit their select! loops. let _ = self.shutdown_tx.send(()); // Drop all senders to ensure receivers close even if a task misses the // broadcast (e.g., it's blocked on `send()`). let set = self.inner.read().await; for entry in set.watchers.values() { // tx is dropped here when `set` is dropped at end of scope. } drop(set); // Give tasks a moment to notice the signal and exit. tokio::time::sleep(Duration::from_millis(200)).await; // Join all handles. let set = self.inner.read().await; for entry in set.watchers.values() { let handle = entry.handle.clone(); let id = entry.id; tokio::spawn(async move { if handle.await.is_err() { tracing::warn!(watcher_id = id, "watcher task panicked during shutdown"); } }); } Ok(()) } /// Clone for sharing between manager and consumer. fn clone_consumer(&self, shutdown_tx: broadcast::Sender<()>) -> Self { Self { inner: self.inner.clone(), shutdown_tx, output_tx: self.output_tx.clone(), } } /// Internal consumer loop that rebuilds its receiver stream when watchers /// are added or removed. async fn run_consumer(self) { // We store receivers in the shared state alongside entries. // Re-implementing with a cleaner design: store (tx, rx) pairs in the set. // For this demo, we'll use a simpler approach with a dedicated receiver vec. // Actually, let me restructure: the consumer will poll the shared state // to collect current receivers and rebuild select_all each iteration. // This is slightly less efficient but correct for a demo. let mut receivers: Vec> = Vec::new(); let mut last_count = 0; loop { // Collect current receivers from shared state. let set = self.inner.read().await; let new_receivers: Vec<_> = set .watchers .values() .map(|entry| { let (tx, rx) = mpsc::channel(32); // We can't actually create new receivers for existing senders. // This approach won't work. drop(tx); rx }) .collect(); drop(set); // This approach is flawed. Let me use a different design. break; } // ── Revised consumer using a receiver coordination channel ───────── // Each watcher stores its receiver in the shared state. The consumer // reads receivers directly from there using select_all rebuilt on // changes. We use a oneshot-per-change notification. } } // ─── Revised architecture (cleaner) ────────────────────────────────────────── // // To avoid the dynamic-receiver problem, we use a single shared mpsc channel // from all watchers to the consumer. Each watcher clones the sender before // spawning. The consumer reads from one receiver. This is the standard // tokio pattern and avoids compile-time select! limitations entirely. // // Per-watcher health is tracked via an Arc stored in the shared // state, updated by the watcher task itself. use std::sync::atomic::{AtomicBool, Ordering}; /// Revised WatcherSet that stores health atomics for per-watcher status. #[derive(Debug)] struct WatcherSetV2 { watchers: HashMap, health_flags: HashMap>, } #[derive(Debug)] struct WatcherEntryV2 { id: u32, handle: tokio::task::JoinHandle<()>, } /// Revised manager using a single shared consumer channel. #[derive(Debug, Clone)] pub struct WatcherManagerV2 { inner: Arc>, /// Single shared sender — cloned for each watcher. Bounded at 32; when /// full, `send()` blocks (backpressure). Dropping all senders closes the /// channel, causing the consumer to exit. shared_tx: mpsc::Sender, shutdown_tx: broadcast::Sender<()>, output_tx: mpsc::Sender, } impl WatcherManagerV2 { /// Creates a new manager and spawns the consumer task. pub fn new() -> (Self, mpsc::Receiver) { let (shared_tx, shared_rx) = mpsc::channel(32); // backpressure bound: 32 let (output_tx, output_rx) = mpsc::channel(1); let (shutdown_tx, _) = broadcast::channel(1); let manager = Self { inner: Arc::new(RwLock::new(WatcherSetV2 { watchers: HashMap::new(), health_flags: HashMap::new(), })), shared_tx: shared_tx.clone(), shutdown_tx: shutdown_tx.clone(), output_tx, }; // Spawn the consumer. tokio::spawn(manager.clone().run_consumer(shared_rx, shutdown_tx)); (manager, output_rx) } /// Returns the number of currently active watchers. pub async fn watcher_count(&self) -> usize { self.inner.read().await.watchers.len() } /// Returns the health status of a watcher, if it exists. pub async fn is_healthy(&self, id: u32) -> Option { self.inner .read() .await .health_flags .get(&id) .map(|flag| flag.load(Ordering::Relaxed)) } /// Adds a new watcher with the given ID. Returns `Ok(())` or an error if /// a watcher with that ID already exists. pub async fn add_watcher(&self, id: u32) -> Result<(), ServiceError> { let mut set = self.inner.write().await; if set.watchers.contains_key(&id) { return Err(ServiceError::WatcherNotFound(id)); } let health_flag = Arc::new(AtomicBool::new(true)); set.health_flags.insert(id, health_flag.clone()); let mut shutdown_rx = self.shutdown_tx.subscribe(); let tx = self.shared_tx.clone(); // clone sender for this watcher let handle = tokio::spawn(watcher_task_v2( id, tx, health_flag, shutdown_rx, )); set.watchers.insert(id, WatcherEntryV2 { id, handle }); Ok(()) } /// Removes a watcher by ID, dropping its sender so the task exits. pub async fn remove_watcher(&self, id: u32) -> Result<(), ServiceError> { let mut set = self.inner.write().await; if !set.watchers.contains_key(&id) { return Err(ServiceError::WatcherNotFound(id)); } let entry = set.watchers.remove(&id).unwrap(); set.health_flags.remove(&id);