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 maud::html;
use maud::Markup;
use prometheus::{register_gauge_vec, Encoder, GaugeVec, TextEncoder};

use lazy_static::lazy_static;
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 {
    vrcdn: Option<String>,
    instances: Option<Vec<VrcInstance>>,
    vrchat_token: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
struct VrcInstance {
    /// Full invite URL. Takes precedence over instance and world.
    url: Option<String>,
    /// Instance ID.
    instance: Option<String>,
    /// World ID.
    world: Option<String>,
    /// Custom name for the instance.
    name: Option<String>
}

#[derive(Clone, Debug, Deserialize)]
struct VrcInstanceData {
    capacity: Option<f64>,
    n_users: Option<f64>,
    userCount: Option<f64>,
}

#[derive(Clone, Debug, Deserialize)]
struct VrCdnData {
    viewers: Vec<VrCdnRegion>,
}
#[derive(Clone, Debug, Deserialize)]
struct VrCdnRegion {
    region: String,
    total: 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> {
    let encoder = TextEncoder::new();
    let client = reqwest::Client::new();
    let auth_cookie = format!("auth={}", &config.vrchat_token.unwrap());

    // Do work.
    for instance in config.instances.ok_or(WsError::Custom("".to_owned()))? {
        let (instance_id, world_id): (String, String) = if instance.url.is_none() {
            (
                instance.instance.ok_or(WsError::Custom("".to_owned()))?,
                instance.world.ok_or(WsError::Custom("".to_owned()))?,
            )
        } else {
            let url = Url::parse(&instance.url.unwrap_or("".to_owned()))?;
            let hash_query: HashMap<_, _> = url.query_pairs().into_owned().collect();
            (
                hash_query
                    .get("instanceId")
                    .ok_or(WsError::Custom("".to_owned()))?
                    .to_string(),
                hash_query
                    .get("worldId")
                    .ok_or(WsError::Custom("".to_owned()))?
                    .to_string(),
            )
        };
        let api_url = format!(
            "https://api.vrchat.cloud/api/1/instances/{0}:{1}",
            world_id, instance_id
        );
        let url = Url::parse(&api_url).unwrap();

        let req = client
            .get(url)
            .header(USER_AGENT, "vr-event-tracker")
            .header("Cookie", &auth_cookie)
            .send()
            .await?;

        let data: VrcInstanceData = req.json().await?;
        let name = instance.name.unwrap_or(instance_id.clone());
        PLAYER_COUNT
            .with_label_values(&[&world_id, &instance_id, &name])
            .set(data.userCount.unwrap_or(0 as f64));
    }

    let vrcdn_url = format!(
        "https://api.vrcdn.live/v1/viewers/{}",
        config.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);
    }

    let metric_families = prometheus::gather();
    let mut buffer = vec![];
    encoder.encode(&metric_families, &mut buffer).unwrap();
    Ok(buffer)
}