use crate::analytics::AnalyticsQueue; use crate::config::Config; use crate::db::Db; use crate::state::AppState; use crate::web::create_router; use std::path::PathBuf; use std::time::Instant; use tracing::info; async fn shutdown_signal() { let ctrl_c = async { tokio::signal::ctrl_c() .await .expect("failed to install Ctrl+C handler"); }; #[cfg(unix)] let terminate = async { tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) .expect("failed to install signal handler") .recv() .await; }; #[cfg(not(unix))] let terminate = std::future::pending::<()>(); tokio::select! { _ = ctrl_c => {}, _ = terminate => {}, } } pub async fn run( host: Option, port: Option, data_dir: Option, mut config: Config, ) -> Result<(), Box> { if let Some(h) = host { config.host = h; } if let Some(p) = port { config.port = p; } if let Some(d) = data_dir { config.data_dir = PathBuf::from(d); } info!("Starting BZOD server on {}:{}", config.host, config.port); info!("Database directory: {:?}", config.data_dir); // Init DBs let db = Db::init(&config)?; let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); let mut join_handles = Vec::new(); // Init Queue let (queue, analytics_handle) = AnalyticsQueue::new(db.clone(), 1000, shutdown_rx.clone()); join_handles.push(("analytics_worker", analytics_handle)); // Spawn background tasks let link_checker_db = db.clone(); let link_checker_interval = config.link_check_interval_mins; let rx = shutdown_rx.clone(); join_handles.push(( "link_checker", tokio::spawn(async move { crate::jobs::run_link_checker(link_checker_db, link_checker_interval, rx).await; }), )); let aggregator_db = db.clone(); let aggregator_interval = config.aggregation_interval_mins; let rx = shutdown_rx.clone(); join_handles.push(( "aggregator", tokio::spawn(async move { crate::jobs::run_aggregator(aggregator_db, aggregator_interval, rx).await; }), )); let retention_db = db.clone(); let retention_days = config.data_retention_days; let rx = shutdown_rx.clone(); join_handles.push(( "retention_cleaner", tokio::spawn(async move { crate::jobs::run_retention_cleaner(retention_db, retention_days, rx).await; }), )); // Spawn optional backup scheduler let backup_db = db.clone(); let backup_config = config.clone(); let rx = shutdown_rx.clone(); join_handles.push(( "backup_scheduler", tokio::spawn(async move { crate::jobs::backup::run_backup_scheduler(backup_db, backup_config, rx).await; }), )); let expiry_db = db.clone(); let rx = shutdown_rx.clone(); join_handles.push(( "expiry_checker", tokio::spawn(async move { crate::jobs::run_expiry_checker(expiry_db, rx).await; }), )); let reconcile_db = db.clone(); let reconcile_interval_hours = { let conn = db.system.lock().unwrap(); conn.query_row( "SELECT value FROM settings WHERE key = 'quota_reconcile_interval_hours';", [], |row| row.get::<_, String>(0), ) .ok() .and_then(|val| val.parse::().ok()) .unwrap_or(24) }; let rx = shutdown_rx.clone(); join_handles.push(( "quota_reconciliation", tokio::spawn(async move { crate::jobs::run_quota_reconciliation(reconcile_db, reconcile_interval_hours, rx).await; }), )); let state = AppState { admin_db: db.admin.clone(), system_db: db.system.clone(), users_db: db.users.clone(), user_dbs: std::sync::Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())), db: db.clone(), config: config.clone(), analytics_queue: queue, start_time: Instant::now(), }; // Run axum server let router = create_router(state); let addr = format!("{}:{}", config.host, config.port); let listener = tokio::net::TcpListener::bind(&addr).await?; info!("Listening for requests on http://{}", addr); axum::serve(listener, router) .with_graceful_shutdown(async move { shutdown_signal().await; info!("Shutdown signal received"); info!("Stopping HTTP server..."); let _ = shutdown_tx.send(true); }) .await?; info!("Stopping background workers..."); let timeout_duration = std::time::Duration::from_secs(10); let deadline = tokio::time::Instant::now() + timeout_duration; for (name, handle) in join_handles { match tokio::time::timeout_at(deadline, handle).await { Ok(Ok(_)) => {} Ok(Err(e)) => tracing::error!("Background task '{}' panicked: {:?}", name, e), Err(_) => tracing::warn!("Background task did not terminate: {}", name), } } info!("Background workers stopped"); info!("BZOD shutdown complete"); Ok(()) }