Expand metrics to include world data

src/main.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::error::Error;
 use std::{env, fmt, fs};
 
@@ -12,7 +13,8 @@ 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;
@@ -65,11 +67,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 +90,6 @@ struct VrcInstance {
     /// Raw location
     location: Option<String>,
     /// Custom name for the instance.
-    #[serde(skip)]
     name: Option<String>,
 }
 
@@ -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();
@@ -148,30 +174,52 @@ async fn metrics(config: Config) -> Result<Vec<u8>, WsError> {
     let encoder = TextEncoder::new();
     let client = reqwest::Client::new();
     let auth_cookie = format!("auth={}", &config.vrchat_token.unwrap());
-    let mut instances: Vec<VrcInstance> = vec![];
+    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);
+        }
+    }
 
-    // Check if we can fetch instances from a group if set.
+    let metric_families = prometheus::gather();
-    if config.group.is_some() {
+    let mut buffer = vec![];
-        let api_url = format!(
+    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,20 +234,21 @@ 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])
             .set(data.user_count.unwrap_or(0 as f64));
     }
 
+    if group.vrcdn.is_some() {
         let vrcdn_url = format!(
             "https://api.vrcdn.live/v1/viewers/{}",
-        config.vrcdn.unwrap()
+            group.vrcdn.unwrap()
         );
         let req = client
             .get(vrcdn_url)
@@ -211,15 +260,37 @@ 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])
                 .set(region.total);
         }
-
+    }
-    let metric_families = prometheus::gather();
+
-    let mut buffer = vec![];
+    Ok(())
-    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 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?.json().await?;
+
+    for instance in world_data.instances {
+        PLAYER_COUNT
+            .with_label_values(&[&instance.0, &id, &name])
+            .set(instance.1);
+    }
+    Ok(())
 }