diff --git a/src/main.rs b/src/main.rs index 957d17a..16f1f69 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,41 +14,51 @@ use std::sync::Arc; use std::time::Duration; use tokio::time::interval; +/// Command line arguments structure #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] struct Args { + /// Interval in seconds between each query to the server #[arg(short, long, default_value = "5", help="Interval in seconds between each query to the server")] update_interval: u64, + /// Hostname and port of the server to query #[arg(short, long, help="Hostname and port of the server to query")] endpoint: String, + /// File containing the bearer token to use for authentication #[arg(short, long, help="File containing the bearer token to use for authentication")] token_file: Option, + /// Allow insecure connections (e.g., to a server with a self-signed certificate) #[arg(short, long, help="Allow insecure connections (e.g., to a server with a self-signed certificate)")] allow_insecure: bool, + /// Address:Port to which the server will listen #[arg(short, long, help="Address:Port to which the server will listen", default_value = "127.0.0.1:3030")] listen: String, } +/// Structure for the query body sent to the server #[derive(Serialize)] struct QueryBody { function: String, } +/// Structure for the server response #[derive(Deserialize)] struct ServerResponse { data: ServerData, } +/// Structure for the server data within the server response #[derive(Deserialize)] #[serde(rename_all = "camelCase")] struct ServerData { server_game_state: ServerGameState, } +/// Structure for the server game state within the server data #[derive(Deserialize)] #[serde(rename_all = "camelCase")] struct ServerGameState { @@ -58,6 +68,7 @@ struct ServerGameState { average_tick_rate: f64, } +/// Structure for the metrics to be collected #[derive(Clone)] struct Metrics { num_connected_players: Gauge, @@ -67,6 +78,7 @@ struct Metrics { } impl Metrics { + /// Creates a new instance of `Metrics` fn new() -> Self { Metrics { num_connected_players: Gauge::new("num_connected_players", "Number of connected players").unwrap(), @@ -76,6 +88,7 @@ impl Metrics { } } + /// Updates the metrics with the provided game state fn update(&self, game_state: &ServerGameState) { self.num_connected_players.set(game_state.num_connected_players as f64); self.tech_tier.set(game_state.tech_tier as f64); @@ -84,8 +97,10 @@ impl Metrics { } } +/// Shared state type alias type SharedState = Arc<(Metrics, Registry)>; +/// Handler for the `/metrics` endpoint async fn metrics_handler(State(state): State) -> impl IntoResponse { let encoder = TextEncoder::new(); let mut buffer = Vec::new(); @@ -93,20 +108,26 @@ async fn metrics_handler(State(state): State) -> impl IntoResponse String::from_utf8(buffer).unwrap() } +/// Main function #[tokio::main] async fn main() -> Result<(), Box> { + // Parse command line arguments let args = Args::parse(); + // Create a new registry and metrics instance let registry = Registry::new(); let metrics = Arc::new(Metrics::new()); + // Register metrics with the registry registry.register(Box::new(metrics.num_connected_players.clone())).unwrap(); registry.register(Box::new(metrics.tech_tier.clone())).unwrap(); registry.register(Box::new(metrics.total_game_duration.clone())).unwrap(); registry.register(Box::new(metrics.average_tick_rate.clone())).unwrap(); + // Create shared state let shared_state: SharedState = Arc::new(((*metrics).clone(), registry)); + // Clone metrics for use in the update loop let metrics_clone = Arc::clone(&metrics); let update_interval = Duration::from_secs(args.update_interval); @@ -122,6 +143,7 @@ async fn main() -> Result<(), Box> { let query_endpoint = format!("https://{}/api/v1", args.endpoint); + // Spawn a task to periodically query the server and update metrics tokio::spawn(async move { let mut interval = interval(update_interval); let query_body = QueryBody { @@ -150,10 +172,12 @@ async fn main() -> Result<(), Box> { } }); + // Build the application router let app = Router::new() .route("/metrics", get(metrics_handler)) .with_state(shared_state); + // Start the server let addr = std::net::SocketAddr::from_str(&args.listen)?; println!("Listening on {}", addr); axum::Server::bind(&addr) @@ -161,4 +185,4 @@ async fn main() -> Result<(), Box> { .await?; Ok(()) -} \ No newline at end of file +}