From 99eceb4db0ceab7650d4eec93f20b223e9fd564e Mon Sep 17 00:00:00 2001
From: Subhra264 <chakrabortysubhradeep556@gmail.com>
Date: Tue, 4 Jun 2024 14:46:12 +0530
Subject: [PATCH] fix: move usage reporting to cron jobs (#3672)

Fixes #3373

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Improved the handling and publishing of usage data, ensuring more
efficient and timely updates.

- **Refactor**
- Streamlined the internal control flow for managing usage data,
enhancing overall system performance.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 src/job/mod.rs           |  4 ++-
 src/service/usage/mod.rs | 77 +++++++++++++++++++++++++---------------
 2 files changed, 52 insertions(+), 29 deletions(-)

diff --git a/src/job/mod.rs b/src/job/mod.rs
index b2145f6adf..6227433836 100644
--- a/src/job/mod.rs
+++ b/src/job/mod.rs
@@ -24,7 +24,7 @@ use crate::{
         infra::config::SYSLOG_ENABLED,
         meta::{organization::DEFAULT_ORG, user::UserRequest},
     },
-    service::{compact::stats::update_stats_from_file_list, db, users},
+    service::{compact::stats::update_stats_from_file_list, db, usage, users},
 };
 
 mod alert_manager;
@@ -95,6 +95,8 @@ pub async fn init() -> Result<(), anyhow::Error> {
         tokio::task::spawn(async move { telemetry::run().await });
     }
 
+    tokio::task::spawn(async move { usage::run().await });
+
     // initialize metadata watcher
     tokio::task::spawn(async move { db::schema::watch().await });
     tokio::task::spawn(async move { db::functions::watch().await });
diff --git a/src/service/usage/mod.rs b/src/service/usage/mod.rs
index 11c87af0b7..ee3baa06b9 100644
--- a/src/service/usage/mod.rs
+++ b/src/service/usage/mod.rs
@@ -35,7 +35,7 @@ use o2_enterprise::enterprise::common::auditor;
 use once_cell::sync::Lazy;
 use proto::cluster_rpc;
 use reqwest::Client;
-use tokio::sync::RwLock;
+use tokio::{sync::RwLock, time};
 
 pub mod ingestion_service;
 pub mod stats;
@@ -45,12 +45,6 @@ pub static USAGE_DATA: Lazy<Arc<RwLock<Vec<UsageData>>>> =
 pub static TRIGGERS_USAGE_DATA: Lazy<Arc<RwLock<Vec<TriggerData>>>> =
     Lazy::new(|| Arc::new(RwLock::new(vec![])));
 
-pub static LAST_PUBLISHED_USAGE: Lazy<Arc<RwLock<i64>>> =
-    Lazy::new(|| Arc::new(RwLock::new(Utc::now().timestamp())));
-
-pub static LAST_PUBLISHED_TRIGGERS: Lazy<Arc<RwLock<i64>>> =
-    Lazy::new(|| Arc::new(RwLock::new(Utc::now().timestamp())));
-
 pub async fn report_request_usage_stats(
     stats: RequestStats,
     org_id: &str,
@@ -145,15 +139,9 @@ pub async fn publish_usage(mut usage: Vec<UsageData>) {
     let mut usages = USAGE_DATA.write().await;
     usages.append(&mut usage);
 
-    let current_ts = Utc::now().timestamp();
-    let r = LAST_PUBLISHED_USAGE.read().await;
-    let last_published = *r;
-    if usages.len() < get_config().common.usage_batch_size
-        && current_ts - last_published < get_config().common.usage_publish_interval
-    {
+    if usages.len() < get_config().common.usage_batch_size {
         return;
     }
-    drop(r);
 
     let curr_usages = std::mem::take(&mut *usages);
     // release the write lock
@@ -171,16 +159,9 @@ pub async fn publish_triggers_usage(trigger: TriggerData) {
     let mut usages = TRIGGERS_USAGE_DATA.write().await;
     usages.push(trigger);
 
-    let current_ts = Utc::now().timestamp();
-    let r = LAST_PUBLISHED_TRIGGERS.read().await;
-    let last_published = *r;
-
-    if usages.len() < cfg.common.usage_batch_size
-        && current_ts - last_published < cfg.common.usage_publish_interval
-    {
+    if usages.len() < cfg.common.usage_batch_size {
         return;
     }
-    drop(r);
 
     let curr_usages = std::mem::take(&mut *usages);
     // release the write lock
@@ -334,9 +315,6 @@ async fn ingest_usages(curr_usages: Vec<UsageData>) {
             drop(usages);
         }
     }
-    let mut w = LAST_PUBLISHED_USAGE.write().await;
-    *w = Utc::now().timestamp();
-    drop(w);
 }
 
 async fn ingest_trigger_usages(curr_usages: Vec<TriggerData>) {
@@ -363,9 +341,52 @@ async fn ingest_trigger_usages(curr_usages: Vec<TriggerData>) {
         usages.append(&mut curr_usages);
         drop(usages);
     }
-    let mut w = LAST_PUBLISHED_TRIGGERS.write().await;
-    *w = Utc::now().timestamp();
-    drop(w);
+}
+
+async fn publish_existing_usage() {
+    let mut usages = USAGE_DATA.write().await;
+    log::debug!("publishing usage reports,len: {}", usages.len());
+    if usages.is_empty() {
+        return;
+    }
+
+    let curr_usages = std::mem::take(&mut *usages);
+    // release the write lock
+    drop(usages);
+
+    ingest_usages(curr_usages).await
+}
+
+async fn publish_existing_triggers_usage() {
+    let mut usages = TRIGGERS_USAGE_DATA.write().await;
+    log::debug!("publishing triggers usage reports,len: {}", usages.len());
+
+    if usages.is_empty() {
+        return;
+    }
+
+    let curr_usages = std::mem::take(&mut *usages);
+    // release the write lock
+    drop(usages);
+
+    ingest_trigger_usages(curr_usages).await
+}
+
+pub async fn run() {
+    let cfg = get_config();
+    if !cfg.common.usage_enabled {
+        return;
+    }
+    let mut usage_interval = time::interval(time::Duration::from_secs(
+        cfg.common.usage_publish_interval.try_into().unwrap(),
+    ));
+    usage_interval.tick().await; // trigger the first run
+    loop {
+        log::debug!("Usage ingestion loop running");
+        usage_interval.tick().await;
+        publish_existing_usage().await;
+        publish_existing_triggers_usage().await;
+    }
 }
 
 #[cfg(feature = "enterprise")]