From 9213542b0a0409fe5e84cd2d951fd5147483843c Mon Sep 17 00:00:00 2001
From: Riccardo Busetti <riccardo.busetti@sentry.io>
Date: Wed, 9 Oct 2024 14:53:25 +0200
Subject: [PATCH] feat(spooler): Add metric to track cached envelopes

---
 .../services/buffer/envelope_repository/sqlite.rs    | 12 ++++++++++++
 relay-server/src/statsd.rs                           |  3 +++
 2 files changed, 15 insertions(+)

diff --git a/relay-server/src/services/buffer/envelope_repository/sqlite.rs b/relay-server/src/services/buffer/envelope_repository/sqlite.rs
index 61b0795349..7c5223dc4e 100644
--- a/relay-server/src/services/buffer/envelope_repository/sqlite.rs
+++ b/relay-server/src/services/buffer/envelope_repository/sqlite.rs
@@ -99,6 +99,9 @@ impl SqliteEnvelopeRepository {
             .push(envelope);
 
         self.cached_envelopes_size += 1;
+        relay_statsd::metric!(
+            histogram(RelayHistograms::BufferCachedEnvelopes) = self.cached_envelopes_size
+        );
 
         Ok(())
     }
@@ -127,6 +130,9 @@ impl SqliteEnvelopeRepository {
             }
 
             self.cached_envelopes_size += envelopes.len() as u64;
+            relay_statsd::metric!(
+                histogram(RelayHistograms::BufferCachedEnvelopes) = self.cached_envelopes_size
+            );
             self.envelope_stacks
                 .entry(project_key_pair)
                 .or_default()
@@ -154,6 +160,9 @@ impl SqliteEnvelopeRepository {
         if let Some(envelope) = envelope {
             // We only decrement the counter when removing data from the in memory buffer.
             self.cached_envelopes_size -= 1;
+            relay_statsd::metric!(
+                histogram(RelayHistograms::BufferCachedEnvelopes) = self.cached_envelopes_size
+            );
             return Ok(Some(envelope));
         }
 
@@ -296,6 +305,9 @@ impl SqliteEnvelopeRepository {
         let envelope_iter = self.envelope_stacks.values_mut().flat_map(|e| {
             e.check_disk = true;
             self.cached_envelopes_size -= e.cached_envelopes.len() as u64;
+            relay_statsd::metric!(
+                histogram(RelayHistograms::BufferCachedEnvelopes) = self.cached_envelopes_size
+            );
             relay_statsd::metric!(
                 histogram(RelayHistograms::BufferInMemoryEnvelopesPerKeyPair) =
                     e.cached_envelopes.len() as u64
diff --git a/relay-server/src/statsd.rs b/relay-server/src/statsd.rs
index fb7489261e..219de058bb 100644
--- a/relay-server/src/statsd.rs
+++ b/relay-server/src/statsd.rs
@@ -183,6 +183,8 @@ pub enum RelayHistograms {
     BufferBackpressureEnvelopesCount,
     /// Number of envelopes in the buffer per key pair of projects.
     BufferInMemoryEnvelopesPerKeyPair,
+    /// Number of cached envelopes in the buffer.
+    BufferCachedEnvelopes,
     /// The number of batches emitted per partition.
     BatchesPerPartition,
     /// The number of buckets in a batch emitted.
@@ -314,6 +316,7 @@ impl HistogramMetric for RelayHistograms {
             RelayHistograms::BufferInMemoryEnvelopesPerKeyPair => {
                 "buffer.in_memory_envelopes_per_key_pair"
             }
+            RelayHistograms::BufferCachedEnvelopes => "buffer.cached_envelopes",
             RelayHistograms::ProjectStatePending => "project_state.pending",
             RelayHistograms::ProjectStateAttempts => "project_state.attempts",
             RelayHistograms::ProjectStateRequestBatchSize => "project_state.request.batch_size",