use crate::api_client::ApiClient; use crate::buffer::StorageError; #[derive(Debug)] pub enum SchedulerError { Api(crate::api_client::ApiError), Storage(StorageError), Snmp(crate::snmp::SnmpError), Executor(super::executor::ExecutorError), } impl std::fmt::Display for SchedulerError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Api(err) => write!(f, "API error: {}", err), Self::Storage(err) => write!(f, "Storage error: {}", err), Self::Snmp(err) => write!(f, "SNMP error: {}", err), Self::Executor(err) => write!(f, "Executor error: {}", err), } } } impl std::error::Error for SchedulerError {} impl From for SchedulerError { fn from(err: crate::api_client::ApiError) -> Self { Self::Api(err) } } impl From for SchedulerError { fn from(err: StorageError) -> Self { Self::Storage(err) } } impl From for SchedulerError { fn from(err: crate::snmp::SnmpError) -> Self { Self::Snmp(err) } } impl From for SchedulerError { fn from(err: super::executor::ExecutorError) -> Self { Self::Executor(err) } } pub type Result = std::result::Result; use crate::buffer::Storage; use crate::config::{AgentConfig, HeartbeatMetadata}; use crate::poller::Executor; use crate::snmp::SnmpClient; use crate::metrics::Timestamp; use log::{error, info, warn}; use std::sync::Arc; use std::time::Duration; use tokio::sync::Semaphore; use tokio::time::interval; /// Main scheduler that orchestrates polling, config refresh, and metrics submission pub struct Scheduler { api_client: ApiClient, storage: Storage, executor: Executor, config_refresh_seconds: u64, current_config: Option, start_time: Timestamp, } impl Scheduler { pub fn new( api_client: ApiClient, storage: Storage, snmp_client: SnmpClient, config_refresh_seconds: u64, ) -> Self { let executor = Executor::new(snmp_client, storage.clone()); Self { api_client, storage, executor, config_refresh_seconds, current_config: None, start_time: Timestamp::now(), } } /// Run the main event loop pub async fn run(&mut self) -> Result<()> { info!("Starting Towerops agent scheduler"); // Fetch initial configuration if let Err(e) = self.refresh_config().await { error!("Failed to fetch initial config: {}", e); } let mut config_ticker = interval(Duration::from_secs(self.config_refresh_seconds)); let mut metrics_ticker = interval(Duration::from_secs(30)); let mut heartbeat_ticker = interval(Duration::from_secs(60)); let mut cleanup_ticker = interval(Duration::from_secs(3600)); // Cleanup every hour let mut poll_ticker = interval(Duration::from_secs(5)); // Check if polling needed every 5s let mut update_ticker = interval(Duration::from_secs(3600)); // Check for updates every hour loop { tokio::select! { _ = config_ticker.tick() => { if let Err(e) = self.refresh_config().await { error!("Failed to refresh config: {}", e); } } _ = metrics_ticker.tick() => { if let Err(e) = self.flush_metrics().await { error!("Failed to flush metrics: {}", e); } } _ = heartbeat_ticker.tick() => { if let Err(e) = self.send_heartbeat().await { warn!("Failed to send heartbeat: {}", e); } } _ = cleanup_ticker.tick() => { if let Err(e) = self.storage.cleanup_old_metrics() { error!("Failed to cleanup old metrics: {}", e); } } _ = poll_ticker.tick() => { if let Err(e) = self.poll_equipment().await { error!("Polling error: {}", e); } } _ = update_ticker.tick() => { self.check_and_update().await; } } } } async fn refresh_config(&mut self) -> Result<()> { info!("Refreshing configuration from API"); match self.api_client.fetch_config().await { Ok(config) => { info!( "Configuration updated: {} equipment items", config.equipment.len() ); self.current_config = Some(config); Ok(()) } Err(e) => { warn!( "Failed to fetch config, continuing with cached config: {}", e ); Err(e.into()) } } } async fn flush_metrics(&self) -> Result<()> { // Process metrics in batches until queue is empty or we hit an error // This handles high-volume scenarios with 10,000+ equipment let mut total_flushed = 0; const BATCH_SIZE: usize = 500; const MAX_BATCHES: usize = 20; // Limit to 10,000 metrics per flush cycle for _ in 0..MAX_BATCHES { let pending = self.storage.get_pending_metrics(BATCH_SIZE)?; if pending.is_empty() { break; } let batch_size = pending.len(); let ids: Vec = pending.iter().map(|(id, _)| *id).collect(); let metrics: Vec<_> = pending.into_iter().map(|(_, m)| m).collect(); match self.api_client.submit_metrics(metrics).await { Ok(_) => { self.storage.mark_metrics_sent(&ids)?; total_flushed += batch_size; } Err(e) => { warn!("Failed to submit batch of {} metrics: {}", batch_size, e); // Don't return error, just log and continue with remaining batches break; } } // If we got less than batch size, we've emptied the queue if batch_size < BATCH_SIZE { break; } } if total_flushed > 0 { info!("Successfully flushed {} metrics to API", total_flushed); } Ok(()) } async fn send_heartbeat(&self) -> Result<()> { let uptime = self.start_time.elapsed_secs() as u64; // Get hostname from environment or system let hostname = std::env::var("HOSTNAME") .or_else(|_| std::fs::read_to_string("/etc/hostname").map(|s| s.trim().to_string())) .unwrap_or_else(|_| "unknown".to_string()); let metadata = HeartbeatMetadata { version: env!("CARGO_PKG_VERSION").to_string(), hostname, uptime_seconds: uptime, }; self.api_client.heartbeat(metadata).await?; Ok(()) } async fn poll_equipment(&self) -> Result<()> { let config = match &self.current_config { Some(c) => c, None => return Ok(()), }; let poll_times = self.storage.get_all_last_poll_times()?; // Collect equipment that needs polling let equipment_to_poll: Vec<_> = config .equipment .iter() .filter(|eq| eq.snmp.enabled) .filter(|eq| { match poll_times.get(&eq.id) { Some(last_poll) => { let elapsed = last_poll.elapsed_secs() as u64; elapsed >= eq.poll_interval_seconds } None => true, // Never polled before } }) .collect(); if equipment_to_poll.is_empty() { return Ok(()); } info!( "Polling {} equipment items in parallel", equipment_to_poll.len() ); // Limit concurrent polling to prevent overwhelming the system // With 10,000+ equipment, we don't want 10,000 concurrent tasks const MAX_CONCURRENT_POLLS: usize = 100; let semaphore = Arc::new(Semaphore::new(MAX_CONCURRENT_POLLS)); // Spawn parallel polling tasks with concurrency limit let mut tasks = Vec::new(); for equipment in equipment_to_poll { let executor = self.executor.clone(); let storage = self.storage.clone(); let equipment = equipment.clone(); let permit = semaphore.clone(); let task = tokio::spawn(async move { // Acquire permit before polling (limits concurrency) let _permit = match permit.acquire().await { Ok(p) => p, Err(e) => { error!("Failed to acquire polling permit: {}", e); return; } }; info!("Polling equipment: {}", equipment.name); // Poll sensors, interfaces, and neighbors in parallel let (sensor_result, interface_result, neighbor_result) = tokio::join!( executor.poll_sensors(&equipment), executor.poll_interfaces(&equipment), executor.poll_neighbors(&equipment) ); if let Err(e) = sensor_result { error!("Failed to poll sensors for {}: {}", equipment.name, e); } if let Err(e) = interface_result { error!("Failed to poll interfaces for {}: {}", equipment.name, e); } if let Err(e) = neighbor_result { error!("Failed to poll neighbors for {}: {}", equipment.name, e); } // Update last poll time if let Err(e) = storage.update_last_poll_time(&equipment.id) { error!("Failed to update last poll time: {}", e); } }); tasks.push(task); } // Wait for all polling tasks to complete for task in tasks { if let Err(e) = task.await { error!("Polling task failed: {}", e); } } Ok(()) } async fn check_and_update(&self) { info!("Checking for agent updates"); // Run version check in blocking thread to avoid blocking event loop let result = tokio::task::spawn_blocking(crate::version::perform_self_update).await; match result { Ok(Ok(true)) => { info!("Update initiated, container will restart with new version"); // perform_self_update calls std::process::exit(0), so we won't reach here } Ok(Ok(false)) => { info!("Already running latest version"); } Ok(Err(e)) => { warn!("Failed to perform self-update: {}", e); } Err(e) => { error!("Update check task failed: {}", e); } } } }