Add prom metrics for database, use redis max conn var (#3359)

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?,