Expand Metrics Tracked #1
117
src/main.rs
117
src/main.rs
|
@ -1,3 +1,4 @@
|
|||
use std::collections::HashMap;
|
||||
use std::error::Error;
|
||||
use std::{env, fmt, fs};
|
||||
|
||||
|
@ -12,12 +13,13 @@ 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;
|
||||
|
||||
lazy_static! {
|
||||
lazy_static!{
|
||||
static ref PLAYER_COUNT: GaugeVec = register_gauge_vec!(
|
||||
"vrchat_playercount",
|
||||
"Current number of players in instance.",
|
||||
|
@ -65,11 +67,20 @@ impl fmt::Display for WsError {
|
|||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
struct Config {
|
||||
vrcdn: Option<String>,
|
||||
group: Option<String>,
|
||||
/// 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.
|
||||
|
@ -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.
|
||||
if config.group.is_some() {
|
||||
let api_url = format!(
|
||||
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 }
|
||||
}).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!(
|
||||
|
@ -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(&[®ion.region])
|
||||
.set(region.total);
|
||||
}
|
||||
}
|
||||
|
||||
let metric_families = prometheus::gather();
|
||||
let mut buffer = vec![];
|
||||
encoder.encode(&metric_families, &mut buffer).unwrap();
|
||||
Ok(buffer)
|
||||
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(())
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue