From b16c3d81ac02ca234f0a04f410a6a375b1a9630f Mon Sep 17 00:00:00 2001
From: Xavier Vello <xavier.vello@gmail.com>
Date: Tue, 2 Jul 2024 10:46:06 +0200
Subject: [PATCH] chore(capture-rs): add capture_billing_limits_loaded_tokens
 metric (#23193)

---
 rust/capture/src/limiters/billing.rs | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/rust/capture/src/limiters/billing.rs b/rust/capture/src/limiters/billing.rs
index b908519dda265..259fbcf0d69ae 100644
--- a/rust/capture/src/limiters/billing.rs
+++ b/rust/capture/src/limiters/billing.rs
@@ -1,3 +1,4 @@
+use metrics::gauge;
 use std::{collections::HashSet, ops::Sub, sync::Arc};
 
 use crate::redis::Client;
@@ -86,9 +87,9 @@ impl BillingLimiter {
     #[instrument(skip_all)]
     async fn fetch_limited(
         client: &Arc<dyn Client + Send + Sync>,
-        resource: QuotaResource,
+        resource: &QuotaResource,
     ) -> anyhow::Result<Vec<String>> {
-        let now = time::OffsetDateTime::now_utc().unix_timestamp();
+        let now = OffsetDateTime::now_utc().unix_timestamp();
 
         client
             .zrangebyscore(
@@ -130,12 +131,17 @@ impl BillingLimiter {
             // On prod atm we call this around 15 times per second at peak times, and it usually
             // completes in <1ms.
 
-            let set = Self::fetch_limited(&self.redis, resource).await;
+            let set = Self::fetch_limited(&self.redis, &resource).await;
 
             tracing::debug!("fetched set from redis, caching");
 
             if let Ok(set) = set {
                 let set = HashSet::from_iter(set.iter().cloned());
+                gauge!(
+                    "capture_billing_limits_loaded_tokens",
+                    "resource" => resource.as_str(),
+                )
+                .set(set.len() as f64);
 
                 let mut limited = self.limited.write().await;
                 *limited = set;