refactor: add timestamp_ns from receipt

gusinacio · gusinacio · commit 0019608d29ff · 2024-10-23T16:57:22.000+02:00
Signed-off-by: Gustavo Inacio &lt;gustavo@semiotic.ai&gt;
diff --git a/tap-agent/src/agent/sender_account.rs b/tap-agent/src/agent/sender_account.rs
@@ -81,7 +81,7 @@ type Balance = U256;
 
 #[derive(Debug)]
 pub enum ReceiptFees {
-    NewReceipt(u128),
+    NewReceipt(u128, u64),
     UpdateValue(UnaggregatedReceipts),
     RavRequestResponse(anyhow::Result<(UnaggregatedReceipts, Option<SignedRAV>)>),
     Retry,
@@ -594,7 +594,7 @@ impl Actor for SenderAccount {
                 }
 
                 match receipt_fees {
-                    ReceiptFees::NewReceipt(value) => {
+                    ReceiptFees::NewReceipt(value, timestamp_ns) => {
                         // If state is denied and received new receipt, sender was removed manually from DB
                         if state.denied {
                             tracing::warn!(
@@ -609,7 +609,9 @@ impl Actor for SenderAccount {
                         }
 
                         // add new value
-                        state.sender_fee_tracker.add(allocation_id, value);
+                        state
+                            .sender_fee_tracker
+                            .add(allocation_id, value, timestamp_ns);
                         UNAGGREGATED_FEES
                             .with_label_values(&[
                                 &state.sender.to_string(),
@@ -640,8 +642,7 @@ impl Actor for SenderAccount {
                 let counter_greater_receipt_limit = total_counter_for_allocation
                     >= state.config.tap.rav_request_receipt_limit
                     && can_trigger_rav;
-                let total_fee_outside_buffer =
-                    state.sender_fee_tracker.get_total_fee_outside_buffer();
+                let total_fee_outside_buffer = state.sender_fee_tracker.get_ravable_total_fee();
                 let total_fee_greater_trigger_value =
                     total_fee_outside_buffer >= state.config.tap.rav_request_trigger_value;
                 let rav_result = match (
@@ -917,7 +918,7 @@ pub mod tests {
     use std::collections::{HashMap, HashSet};
     use std::sync::atomic::AtomicU32;
     use std::sync::{Arc, Mutex};
-    use std::time::Duration;
+    use std::time::{Duration, SystemTime, UNIX_EPOCH};
     use wiremock::matchers::{body_string_contains, method};
     use wiremock::{Mock, MockServer, ResponseTemplate};
 
@@ -931,7 +932,9 @@ pub mod tests {
                 (Self::UpdateReceiptFees(l0, l1), Self::UpdateReceiptFees(r0, r1)) => {
                     l0 == r0
                         && match (l1, r1) {
-                            (ReceiptFees::NewReceipt(l), ReceiptFees::NewReceipt(r)) => r == l,
+                            (ReceiptFees::NewReceipt(l1, l2), ReceiptFees::NewReceipt(r1, r2)) => {
+                                r1 == l1 && r2 == l2
+                            }
                             (ReceiptFees::UpdateValue(l), ReceiptFees::UpdateValue(r)) => r == l,
                             (
                                 ReceiptFees::RavRequestResponse(l),
@@ -1278,6 +1281,13 @@ pub mod tests {
         )
     }
 
+    fn get_current_timestamp_u64_ns() -> u64 {
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos() as u64
+    }
+
     #[sqlx::test(migrations = "../migrations")]
     async fn test_update_receipt_fees_no_rav(pgpool: PgPool) {
         let (sender_account, handle, prefix, _) = create_sender_account(
@@ -1303,7 +1313,7 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::NewReceipt(TRIGGER_VALUE - 1),
+                ReceiptFees::NewReceipt(TRIGGER_VALUE - 1, get_current_timestamp_u64_ns()),
             ))
             .unwrap();
 
@@ -1346,7 +1356,7 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::NewReceipt(TRIGGER_VALUE),
+                ReceiptFees::NewReceipt(TRIGGER_VALUE, get_current_timestamp_u64_ns()),
             ))
             .unwrap();
 
@@ -1406,7 +1416,7 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::NewReceipt(1),
+                ReceiptFees::NewReceipt(1, get_current_timestamp_u64_ns()),
             ))
             .unwrap();
 
@@ -1419,7 +1429,7 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::NewReceipt(1),
+                ReceiptFees::NewReceipt(1, get_current_timestamp_u64_ns()),
             ))
             .unwrap();
 
@@ -1547,7 +1557,7 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::NewReceipt(TRIGGER_VALUE),
+                ReceiptFees::NewReceipt(TRIGGER_VALUE, get_current_timestamp_u64_ns()),
             ))
             .unwrap();
         tokio::time::sleep(Duration::from_millis(200)).await;
@@ -1941,7 +1951,7 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::NewReceipt(TRIGGER_VALUE),
+                ReceiptFees::NewReceipt(TRIGGER_VALUE, get_current_timestamp_u64_ns()),
             ))
             .unwrap();
         tokio::time::sleep(Duration::from_millis(100)).await;
diff --git a/tap-agent/src/agent/sender_allocation.rs b/tap-agent/src/agent/sender_allocation.rs
@@ -229,7 +229,10 @@ impl Actor for SenderAllocation {
         match message {
             SenderAllocationMessage::NewReceipt(notification) => {
                 let NewReceiptNotification {
-                    id, value: fees, ..
+                    id,
+                    value: fees,
+                    timestamp_ns,
+                    ..
                 } = notification;
                 if id <= unaggregated_fees.last_id {
                     // our world assumption is wrong
@@ -260,7 +263,7 @@ impl Actor for SenderAllocation {
                     .sender_account_ref
                     .cast(SenderAccountMessage::UpdateReceiptFees(
                         state.allocation_id,
-                        ReceiptFees::NewReceipt(fees),
+                        ReceiptFees::NewReceipt(fees, timestamp_ns),
                     ))?;
             }
             SenderAllocationMessage::TriggerRAVRequest => {
@@ -1109,14 +1112,19 @@ pub mod tests {
         )
         .unwrap();
 
+        let timestamp_ns = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos() as u64;
+
         cast!(
             sender_allocation,
             SenderAllocationMessage::NewReceipt(NewReceiptNotification {
                 id: 1,
                 value: 20,
                 allocation_id: *ALLOCATION_ID_0,
                 signer_address: SIGNER.1,
-                timestamp_ns: 0,
+                timestamp_ns,
             })
         )
         .unwrap();
@@ -1126,7 +1134,7 @@ pub mod tests {
         // should emit update aggregate fees message to sender account
         let expected_message = SenderAccountMessage::UpdateReceiptFees(
             *ALLOCATION_ID_0,
-            ReceiptFees::NewReceipt(20u128),
+            ReceiptFees::NewReceipt(20u128, timestamp_ns),
         );
         let startup_load_msg = message_receiver.recv().await.unwrap();
         assert_eq!(
diff --git a/tap-agent/src/tracker/generic_tracker.rs b/tap-agent/src/tracker/generic_tracker.rs
@@ -119,7 +119,7 @@ impl GenericTracker<GlobalFeeTracker, SenderFeeStats, DurationInfo, Unaggregated
     /// It's important to notice that `value` cannot be less than
     /// zero, so the only way to make this counter lower is by using
     /// `update` function
-    pub fn add(&mut self, id: Address, value: u128) {
+    pub fn add(&mut self, id: Address, value: u128, timestamp_ns: u64) {
         let contains_buffer = self.contains_buffer();
         let entry = self
             .id_to_fee
@@ -134,15 +134,15 @@ impl GenericTracker<GlobalFeeTracker, SenderFeeStats, DurationInfo, Unaggregated
         entry.count += 1;
 
         if contains_buffer {
-            entry.buffer_info.new_entry(value);
+            entry.buffer_info.new_entry(value, timestamp_ns);
         }
     }
 
     fn contains_buffer(&self) -> bool {
         self.extra_data.buffer_duration > Duration::ZERO
     }
 
-    pub fn get_total_fee_outside_buffer(&mut self) -> u128 {
+    pub fn get_ravable_total_fee(&mut self) -> u128 {
         self.get_total_fee() - self.get_buffered_fee().min(self.global.total_fee)
     }
 
@@ -152,13 +152,10 @@ impl GenericTracker<GlobalFeeTracker, SenderFeeStats, DurationInfo, Unaggregated
             .fold(0u128, |acc, expiring| acc + expiring.buffer_info.get_sum())
     }
 
-    pub fn get_count_outside_buffer_for_allocation(
-        &mut self,
-        allocation_id: &Address,
-    ) -> u64 {
+    pub fn get_count_outside_buffer_for_allocation(&mut self, allocation_id: &Address) -> u64 {
         self.id_to_fee
             .get_mut(allocation_id)
-            .map(|alloc| alloc.get_count_outside_buffer())
+            .map(|alloc| alloc.ravable_count())
             .unwrap_or_default()
     }
 
@@ -186,15 +183,15 @@ impl GenericTracker<GlobalFeeTracker, SenderFeeStats, DurationInfo, Unaggregated
             .id_to_fee
             .entry(allocation_id)
             .or_insert(SenderFeeStats::default_from_extra(&self.extra_data));
-        entry.failed_info.ok();
+        entry.backoff_info.ok();
     }
 
     pub fn failed_rav_backoff(&mut self, allocation_id: Address) {
         let entry = self
             .id_to_fee
             .entry(allocation_id)
             .or_insert(SenderFeeStats::default_from_extra(&self.extra_data));
-        entry.failed_info.fail();
+        entry.backoff_info.fail();
     }
 }
 
diff --git a/tap-agent/src/tracker/sender_fee_stats.rs b/tap-agent/src/tracker/sender_fee_stats.rs
@@ -1,6 +1,6 @@
 use std::{
     collections::VecDeque,
-    time::{Duration, Instant},
+    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
 };
 // use tracing::error;
 
@@ -22,11 +22,11 @@ pub struct SenderFeeStats {
     pub(super) buffer_info: BufferInfo,
 
     // Backoff info
-    pub(super) failed_info: FailedRavInfo,
+    pub(super) backoff_info: BackoffInfo,
 }
 
 impl SenderFeeStats {
-    pub(crate) fn get_count_outside_buffer(&mut self) -> u64 {
+    pub(super) fn ravable_count(&mut self) -> u64 {
         let allocation_counter = self.count;
         let counter_in_buffer = self.buffer_info.get_count();
         allocation_counter - counter_in_buffer
@@ -35,15 +35,20 @@ impl SenderFeeStats {
 
 #[derive(Debug, Clone, Default)]
 pub struct BufferInfo {
-    pub entries: VecDeque<(Instant, u128)>,
+    pub entries: VecDeque<(SystemTime, u128)>,
     pub fee_in_buffer: u128,
     pub duration: Duration,
 }
 
 impl BufferInfo {
-    pub(super) fn new_entry(&mut self, value: u128) {
-        let now = Instant::now();
-        self.entries.push_back((now, value));
+    pub(super) fn new_entry(&mut self, value: u128, timestamp_ns: u64) {
+        let duration_since_epoch = Duration::from_nanos(timestamp_ns);
+        // Create a SystemTime from the UNIX_EPOCH plus the duration
+        let system_time = UNIX_EPOCH + duration_since_epoch;
+        let system_time = system_time
+            .checked_add(self.duration)
+            .expect("Should be within bounds");
+        self.entries.push_back((system_time, value));
         self.fee_in_buffer += value;
     }
 
@@ -57,26 +62,33 @@ impl BufferInfo {
         self.entries.len() as u64
     }
 
-    fn cleanup(&mut self) {
-        let now = Instant::now();
+    // O(Receipts expired)
+    fn cleanup(&mut self) -> (u128, u64) {
+        let now = SystemTime::now();
+
+        let mut total_value_expired = 0;
+        let mut total_count_expired = 0;
         while let Some(&(timestamp, value)) = self.entries.front() {
-            if now.duration_since(timestamp) >= self.duration {
+            if now >= timestamp {
                 self.entries.pop_front();
-                self.fee_in_buffer -= value;
+                total_value_expired += value;
+                total_count_expired += 1;
             } else {
                 break;
             }
         }
+        self.fee_in_buffer -= total_value_expired;
+        (total_value_expired, total_count_expired)
     }
 }
 
 #[derive(Debug, Clone)]
-pub struct FailedRavInfo {
+pub struct BackoffInfo {
     failed_ravs_count: u32,
     failed_rav_backoff_time: Instant,
 }
 
-impl FailedRavInfo {
+impl BackoffInfo {
     pub fn ok(&mut self) {
         self.failed_ravs_count = 0;
     }
@@ -107,7 +119,7 @@ impl DefaultFromExtra<DurationInfo> for SenderFeeStats {
     }
 }
 
-impl Default for FailedRavInfo {
+impl Default for BackoffInfo {
     fn default() -> Self {
         Self {
             failed_ravs_count: 0,
@@ -137,7 +149,7 @@ impl AllocationStats<UnaggregatedReceipts> for SenderFeeStats {
     }
 
     fn is_allowed_to_trigger_rav_request(&self) -> bool {
-        !self.failed_info.in_backoff() && !self.blocked && !self.requesting
+        !self.backoff_info.in_backoff() && !self.blocked && !self.requesting
     }
 
     fn get_stats(&self) -> UnaggregatedReceipts {
diff --git a/tap-agent/src/tracker/tracker_tests.rs b/tap-agent/src/tracker/tracker_tests.rs
