From 1ad9f6df1eabca2a191e9ae455ee78ebdafa5fbb Mon Sep 17 00:00:00 2001
From: David Herberth <david.herberth@sentry.io>
Date: Thu, 11 Jul 2024 15:16:25 +0200
Subject: [PATCH] ref(redis): Expose redis pool state via statsd (#3812)

Need more insight into the Redis pool, we should make sure it's never
exhausted.
---
 relay-redis/src/real.rs            | 21 ++++++++++++++++++++
 relay-server/src/service.rs        | 10 ++++++++--
 relay-server/src/services/stats.rs | 32 ++++++++++++++++++++++++++++--
 relay-server/src/statsd.rs         | 10 ++++++++++
 4 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/relay-redis/src/real.rs b/relay-redis/src/real.rs
index 48baa9b415..b6adbc5a5f 100644
--- a/relay-redis/src/real.rs
+++ b/relay-redis/src/real.rs
@@ -208,4 +208,25 @@ impl RedisPool {
             inner,
         })
     }
+
+    /// Returns information about the current state of the pool.
+    pub fn stats(&self) -> Stats {
+        let s = match &self.inner {
+            RedisPoolInner::Cluster(p) => p.state(),
+            RedisPoolInner::Single(p) => p.state(),
+        };
+
+        Stats {
+            connections: s.connections,
+            idle_connections: s.idle_connections,
+        }
+    }
+}
+
+/// Stats about how the [`RedisPool`] is performing.
+pub struct Stats {
+    /// The number of connections currently being managed by the pool.
+    pub connections: u32,
+    /// The number of idle connections.
+    pub idle_connections: u32,
 }
diff --git a/relay-server/src/service.rs b/relay-server/src/service.rs
index 7a7b0e9053..3e5d6ba388 100644
--- a/relay-server/src/service.rs
+++ b/relay-server/src/service.rs
@@ -198,7 +198,7 @@ impl ServiceState {
             buffer_guard.clone(),
             project_cache_services,
             metric_outcomes,
-            redis_pool,
+            redis_pool.clone(),
         )
         .spawn_handler(project_cache_rx);
 
@@ -210,7 +210,13 @@ impl ServiceState {
         )
         .start();
 
-        RelayStats::new(config.clone(), upstream_relay.clone()).start();
+        RelayStats::new(
+            config.clone(),
+            upstream_relay.clone(),
+            #[cfg(feature = "processing")]
+            redis_pool,
+        )
+        .start();
 
         let relay_cache = RelayCacheService::new(config.clone(), upstream_relay.clone()).start();
 
diff --git a/relay-server/src/services/stats.rs b/relay-server/src/services/stats.rs
index 81beff2398..a2a1a960d7 100644
--- a/relay-server/src/services/stats.rs
+++ b/relay-server/src/services/stats.rs
@@ -1,6 +1,8 @@
 use std::sync::Arc;
 
 use relay_config::{Config, RelayMode};
+#[cfg(feature = "processing")]
+use relay_redis::RedisPool;
 use relay_statsd::metric;
 use relay_system::{Addr, Service};
 use tokio::time::interval;
@@ -14,13 +16,21 @@ use crate::statsd::{RelayGauges, TokioGauges};
 pub struct RelayStats {
     config: Arc<Config>,
     upstream_relay: Addr<UpstreamRelay>,
+    #[cfg(feature = "processing")]
+    redis_pool: Option<RedisPool>,
 }
 
 impl RelayStats {
-    pub fn new(config: Arc<Config>, upstream_relay: Addr<UpstreamRelay>) -> Self {
+    pub fn new(
+        config: Arc<Config>,
+        upstream_relay: Addr<UpstreamRelay>,
+        #[cfg(feature = "processing")] redis_pool: Option<RedisPool>,
+    ) -> Self {
         Self {
             config,
             upstream_relay,
+            #[cfg(feature = "processing")]
+            redis_pool,
         }
     }
 
@@ -89,6 +99,20 @@ impl RelayStats {
             }
         }
     }
+
+    #[cfg(not(feature = "processing"))]
+    async fn redis_pool(&self) {}
+
+    #[cfg(feature = "processing")]
+    async fn redis_pool(&self) {
+        let Some(ref redis_pool) = self.redis_pool else {
+            return;
+        };
+
+        let state = redis_pool.stats();
+        metric!(gauge(RelayGauges::RedisPoolConnections) = u64::from(state.connections));
+        metric!(gauge(RelayGauges::RedisPoolIdleConnections) = u64::from(state.idle_connections));
+    }
 }
 
 impl Service for RelayStats {
@@ -101,7 +125,11 @@ impl Service for RelayStats {
 
         tokio::spawn(async move {
             loop {
-                let _ = tokio::join!(self.tokio_metrics(), self.upstream_status());
+                let _ = tokio::join!(
+                    self.upstream_status(),
+                    self.tokio_metrics(),
+                    self.redis_pool(),
+                );
                 ticker.tick().await;
             }
         });
diff --git a/relay-server/src/statsd.rs b/relay-server/src/statsd.rs
index efe378eff8..1b5747cafa 100644
--- a/relay-server/src/statsd.rs
+++ b/relay-server/src/statsd.rs
@@ -34,6 +34,12 @@ pub enum RelayGauges {
     ///
     /// Relay uses the same value for its memory health check.
     SystemMemoryTotal,
+    /// The number of connections currently being managed by the Redis Pool.
+    #[cfg(feature = "processing")]
+    RedisPoolConnections,
+    /// The number of idle connections in the Redis Pool.
+    #[cfg(feature = "processing")]
+    RedisPoolIdleConnections,
 }
 
 impl GaugeMetric for RelayGauges {
@@ -46,6 +52,10 @@ impl GaugeMetric for RelayGauges {
             RelayGauges::BufferPeriodicUnspool => "buffer.unspool.periodic",
             RelayGauges::SystemMemoryUsed => "health.system_memory.used",
             RelayGauges::SystemMemoryTotal => "health.system_memory.total",
+            #[cfg(feature = "processing")]
+            RelayGauges::RedisPoolConnections => "redis.pool.connections",
+            #[cfg(feature = "processing")]
+            RelayGauges::RedisPoolIdleConnections => "redis.pool.idle_connections",
         }
     }
 }