diff --git a/Cargo.lock b/Cargo.lock index 6aff0823..ff38ac48 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4288,6 +4288,7 @@ dependencies = [ "maxminddb", "meilisearch-sdk", "murmur2", + "prometheus", "rand 0.8.5", "rand_chacha 0.3.1", "redis", @@ -6133,9 +6134,16 @@ dependencies = [ "memchr", "parking_lot 0.12.3", "procfs", + "protobuf", "thiserror 1.0.64", ] +[[package]] +name = "protobuf" +version = "2.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "106dd99e98437432fed6519dedecfade6a06a73bb7b2a1e019fdd2bee5778d94" + [[package]] name = "psm" version = "0.1.23" diff --git a/apps/labrinth/Cargo.toml b/apps/labrinth/Cargo.toml index 856cd8f2..ce023ba3 100644 --- a/apps/labrinth/Cargo.toml +++ b/apps/labrinth/Cargo.toml @@ -17,6 +17,7 @@ actix-multipart = "0.6.1" actix-cors = "0.7.0" actix-ws = "0.3.0" actix-files = "0.6.5" +prometheus = "0.13.4" actix-web-prom = { version = "0.9.0", features = ["process"] } governor = "0.6.3" diff --git a/apps/labrinth/src/database/mod.rs b/apps/labrinth/src/database/mod.rs index 2bba7dca..a1c51f2a 100644 --- a/apps/labrinth/src/database/mod.rs +++ b/apps/labrinth/src/database/mod.rs @@ -6,3 +6,4 @@ pub use models::Project; pub use models::Version; pub use postgres_database::check_for_migrations; pub use postgres_database::connect; +pub use postgres_database::register_and_set_metrics; diff --git a/apps/labrinth/src/database/postgres_database.rs b/apps/labrinth/src/database/postgres_database.rs index 65601bde..c90e37f9 100644 --- a/apps/labrinth/src/database/postgres_database.rs +++ b/apps/labrinth/src/database/postgres_database.rs @@ -1,4 +1,5 @@ use log::info; +use prometheus::{IntGauge, Registry}; use sqlx::migrate::MigrateDatabase; use sqlx::postgres::{PgPool, PgPoolOptions}; use sqlx::{Connection, PgConnection, Postgres}; @@ -45,3 +46,30 @@ pub async fn check_for_migrations() -> Result<(), sqlx::Error> { Ok(()) } + +pub async fn register_and_set_metrics( + pool: &PgPool, + registry: &Registry, +) -> Result<(), prometheus::Error> { + let pg_pool_size = + IntGauge::new("labrinth_pg_pool_size", "Size of Postgres pool")?; + let pg_pool_idle = IntGauge::new( + "labrinth_pg_pool_idle", + "Number of idle Postgres connections", + )?; + + registry.register(Box::new(pg_pool_size.clone()))?; + registry.register(Box::new(pg_pool_idle.clone()))?; + + let pool_ref = pool.clone(); + tokio::spawn(async move { + loop { + pg_pool_size.set(pool_ref.size() as i64); + pg_pool_idle.set(pool_ref.num_idle() as i64); + + tokio::time::sleep(Duration::from_secs(5)).await; + } + }); + + Ok(()) +} diff --git a/apps/labrinth/src/database/redis.rs b/apps/labrinth/src/database/redis.rs index fa77707e..3fe92a53 100644 --- a/apps/labrinth/src/database/redis.rs +++ b/apps/labrinth/src/database/redis.rs @@ -3,6 +3,7 @@ use ariadne::ids::base62_impl::{parse_base62, to_base62}; use chrono::{TimeZone, Utc}; use dashmap::DashMap; use deadpool_redis::{Config, Runtime}; +use prometheus::{IntGauge, Registry}; use redis::{cmd, Cmd, ExistenceCheck, SetExpiry, SetOptions}; use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize}; @@ -38,7 +39,7 @@ impl RedisPool { .builder() .expect("Error building Redis pool") .max_size( - dotenvy::var("DATABASE_MAX_CONNECTIONS") + dotenvy::var("REDIS_MAX_CONNECTIONS") .ok() .and_then(|x| x.parse().ok()) .unwrap_or(10000), @@ -53,6 +54,48 @@ impl RedisPool { } } + pub async fn register_and_set_metrics( + &self, + registry: &Registry, + ) -> Result<(), prometheus::Error> { + let redis_max_size = IntGauge::new( + "labrinth_redis_pool_max_size", + "Maximum size of Redis pool", + )?; + let redis_size = IntGauge::new( + "labrinth_redis_pool_size", + "Current size of Redis pool", + )?; + let redis_available = IntGauge::new( + "labrinth_redis_pool_available", + "Available connections in Redis pool", + )?; + let redis_waiting = IntGauge::new( + "labrinth_redis_pool_waiting", + "Number of futures waiting for a Redis connection", + )?; + + registry.register(Box::new(redis_max_size.clone()))?; + registry.register(Box::new(redis_size.clone()))?; + registry.register(Box::new(redis_available.clone()))?; + registry.register(Box::new(redis_waiting.clone()))?; + + let redis_pool_ref = self.pool.clone(); + tokio::spawn(async move { + loop { + let status = redis_pool_ref.status(); + redis_max_size.set(status.max_size as i64); + redis_size.set(status.size as i64); + redis_available.set(status.available as i64); + redis_waiting.set(status.waiting as i64); + + tokio::time::sleep(Duration::from_secs(5)).await; + } + }); + + Ok(()) + } + pub async fn connect(&self) -> Result { Ok(RedisConnection { connection: self.pool.get().await?, diff --git a/apps/labrinth/src/main.rs b/apps/labrinth/src/main.rs index 9310a088..d9b0b7b1 100644 --- a/apps/labrinth/src/main.rs +++ b/apps/labrinth/src/main.rs @@ -99,6 +99,14 @@ async fn main() -> std::io::Result<()> { .build() .expect("Failed to create prometheus metrics middleware"); + database::register_and_set_metrics(&pool, &prometheus.registry) + .await + .expect("Failed to register database metrics"); + redis_pool + .register_and_set_metrics(&prometheus.registry) + .await + .expect("Failed to register redis metrics"); + let search_config = search::SearchConfig::new(None); let labrinth_config = labrinth::app_setup(