Add prom metrics for database, use redis max conn var

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"
 

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;

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(())
+}

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<RedisConnection, DatabaseError> {
         Ok(RedisConnection {
             connection: self.pool.get().await?,

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(