Track world occupants, visits and favorites

src/main.rs

@@ -13,8 +13,9 @@ use axum::{
 use lazy_static::lazy_static;
 use maud::html;
 use maud::Markup;
-use prometheus::core::{MetricVec, AtomicF64};
+use prometheus::{
-use prometheus::{register_gauge_vec, Encoder, GaugeVec, TextEncoder, Registry, Opts};
+    register_gauge_vec, register_int_counter_vec, Encoder, GaugeVec, IntCounterVec, TextEncoder,
+};
 use reqwest::header::USER_AGENT;
 use serde::Deserialize;
 use url::Url;
@@ -23,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();
 }
@@ -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")]
@@ -171,9 +190,11 @@ async fn homepage() -> Markup {
 async fn metrics(config: Config) -> Result<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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!(
+        let vrcdn_url = format!("https://api.vrcdn.live/v1/viewers/{}", group.vrcdn.unwrap());
-            "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(&[&region.region])
+                .with_label_values(&[&region.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> {
+async fn world_metrics(
-    let api_url = format!(
+    client: &reqwest::Client,
-        "https://api.vrchat.cloud/api/1/worlds/{}",
+    auth_cookie: &String,
-        &id
+    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
-    let world_data: VrcWorldData = 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(())
 }