Expand Metrics Tracked (#1)

Reviewed-on: #1

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<String>,
+    /// Groups we want to track.
-    group: Option<String>,
+    groups: Option<HashMap<String, VrcGroup>>,
+    /// List of worlds we want to track.
+    worlds: Option<HashMap<String, String>>,
+    #[serde(skip)]
     vrchat_token: Option<String>,
 }
 
+#[derive(Clone, Debug, Deserialize)]
+struct VrcGroup {
+    id: String,
+    vrcdn: Option<String>,
+}
+
 #[derive(Clone, Debug, Deserialize)]
 struct VrcInstance {
     /// Instance ID.
@@ -79,7 +109,6 @@ struct VrcInstance {
     /// Raw location
     location: Option<String>,
     /// Custom name for the instance.
-    #[serde(skip)]
     name: Option<String>,
 }
 
@@ -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<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());
-    let mut instances: Vec<VrcInstance> = vec![];
 
-    // Check if we can fetch instances from a group if set.
+    if config.groups.is_some() {
-    if config.group.is_some() {
+        for (name, group) in config.groups.unwrap() {
-        let api_url = format!(
+            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",
-            config.group.unwrap()
+        group.id
     );
-        let url = Url::parse(&api_url).unwrap();
+    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)
+        .header("Cookie", auth_cookie)
         .send()
         .await?;
+
     let data: Vec<VrcGroupInstance> = req.json().await?;
-        instances = data.into_iter().map(|f| {
+    let instances: Vec<VrcInstance> = 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  }
+            VrcInstance {
-        }).collect();
+                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,21 +260,24 @@ async fn metrics(config: Config) -> Result<Vec<u8>, 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!(
+    if group.vrcdn.is_some() {
-        "https://api.vrcdn.live/v1/viewers/{}",
+        let vrcdn_url = format!("https://api.vrcdn.live/v1/viewers/{}", group.vrcdn.unwrap());
-        config.vrcdn.unwrap()
-    );
         let req = client
             .get(vrcdn_url)
             .header(
@@ -211,15 +288,64 @@ async fn metrics(config: Config) -> Result<Vec<u8>, WsError> {
             .await
             .unwrap();
         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);
         }
+    }
 
-    let metric_families = prometheus::gather();
+    Ok(())
-    let mut buffer = vec![];
+}
-    encoder.encode(&metric_families, &mut buffer).unwrap();
+
-    Ok(buffer)
+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(url)
+        .header(
+            USER_AGENT,
+            "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)",
+        )
+        .header("Cookie", auth_cookie)
+        .send()
+        .await?;
+    let world_data: VrcWorldData = req.json().await?;
+
+    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);
+
+    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(())
 }