use anyhow::{Context, Result}; use clap::{Parser, Subcommand}; use mirror_tool::{auth::Auth, config::Config, control::Controller, web}; use std::future::IntoFuture; use std::path::PathBuf; use tokio_util::sync::CancellationToken; #[derive(Parser)] #[command(version, about = "Async TCP mirror and ForcAD-aware upstream router")] struct Cli { #[arg(short, long, default_value = "mirror.toml", global = true)] config: PathBuf, #[command(subcommand)] command: Option, } #[derive(Subcommand)] enum Command { /// Validate TOML and any existing saved topology without binding or network I/O. Check, /// Print a portable sample configuration (no secrets). Example, } fn main() -> Result<()> { let runtime = tokio::runtime::Builder::new_multi_thread() .enable_all() .build()?; let result = runtime.block_on(run()); // A filesystem worker can remain blocked in the OS after the async deadline. // Do not make runtime destruction an unbounded second shutdown phase. runtime.shutdown_timeout(std::time::Duration::from_millis(100)); result } async fn run() -> Result<()> { tracing_subscriber::fmt() .with_env_filter( tracing_subscriber::EnvFilter::try_from_default_env() .unwrap_or_else(|_| "mirror_tool=info".into()), ) .init(); let cli = Cli::parse(); if matches!(cli.command, Some(Command::Example)) { print!("{}", include_str!("../mirror.example.toml")); return Ok(()); } let config = Config::load(&cli.config)?; if matches!(cli.command, Some(Command::Check)) { let topology = mirror_tool::store::Topology::load(&config)?; println!( "Configuration valid: {} teams, {} services", topology.teams.len(), topology.services.len() ); return Ok(()); } let auth = Auth::from_env(&config.web)?; let listen = config.web.listen; // Reserve control listener before starting data listeners, so startup cannot half-succeed. let listener = tokio::net::TcpListener::bind(listen) .await .context("bind web listener")?; let controller = Controller::start(config).await?; let cancel = CancellationToken::new(); let monitor = tokio::spawn(controller.clone().monitor(cancel.clone())); let board = controller .config .scoreboard .clone() .map(|config| tokio::spawn(controller.board.clone().run(config, cancel.clone()))); let router = web::router(controller.clone(), auth, cancel.clone()); tracing::info!(%listen,"Mirror Tool ready"); let server = axum::serve( listener, router.into_make_service_with_connect_info::(), ); let signal_cancel = cancel.clone(); let server = server .with_graceful_shutdown(signal_cancel.cancelled_owned()) .into_future(); tokio::pin!(server); let result = tokio::select! { result = &mut server => result, _ = shutdown_signal() => { cancel.cancel(); // HTTP body extraction has no idle deadline during normal operation. // Shutdown alone bounds it, while serialized control transactions finish first. let shutdown = async { controller.shutdown().await; let _ = monitor.await; if let Some(board) = board { let _ = board.await; } (&mut server).await }; match tokio::time::timeout( std::time::Duration::from_millis(controller.config.network.drain_timeout_ms) + std::time::Duration::from_secs(1), shutdown).await { Ok(result) => result, Err(_) => { tracing::warn!("shutdown deadline reached"); return Ok(()); } } } }; cancel.cancel(); controller.shutdown().await; result.context("HTTP server")?; Ok(()) } async fn shutdown_signal() { #[cfg(unix)] { let mut term = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) .expect("install SIGTERM handler"); tokio::select! {_=tokio::signal::ctrl_c()=>{},_=term.recv()=>{}} } #[cfg(not(unix))] { let _ = tokio::signal::ctrl_c().await; } }