commit
dd84d26ee9
186
src/main.rs
186
src/main.rs
|
@ -1,3 +1,4 @@
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::error::Error;
|
use std::error::Error;
|
||||||
use std::{env, fmt, fs};
|
use std::{env, fmt, fs};
|
||||||
|
|
||||||
|
@ -12,7 +13,9 @@ use axum::{
|
||||||
use lazy_static::lazy_static;
|
use lazy_static::lazy_static;
|
||||||
use maud::html;
|
use maud::html;
|
||||||
use maud::Markup;
|
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 reqwest::header::USER_AGENT;
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
@ -21,13 +24,31 @@ lazy_static! {
|
||||||
static ref PLAYER_COUNT: GaugeVec = register_gauge_vec!(
|
static ref PLAYER_COUNT: GaugeVec = register_gauge_vec!(
|
||||||
"vrchat_playercount",
|
"vrchat_playercount",
|
||||||
"Current number of players in instance.",
|
"Current number of players in instance.",
|
||||||
&["instance", "world", "name"],
|
&["instance", "world", "name", "group"],
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
static ref VRCDN_VIEWERS: GaugeVec = register_gauge_vec!(
|
static ref VRCDN_VIEWERS: GaugeVec = register_gauge_vec!(
|
||||||
"vrcdn_viewers",
|
"vrcdn_viewers",
|
||||||
"Current number viewers according to VRCDN's API.",
|
"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();
|
.unwrap();
|
||||||
}
|
}
|
||||||
|
@ -65,11 +86,20 @@ impl fmt::Display for WsError {
|
||||||
|
|
||||||
#[derive(Clone, Debug, Deserialize)]
|
#[derive(Clone, Debug, Deserialize)]
|
||||||
struct Config {
|
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>,
|
vrchat_token: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Deserialize)]
|
||||||
|
struct VrcGroup {
|
||||||
|
id: String,
|
||||||
|
vrcdn: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Clone, Debug, Deserialize)]
|
#[derive(Clone, Debug, Deserialize)]
|
||||||
struct VrcInstance {
|
struct VrcInstance {
|
||||||
/// Instance ID.
|
/// Instance ID.
|
||||||
|
@ -79,7 +109,6 @@ struct VrcInstance {
|
||||||
/// Raw location
|
/// Raw location
|
||||||
location: Option<String>,
|
location: Option<String>,
|
||||||
/// Custom name for the instance.
|
/// Custom name for the instance.
|
||||||
#[serde(skip)]
|
|
||||||
name: Option<String>,
|
name: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -104,6 +133,22 @@ struct VrcGroupInstance {
|
||||||
location: String,
|
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]
|
#[tokio::main]
|
||||||
async fn main() -> Result<(), ()> {
|
async fn main() -> Result<(), ()> {
|
||||||
let content = fs::read_to_string("config.toml").unwrap();
|
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> {
|
async fn metrics(config: Config) -> Result<Vec<u8>, WsError> {
|
||||||
PLAYER_COUNT.reset();
|
PLAYER_COUNT.reset();
|
||||||
VRCDN_VIEWERS.reset();
|
VRCDN_VIEWERS.reset();
|
||||||
|
WORLD_VISITS.reset();
|
||||||
let encoder = TextEncoder::new();
|
let encoder = TextEncoder::new();
|
||||||
let client = reqwest::Client::new();
|
let client = reqwest::Client::new();
|
||||||
let auth_cookie = format!("auth={}", &config.vrchat_token.unwrap());
|
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",
|
"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
|
let req = client
|
||||||
.get(url)
|
.get(url)
|
||||||
.header(
|
.header(
|
||||||
USER_AGENT,
|
USER_AGENT,
|
||||||
"vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)",
|
"vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)",
|
||||||
)
|
)
|
||||||
.header("Cookie", &auth_cookie)
|
.header("Cookie", auth_cookie)
|
||||||
.send()
|
.send()
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let data: Vec<VrcGroupInstance> = req.json().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();
|
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 {
|
for instance in instances {
|
||||||
let api_url = format!(
|
let api_url = format!(
|
||||||
|
@ -186,21 +260,24 @@ async fn metrics(config: Config) -> Result<Vec<u8>, WsError> {
|
||||||
USER_AGENT,
|
USER_AGENT,
|
||||||
"vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)",
|
"vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)",
|
||||||
)
|
)
|
||||||
.header("Cookie", &auth_cookie)
|
.header("Cookie", auth_cookie)
|
||||||
.send()
|
.send()
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let data: VrcInstanceData = req.json().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
|
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));
|
.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
|
let req = client
|
||||||
.get(vrcdn_url)
|
.get(vrcdn_url)
|
||||||
.header(
|
.header(
|
||||||
|
@ -211,15 +288,64 @@ async fn metrics(config: Config) -> Result<Vec<u8>, WsError> {
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
let vrcdn_data: VrCdnData = req.json().await.unwrap();
|
let vrcdn_data: VrCdnData = req.json().await.unwrap();
|
||||||
|
|
||||||
for region in vrcdn_data.viewers {
|
for region in vrcdn_data.viewers {
|
||||||
VRCDN_VIEWERS
|
VRCDN_VIEWERS
|
||||||
.with_label_values(&[®ion.region])
|
.with_label_values(&[®ion.region, &name])
|
||||||
.set(region.total);
|
.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 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(())
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in a new issue