refresh cache usages every 16MB (#958)

Author: dinmukhamedm

app-server/src/cache/keys.rs

@@ -6,3 +6,5 @@ pub const PROJECT_API_KEY_CACHE_KEY: &str = "project_api_key";
 pub const PROJECT_CACHE_KEY: &str = "project";
 pub const WORKSPACE_LIMITS_CACHE_KEY: &str = "workspace_limits";
 pub const PROJECT_EVALUATORS_BY_PATH_CACHE_KEY: &str = "project_evaluators_by_path";
+
+pub const WORKSPACE_PARTIAL_USAGE_CACHE_KEY: &str = "workspace_partial_usage";

app-server/src/ch/events.rs

@@ -54,20 +54,23 @@ impl CHEvent {
     }
 }
 
-pub async fn insert_events(clickhouse: clickhouse::Client, events: Vec<CHEvent>) -> Result<()> {
+/// Insert events into ClickHouse and return the number of bytes inserted
+pub async fn insert_events(clickhouse: clickhouse::Client, events: Vec<CHEvent>) -> Result<usize> {
     if events.is_empty() {
-        return Ok(());
+        return Ok(0);
     }
 
     let ch_insert = clickhouse.insert("events");
     match ch_insert {
         Ok(mut ch_insert) => {
+            let mut total_size_bytes = 0;
             for event in events {
                 ch_insert.write(&event).await?;
+                total_size_bytes += event.size_bytes as usize;
             }
             let ch_insert_end_res = ch_insert.end().await;
             match ch_insert_end_res {
-                Ok(_) => Ok(()),
+                Ok(_) => Ok(total_size_bytes),
                 Err(e) => Err(anyhow::anyhow!(
                     "Clickhouse events insertion failed: {:?}",
                     e

app-server/src/traces/consumer.rs

@@ -369,20 +369,28 @@ async fn process_batch(
         log::error!("Failed to ack MQ delivery (batch): {:?}", e);
     });
 
-    match record_events(clickhouse.clone(), &all_events).await {
-        Ok(_) => {}
+    let total_events_ingested_bytes = match record_events(clickhouse.clone(), &all_events).await {
+        Ok(bytes) => bytes,
         Err(e) => {
             log::error!("Failed to record events: {:?}", e);
+            0
         }
     };
 
+    let total_ingested_bytes = spans_ingested_bytes
+        .iter()
+        .map(|b| b.span_bytes)
+        .sum::<usize>()
+        + total_events_ingested_bytes;
+
     for project_id in project_ids {
         if is_feature_enabled(Feature::UsageLimit) {
             if let Err(e) = update_workspace_limit_exceeded_by_project_id(
                 db.clone(),
                 clickhouse.clone(),
                 cache.clone(),
                 project_id,
+                total_ingested_bytes,
             )
             .await
             {

app-server/src/traces/events.rs

@@ -5,10 +5,11 @@ use crate::{
     db::events::Event,
 };
 
+///
 pub async fn record_events(
     clickhouse: clickhouse::Client,
     event_payloads: &Vec<Event>,
-) -> Result<()> {
+) -> Result<usize> {
     let ch_events = event_payloads
         .iter()
         .map(|e| CHEvent::from_db_event(e))

app-server/src/traces/limits.rs

@@ -9,12 +9,15 @@ use uuid::Uuid;
 use crate::{
     cache::{
         Cache, CacheTrait,
-        keys::{PROJECT_CACHE_KEY, WORKSPACE_LIMITS_CACHE_KEY},
+        keys::{PROJECT_CACHE_KEY, WORKSPACE_LIMITS_CACHE_KEY, WORKSPACE_PARTIAL_USAGE_CACHE_KEY},
     },
     ch,
     db::{self, DB, projects::ProjectWithWorkspaceBillingInfo, stats::WorkspaceLimitsExceeded},
 };
 
+// Threshold in bytes (16MB) - only recompute workspace limits after this much data is written
+const RECOMPUTE_THRESHOLD_BYTES: usize = 16 * 1024 * 1024; // 16MB
+
 pub async fn get_workspace_limit_exceeded_by_project_id(
     db: Arc<DB>,
     clickhouse: clickhouse::Client,
@@ -56,6 +59,7 @@ pub async fn update_workspace_limit_exceeded_by_project_id(
     clickhouse: clickhouse::Client,
     cache: Arc<Cache>,
     project_id: Uuid,
+    written_bytes: usize,
 ) -> Result<()> {
     tokio::spawn(async move {
         let project_info = get_workspace_info_for_project_id(db.clone(), cache.clone(), project_id)
@@ -74,26 +78,57 @@ pub async fn update_workspace_limit_exceeded_by_project_id(
             return;
         }
 
-        let cache_key = format!("{WORKSPACE_LIMITS_CACHE_KEY}:{workspace_id}");
-        let workspace_limits_exceeded = is_workspace_over_limit(
-            clickhouse,
-            project_info.workspace_project_ids,
-            project_info.bytes_limit,
-            project_info.reset_time,
-        )
-        .await
-        .map_err(|e| {
-            log::error!(
-                "Failed to update workspace limit exceeded for project [{}]: {:?}",
-                project_id,
-                e
-            );
-        })
-        .unwrap();
-        cache
-            .insert::<WorkspaceLimitsExceeded>(&cache_key, workspace_limits_exceeded.clone())
+        let partial_usage_cache_key = format!("{WORKSPACE_PARTIAL_USAGE_CACHE_KEY}:{workspace_id}");
+        let limits_cache_key = format!("{WORKSPACE_LIMITS_CACHE_KEY}:{workspace_id}");
+
+        // Get current partial usage from cache
+        let cache_result = cache.get::<usize>(&partial_usage_cache_key).await;
+
+        // If cache is missing or errored, we should recompute
+        let (current_partial_usage, cache_available) = match cache_result {
+            Ok(Some(value)) => (value, true),
+            Ok(None) | Err(_) => (0, false),
+        };
+
+        let new_partial_usage = current_partial_usage + written_bytes;
+
+        // Recompute if: cache was unavailable, or we've accumulated at least RECOMPUTE_THRESHOLD_BYTES
+        let should_recompute = !cache_available || new_partial_usage >= RECOMPUTE_THRESHOLD_BYTES;
+
+        if should_recompute {
+            // Perform the heavy computation
+            let workspace_limits_exceeded = is_workspace_over_limit(
+                clickhouse,
+                project_info.workspace_project_ids,
+                project_info.bytes_limit,
+                project_info.reset_time,
+            )
             .await
+            .map_err(|e| {
+                log::error!(
+                    "Failed to update workspace limit exceeded for project [{}]: {:?}",
+                    project_id,
+                    e
+                );
+            })
             .unwrap();
+
+            // Update the limits cache
+            let _ = cache
+                .insert::<WorkspaceLimitsExceeded>(
+                    &limits_cache_key,
+                    workspace_limits_exceeded.clone(),
+                )
+                .await;
+
+            // Reset the partial usage counter
+            let _ = cache.insert::<usize>(&partial_usage_cache_key, 0).await;
+        } else {
+            // Just update the partial usage counter
+            let _ = cache
+                .insert::<usize>(&partial_usage_cache_key, new_partial_usage)
+                .await;
+        }
     });
 
     Ok(())