vrchat-prometheus-adapter/src/main.rs

use std::collections::HashMap;
use std::error::Error;
use std::{env, fmt, fs};

use axum::extract::State;
use axum::{
    body::Full,
    http::{Response, StatusCode},
    response::IntoResponse,
    routing::get,
    Router,
};
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 reqwest::header::USER_AGENT;
use serde::Deserialize;
use url::Url;

lazy_static!{
    static ref PLAYER_COUNT: GaugeVec = register_gauge_vec!(
        "vrchat_playercount",
        "Current number of players in instance.",
        &["instance", "world", "name"],
    )
        .unwrap();
    static ref VRCDN_VIEWERS: GaugeVec = register_gauge_vec!(
        "vrcdn_viewers",
        "Current number viewers according to VRCDN's API.",
        &["region"],
    )
        .unwrap();
}

#[derive(Debug)]
enum WsError {
    Reqwest(reqwest::Error),
    Url(url::ParseError),
    Custom(String),
}

impl From<url::ParseError> for WsError {
    fn from(value: url::ParseError) -> Self {
        Self::Url(value)
    }
}

impl From<reqwest::Error> for WsError {
    fn from(value: reqwest::Error) -> Self {
        Self::Reqwest(value)
    }
}

impl Error for WsError {}

impl fmt::Display for WsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            WsError::Reqwest(e) => write!(f, "Reqwest error: {}", e),
            WsError::Custom(e) => write!(f, "Error: {}", e),
            WsError::Url(_) => todo!(),
        }
    }
}

#[derive(Clone, Debug, Deserialize)]
struct Config {
    /// Groups we want to track.
    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.
    instance: Option<String>,
    /// World ID.
    world: Option<String>,
    /// Raw location
    location: Option<String>,
    /// Custom name for the instance.
    name: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
struct VrcInstanceData {
    #[serde(rename = "userCount")]
    user_count: Option<f64>,
}

#[derive(Clone, Debug, Deserialize)]
struct VrCdnData {
    viewers: Vec<VrCdnRegion>,
}
#[derive(Clone, Debug, Deserialize)]
struct VrCdnRegion {
    region: String,
    total: f64,
}

#[derive(Clone, Debug, Deserialize)]
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();
    let mut config: Config = toml::from_str(&content).unwrap();
    config.vrchat_token = Some(env::var("VRCHAT_AUTH_TOKEN").unwrap());

    let app = Router::new()
        .route("/", get(homepage))
        .route("/metrics", get(metrics_handler))
        .with_state(config);

    let _ = axum::Server::bind(&"0.0.0.0:6534".parse().unwrap())
        .serve(app.into_make_service())
        .await;
    Ok(())
}

async fn metrics_handler(State(config): State<Config>) -> Result<impl IntoResponse, StatusCode> {
    match metrics(config).await {
        Ok(b) => Ok(Response::builder()
                    .header("content-type", "text/plain")
                    .status(StatusCode::OK)
                    .body(Full::from(b))
                    .unwrap()),
        Err(_e) => Err(StatusCode::INTERNAL_SERVER_ERROR),
    }
}

async fn homepage() -> Markup {
    html! {
        body {
            main {
                a href = "/metrics" { "metrics" }
            }
        }
    }
}

async fn metrics(config: Config) -> Result<Vec<u8>, WsError> {
    PLAYER_COUNT.reset();
    VRCDN_VIEWERS.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;
        }
    }
    if config.worlds.is_some() {
        for (name, id) in config.worlds.unwrap() {
            let _ = world_metrics(&client, &auth_cookie, name, id);
        }
    }

    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<VrcGroupInstance> = req.json().await?;
    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,
            }
        })
        .collect();

    for instance in instances {
        let api_url = format!(
            "https://api.vrchat.cloud/api/1/instances/{}",
            &instance.location.clone().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: 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])
            .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 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(&[&region.region])
                .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 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(())
}