main.rs
1 mod auth; 2 mod collect_info; 3 mod config; 4 mod db; 5 mod endpoints; 6 mod models; 7 mod alerts; 8 mod logging; 9 10 use alerts::check_alerts; 11 use axum::{ 12 routing::{delete, get, post}, Router 13 }; 14 use std::net::SocketAddr; 15 use db::db_update; 16 use endpoints::{ 17 add_alert, add_notif_method, delete_alert, delete_notif_method, get_alert_vars, get_alerts, get_container_logs, get_notif_methods, historical_data, req_info, serve_static, ws_handler_d, ws_handler_g, ws_handler_p 18 }; 19 use log::{error, info, debug}; 20 use std::sync::{Arc, Mutex}; 21 use sysinfo::System; 22 use tokio::{self, time::Duration}; 23 use tower_http::compression::CompressionLayer; 24 25 async fn sys_refresh(sys: Arc<Mutex<System>>, update_interval: u64) { 26 loop { 27 { 28 let mut sys_write = sys.lock().unwrap(); 29 sys_write.refresh_cpu_usage(); 30 sys_write.refresh_memory(); 31 } 32 tokio::time::sleep(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL).await; 33 { 34 let mut sys_write = sys.lock().unwrap(); 35 sys_write.refresh_cpu_usage(); 36 sys_write.refresh_memory(); 37 } 38 tokio::time::sleep(Duration::from_secs(update_interval)).await; 39 } 40 } 41 42 #[tokio::main] 43 async fn main() { 44 logging::setup(); 45 46 // Parse command line arguments 47 let config = config::parse_config(); 48 let update_interval = config.update_interval; 49 info!("Update interval: {} seconds", update_interval); 50 51 // Create system instance for the main thread and web API 52 let sys = System::new_all(); 53 54 let shared_sys = Arc::new(Mutex::new(sys)); 55 56 let bg_sys = shared_sys.clone(); 57 let db_sys = shared_sys.clone(); 58 59 // System refresh background task with restart on panic 60 tokio::spawn(async move { 61 loop { 62 let result = tokio::task::spawn(sys_refresh(bg_sys.clone(), update_interval)).await; 63 match result { 64 Err(e) => { 65 error!("System refresh task panicked: {}", e); 66 tokio::time::sleep(Duration::from_secs(5)).await; 67 info!("Restarting system refresh task"); 68 // Continue the loop to restart the task 69 } 70 _ => { 71 break; // This should not happen as sys_refresh runs indefinitely 72 } 73 } 74 } 75 }); 76 debug!("System refresh background task started"); 77 78 // Database update background task with restart on panic 79 let db_path = config.db_path.clone(); 80 tokio::spawn(async move { 81 loop { 82 let db_path = db_path.clone(); 83 let db_sys = db_sys.clone(); 84 let result = tokio::task::spawn(async move { db_update(db_sys, &db_path).await }).await; 85 match result { 86 Err(e) => { 87 error!("Database update task panicked: {}", e); 88 tokio::time::sleep(Duration::from_secs(5)).await; 89 info!("Restarting database update task"); 90 // Continue the loop to restart the task 91 } 92 _ => { 93 break; // This should not happen as db_update runs indefinitely 94 } 95 } 96 } 97 }); 98 debug!("Database update background task started"); 99 100 let db_path = config.db_path.clone(); 101 // Check alerts background task with restart on panic 102 tokio::spawn(async move { 103 loop { 104 let db_path = db_path.clone(); 105 let result = tokio::task::spawn(async move { check_alerts(&db_path).await }).await; 106 match result { 107 Err(e) => { 108 error!("Check alerts task panicked: {}", e); 109 tokio::time::sleep(Duration::from_secs(5)).await; 110 info!("Restarting check alerts task"); 111 // Continue the loop to restart the task 112 } 113 _ => { 114 break; // This should not happen as db_update runs indefinitely 115 } 116 } 117 } 118 }); 119 debug!("Alerts checking background task started"); 120 121 let mut app = Router::new() 122 .route("/", get(serve_static)) 123 .route("/favicon.png", get(serve_static)) 124 .route("/Inter-Regular.woff", get(serve_static)) 125 .route("/Inter-Regular.woff2", get(serve_static)) 126 .route("/RobotoMono-Regular.woff", get(serve_static)) 127 .route("/RobotoMono-Regular.woff2", get(serve_static)) 128 .route("/auth", get(serve_static)) 129 .route("/auth", post(auth::auth_handler)) 130 .route("/ws/g", get(ws_handler_g)) 131 .route("/ws/p", get(ws_handler_p)) 132 .route("/ws/d", get(ws_handler_d)) 133 .route("/container_logs/{continer_id}", get(get_container_logs)) 134 .route("/reqinfo", get(req_info)) 135 .route("/api/historical", get(historical_data)) 136 .route("/api/notif_methods", post(add_notif_method)) 137 .route("/api/notif_methods", get(get_notif_methods)) 138 .route("/api/notif_methods/{id}", delete(delete_notif_method)) 139 .route("/api/alerts", post(add_alert)) 140 .route("/api/alerts", get(get_alerts)) 141 .route("/api/alerts/{id}", delete(delete_alert)) 142 .route("/api/alert_vars", get(get_alert_vars)) 143 .fallback(get(serve_static)) 144 .with_state((shared_sys, config.clone())); 145 146 if let Some(_) = &config.password_hash { 147 app = auth::apply_auth_middleware(app, config.clone()); 148 info!("Running with authentication"); 149 } else { 150 info!("Running without authentication"); 151 } 152 app = app.layer(CompressionLayer::new()); 153 154 info!( 155 "Server running on http://{}:{}", 156 config.address, config.port 157 ); 158 159 let listener = tokio::net::TcpListener::bind(config.socket_address()) 160 .await 161 .unwrap(); 162 axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>()).await.unwrap(); 163 }