diff --git a/src/main.rs b/src/main.rs index 60f8ebb..2ee8938 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,7 +13,9 @@ use axum::{ use lazy_static::lazy_static; use maud::html; use maud::Markup; -use prometheus::{register_gauge_vec, Encoder, GaugeVec, TextEncoder}; +use prometheus::{ + register_gauge_vec, register_int_counter_vec, Encoder, GaugeVec, IntCounterVec, TextEncoder, +}; use reqwest::header::USER_AGENT; use serde::Deserialize; use url::Url; @@ -21,13 +24,31 @@ lazy_static! { static ref PLAYER_COUNT: GaugeVec = register_gauge_vec!( "vrchat_playercount", "Current number of players in instance.", - &["instance", "world", "name"], + &["instance", "world", "name", "group"], ) .unwrap(); static ref VRCDN_VIEWERS: GaugeVec = register_gauge_vec!( "vrcdn_viewers", "Current number viewers according to VRCDN's API.", - &["region"], + &["region", "group"], + ) + .unwrap(); + static ref WORLD_VISITS: IntCounterVec = register_int_counter_vec!( + "vrchat_world_visits", + "Number of times a world has been visited.", + &["world", "name"], + ) + .unwrap(); + static ref WORLD_OCCUPANTS: GaugeVec = register_gauge_vec!( + "vrchat_world_occupants", + "Occupants currently in the world", + &["world", "name", "type"] + ) + .unwrap(); + static ref WORLD_FAVORITES: IntCounterVec = register_int_counter_vec!( + "vrchat_world_favorites", + "Number of times a world has been favorited.", + &["world", "name"], ) .unwrap(); } @@ -65,11 +86,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 +109,6 @@ struct VrcInstance { /// Raw location location: Option, /// Custom name for the instance. - #[serde(skip)] name: Option, } @@ -104,6 +133,22 @@ struct VrcGroupInstance { location: String, } +#[derive(Clone, Debug, Deserialize)] +struct VrcWorldData { + favorites: u64, + version: f64, + visits: u64, + 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(); @@ -145,33 +190,62 @@ async fn homepage() -> Markup { async fn metrics(config: Config) -> Result, WsError> { PLAYER_COUNT.reset(); VRCDN_VIEWERS.reset(); + WORLD_VISITS.reset(); 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).await; + } + } + + 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 +260,92 @@ 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]) + .with_label_values(&[ + &instance.world.unwrap(), + &instance.instance.unwrap(), + &instance_name, + &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, &name]) + .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 url = Url::parse(&api_url).unwrap(); let req = client - .get(vrcdn_url) + .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?; + let world_data: VrcWorldData = req.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 current_visits: u64 = match WORLD_VISITS.get_metric_with_label_values(&[&id, &name]) { + Ok(v) => v.get(), + Err(_) => 0, + }; + WORLD_VISITS + .with_label_values(&[&id, &name]) + .inc_by(world_data.visits - current_visits); + let current_favorites: u64 = match WORLD_FAVORITES.get_metric_with_label_values(&[&id, &name]) { + Ok(v) => v.get(), + Err(_) => 0, + }; + WORLD_FAVORITES + .with_label_values(&[&id, &name]) + .inc_by(world_data.favorites - current_favorites); - let metric_families = prometheus::gather(); - let mut buffer = vec![]; - encoder.encode(&metric_families, &mut buffer).unwrap(); - Ok(buffer) + WORLD_OCCUPANTS + .with_label_values(&[&id, &name, "private"]) + .set(world_data.private_occupants); + WORLD_OCCUPANTS + .with_label_values(&[&id, &name, "public"]) + .set(world_data.public_occupants); + WORLD_OCCUPANTS + .with_label_values(&[&id, &name, "total"]) + .set(world_data.total_occupants); + + Ok(()) }