From 305414b547ef8312178ead4c190c6b57fee45fc1 Mon Sep 17 00:00:00 2001 From: Gabriel Simmer Date: Tue, 24 Oct 2023 13:03:16 +0100 Subject: [PATCH] Track world occupants, visits and favorites --- src/main.rs | 115 ++++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 85 insertions(+), 30 deletions(-) diff --git a/src/main.rs b/src/main.rs index a22b261..2ee8938 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,25 +13,44 @@ use axum::{ use lazy_static::lazy_static; use maud::html; use maud::Markup; -use prometheus::core::{MetricVec, AtomicF64}; -use prometheus::{register_gauge_vec, Encoder, GaugeVec, TextEncoder, Registry, Opts}; +use prometheus::{ + register_gauge_vec, register_int_counter_vec, Encoder, GaugeVec, IntCounterVec, TextEncoder, +}; 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"], + &["instance", "world", "name", "group"], ) - .unwrap(); + .unwrap(); static ref VRCDN_VIEWERS: GaugeVec = register_gauge_vec!( "vrcdn_viewers", "Current number viewers according to VRCDN's API.", - &["region"], + &["region", "group"], ) - .unwrap(); + .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(); } #[derive(Debug)] @@ -116,9 +135,9 @@ struct VrcGroupInstance { #[derive(Clone, Debug, Deserialize)] struct VrcWorldData { - favorites: f64, + favorites: u64, version: f64, - visits: f64, + visits: u64, popularity: f64, heat: f64, #[serde(rename = "publicOccupants")] @@ -150,10 +169,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), } } @@ -171,9 +190,11 @@ 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()); + if config.groups.is_some() { for (name, group) in config.groups.unwrap() { let _ = group_metrics(&client, &auth_cookie, name, group).await; @@ -181,7 +202,7 @@ async fn metrics(config: Config) -> Result, WsError> { } if config.worlds.is_some() { for (name, id) in config.worlds.unwrap() { - let _ = world_metrics(&client, &auth_cookie, name, id); + let _ = world_metrics(&client, &auth_cookie, name, id).await; } } @@ -191,7 +212,12 @@ async fn metrics(config: Config) -> Result, WsError> { Ok(buffer) } -async fn group_metrics(client: &reqwest::Client, auth_cookie: &String, name: String, group: VrcGroup) -> Result<(), WsError> { +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 @@ -241,15 +267,17 @@ async fn group_metrics(client: &reqwest::Client, auth_cookie: &String, name: Str let data: VrcInstanceData = req.json().await?; 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)); } if group.vrcdn.is_some() { - let vrcdn_url = format!( - "https://api.vrcdn.live/v1/viewers/{}", - group.vrcdn.unwrap() - ); + let vrcdn_url = format!("https://api.vrcdn.live/v1/viewers/{}", group.vrcdn.unwrap()); let req = client .get(vrcdn_url) .header( @@ -262,7 +290,7 @@ async fn group_metrics(client: &reqwest::Client, auth_cookie: &String, name: Str let vrcdn_data: VrCdnData = req.json().await.unwrap(); for region in vrcdn_data.viewers { VRCDN_VIEWERS - .with_label_values(&[®ion.region]) + .with_label_values(&[®ion.region, &name]) .set(region.total); } } @@ -270,14 +298,15 @@ async fn group_metrics(client: &reqwest::Client, auth_cookie: &String, name: Str 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 - ); +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 world_data: VrcWorldData = client + let req = client .get(url) .header( USER_AGENT, @@ -285,12 +314,38 @@ async fn world_metrics(client: &reqwest::Client, auth_cookie: &String, name: Str ) .header("Cookie", auth_cookie) .send() - .await?.json().await?; + .await?; + let world_data: VrcWorldData = req.json().await?; for instance in world_data.instances { PLAYER_COUNT - .with_label_values(&[&instance.0, &id, &name]) + .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); + + 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(()) }