diff --git a/src/main.rs b/src/main.rs index 60f8ebb..a22b261 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::error::Error; use std::{env, fmt, fs}; @@ -12,24 +13,25 @@ use axum::{ use lazy_static::lazy_static; use maud::html; use maud::Markup; -use prometheus::{register_gauge_vec, Encoder, GaugeVec, TextEncoder}; +use prometheus::core::{MetricVec, AtomicF64}; +use prometheus::{register_gauge_vec, Encoder, GaugeVec, TextEncoder, Registry, Opts}; use reqwest::header::USER_AGENT; use serde::Deserialize; use url::Url; -lazy_static! { +lazy_static!{ static ref PLAYER_COUNT: GaugeVec = register_gauge_vec!( "vrchat_playercount", "Current number of players in instance.", &["instance", "world", "name"], ) - .unwrap(); + .unwrap(); static ref VRCDN_VIEWERS: GaugeVec = register_gauge_vec!( "vrcdn_viewers", "Current number viewers according to VRCDN's API.", &["region"], ) - .unwrap(); + .unwrap(); } #[derive(Debug)] @@ -65,11 +67,20 @@ impl fmt::Display for WsError { #[derive(Clone, Debug, Deserialize)] struct Config { - vrcdn: Option, - group: Option, + /// Groups we want to track. + groups: Option>, + /// List of worlds we want to track. + worlds: Option>, + #[serde(skip)] vrchat_token: Option, } +#[derive(Clone, Debug, Deserialize)] +struct VrcGroup { + id: String, + vrcdn: Option, +} + #[derive(Clone, Debug, Deserialize)] struct VrcInstance { /// Instance ID. @@ -79,7 +90,6 @@ struct VrcInstance { /// Raw location location: Option, /// Custom name for the instance. - #[serde(skip)] name: Option, } @@ -104,6 +114,22 @@ struct VrcGroupInstance { location: String, } +#[derive(Clone, Debug, Deserialize)] +struct VrcWorldData { + favorites: f64, + version: f64, + visits: f64, + popularity: f64, + heat: f64, + #[serde(rename = "publicOccupants")] + public_occupants: f64, + #[serde(rename = "privateOccupants")] + private_occupants: f64, + #[serde(rename = "occupants")] + total_occupants: f64, + instances: Vec<(String, f64)>, +} + #[tokio::main] async fn main() -> Result<(), ()> { let content = fs::read_to_string("config.toml").unwrap(); @@ -124,10 +150,10 @@ async fn main() -> Result<(), ()> { async fn metrics_handler(State(config): State) -> Result { match metrics(config).await { Ok(b) => Ok(Response::builder() - .header("content-type", "text/plain") - .status(StatusCode::OK) - .body(Full::from(b)) - .unwrap()), + .header("content-type", "text/plain") + .status(StatusCode::OK) + .body(Full::from(b)) + .unwrap()), Err(_e) => Err(StatusCode::INTERNAL_SERVER_ERROR), } } @@ -148,30 +174,52 @@ async fn metrics(config: Config) -> Result, WsError> { let encoder = TextEncoder::new(); let client = reqwest::Client::new(); let auth_cookie = format!("auth={}", &config.vrchat_token.unwrap()); - let mut instances: Vec = vec![]; - - // Check if we can fetch instances from a group if set. - if config.group.is_some() { - let api_url = format!( - "https://api.vrchat.cloud/api/1/groups/{}/instances", - config.group.unwrap() - ); - let url = Url::parse(&api_url).unwrap(); - let req = client - .get(url) - .header( - USER_AGENT, - "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)", - ) - .header("Cookie", &auth_cookie) - .send() - .await?; - let data: Vec = req.json().await?; - instances = data.into_iter().map(|f| { - let spl: Vec<&str> = f.location.split(":").collect(); - VrcInstance{ instance: Some(spl[0].to_owned()), world: Some(spl[1].to_owned()), location: Some(f.location), name: None } - }).collect(); + if config.groups.is_some() { + for (name, group) in config.groups.unwrap() { + let _ = group_metrics(&client, &auth_cookie, name, group).await; + } } + if config.worlds.is_some() { + for (name, id) in config.worlds.unwrap() { + let _ = world_metrics(&client, &auth_cookie, name, id); + } + } + + let metric_families = prometheus::gather(); + let mut buffer = vec![]; + encoder.encode(&metric_families, &mut buffer).unwrap(); + Ok(buffer) +} + +async fn group_metrics(client: &reqwest::Client, auth_cookie: &String, name: String, group: VrcGroup) -> Result<(), WsError> { + let instance_list_url = format!( + "https://api.vrchat.cloud/api/1/groups/{}/instances", + group.id + ); + let url = Url::parse(&instance_list_url).unwrap(); + let req = client + .get(url) + .header( + USER_AGENT, + "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)", + ) + .header("Cookie", auth_cookie) + .send() + .await?; + + let data: Vec = req.json().await?; + let instances: Vec = data + .into_iter() + .map(|f| { + let spl: Vec<&str> = f.location.split(":").collect(); + VrcInstance { + instance: Some(spl[0].to_owned()), + world: Some(spl[1].to_owned()), + location: Some(f.location), + name: None, + } + }) + .collect(); for instance in instances { let api_url = format!( @@ -186,40 +234,63 @@ async fn metrics(config: Config) -> Result, WsError> { USER_AGENT, "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)", ) - .header("Cookie", &auth_cookie) + .header("Cookie", auth_cookie) .send() .await?; let data: VrcInstanceData = req.json().await?; - let name = instance.name.unwrap_or(instance.location.unwrap()); + let instance_name = instance.name.unwrap_or(instance.location.unwrap()); PLAYER_COUNT .with_label_values(&[&instance.world.unwrap(), &instance.instance.unwrap(), &name]) .set(data.user_count.unwrap_or(0 as f64)); } - let vrcdn_url = format!( - "https://api.vrcdn.live/v1/viewers/{}", - config.vrcdn.unwrap() + if group.vrcdn.is_some() { + let vrcdn_url = format!( + "https://api.vrcdn.live/v1/viewers/{}", + group.vrcdn.unwrap() + ); + let req = client + .get(vrcdn_url) + .header( + USER_AGENT, + "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)", + ) + .send() + .await + .unwrap(); + let vrcdn_data: VrCdnData = req.json().await.unwrap(); + for region in vrcdn_data.viewers { + VRCDN_VIEWERS + .with_label_values(&[®ion.region]) + .set(region.total); + } + } + + Ok(()) +} + +async fn world_metrics(client: &reqwest::Client, auth_cookie: &String, name: String, id: String) -> Result<(), WsError> { + let api_url = format!( + "https://api.vrchat.cloud/api/1/worlds/{}", + &id ); - let req = client - .get(vrcdn_url) + let url = Url::parse(&api_url).unwrap(); + + let world_data: VrcWorldData = client + .get(url) .header( USER_AGENT, "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)", ) + .header("Cookie", auth_cookie) .send() - .await - .unwrap(); - let vrcdn_data: VrCdnData = req.json().await.unwrap(); + .await?.json().await?; - for region in vrcdn_data.viewers { - VRCDN_VIEWERS - .with_label_values(&[®ion.region]) - .set(region.total); + for instance in world_data.instances { + PLAYER_COUNT + .with_label_values(&[&instance.0, &id, &name]) + .set(instance.1); } - - let metric_families = prometheus::gather(); - let mut buffer = vec![]; - encoder.encode(&metric_families, &mut buffer).unwrap(); - Ok(buffer) + Ok(()) }