test(tap-agent): buffer window fee

Signed-off-by: Gustavo Inacio <[email protected]>

Author: gusinacio

tap-agent/src/agent/sender_account.rs

@@ -601,7 +601,6 @@ impl Actor for SenderAccount {
                     state.add_to_denylist().await;
                 }
 
-                // TODO if the total_fee_outside_the_buffer
                 if state.sender_fee_tracker.get_total_fee_outsite_buffer()
                     >= state.config.tap.rav_request_trigger_value
                 {
@@ -904,6 +903,7 @@ pub mod tests {
     const DUMMY_URL: &str = "http://localhost:1234";
     const TRIGGER_VALUE: u128 = 500;
     const ESCROW_VALUE: u128 = 1000;
+    const BUFFER_MS: u64 = 100;
 
     async fn create_sender_account(
         pgpool: PgPool,
@@ -924,7 +924,7 @@ pub mod tests {
             },
             tap: config::Tap {
                 rav_request_trigger_value,
-                rav_request_timestamp_buffer_ms: 1,
+                rav_request_timestamp_buffer_ms: BUFFER_MS,
                 rav_request_timeout_secs: 5,
                 max_unnaggregated_fees_per_sender,
                 ..Default::default()
@@ -1211,14 +1211,11 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::UpdateValue(UnaggregatedReceipts {
-                    value: TRIGGER_VALUE - 1,
-                    last_id: 10,
-                }),
+                ReceiptFees::NewReceipt(TRIGGER_VALUE - 1),
             ))
             .unwrap();
 
-        tokio::time::sleep(Duration::from_millis(10)).await;
+        tokio::time::sleep(Duration::from_millis(BUFFER_MS)).await;
 
         assert_eq!(
             triggered_rav_request.load(std::sync::atomic::Ordering::SeqCst),
@@ -1250,10 +1247,24 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::UpdateValue(UnaggregatedReceipts {
-                    value: TRIGGER_VALUE,
-                    last_id: 10,
-                }),
+                ReceiptFees::NewReceipt(TRIGGER_VALUE),
+            ))
+            .unwrap();
+
+        tokio::time::sleep(Duration::from_millis(20)).await;
+
+        assert_eq!(
+            triggered_rav_request.load(std::sync::atomic::Ordering::SeqCst),
+            0
+        );
+
+        // wait for it to be outside buffer
+        tokio::time::sleep(Duration::from_millis(BUFFER_MS)).await;
+
+        sender_account
+            .cast(SenderAccountMessage::UpdateReceiptFees(
+                *ALLOCATION_ID_0,
+                ReceiptFees::Retry,
             ))
             .unwrap();
 
@@ -1362,10 +1373,7 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::UpdateValue(UnaggregatedReceipts {
-                    value: TRIGGER_VALUE,
-                    last_id: 11,
-                }),
+                ReceiptFees::NewReceipt(TRIGGER_VALUE),
             ))
             .unwrap();
         tokio::time::sleep(Duration::from_millis(200)).await;
@@ -1750,10 +1758,7 @@ pub mod tests {
         sender_account
             .cast(SenderAccountMessage::UpdateReceiptFees(
                 *ALLOCATION_ID_0,
-                ReceiptFees::UpdateValue(UnaggregatedReceipts {
-                    value: TRIGGER_VALUE,
-                    last_id: 11,
-                }),
+                ReceiptFees::NewReceipt(TRIGGER_VALUE),
             ))
             .unwrap();
         tokio::time::sleep(Duration::from_millis(100)).await;

tap-agent/src/agent/sender_fee_tracker.rs

@@ -34,26 +34,30 @@ pub struct SenderFeeTracker {
     id_to_fee: HashMap<Address, u128>,
     total_fee: u128,
 
-    buffer_fee: HashMap<Address, ExpiringSum>,
-
-    buffer_duration: Duration,
+    buffer_window_fee: HashMap<Address, ExpiringSum>,
+    buffer_window_duration: Duration,
     // there are some allocations that we don't want it to be
     // heaviest allocation, because they are already marked for finalization,
     // and thus requesting RAVs on their own in their `post_stop` routine.
     blocked_addresses: HashSet<Address>,
 }
 
 impl SenderFeeTracker {
-    pub fn new(buffer_duration: Duration) -> Self {
+    pub fn new(buffer_window_duration: Duration) -> Self {
         Self {
-            buffer_duration,
+            buffer_window_duration,
             ..Default::default()
         }
     }
+    /// Adds into the total_fee entry and buffer window totals
+    ///
+    /// 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) {
-        if self.buffer_duration > Duration::ZERO {
+        if self.buffer_window_duration > Duration::ZERO {
             let now = Instant::now();
-            let expiring_sum = self.buffer_fee.entry(id).or_default();
+            let expiring_sum = self.buffer_window_fee.entry(id).or_default();
             expiring_sum.entries.push_back((now, value));
             expiring_sum.sum += value;
         }
@@ -62,6 +66,10 @@ impl SenderFeeTracker {
         *entry += value;
     }
 
+    /// Updates and overwrite the fee counter into the specific
+    /// value provided.
+    ///
+    /// IMPORTANT: This function does not affect the buffer window fee
     pub fn update(&mut self, id: Address, fee: u128) {
         if fee > 0 {
             // insert or update, if update remove old fee from total
@@ -100,12 +108,13 @@ impl SenderFeeTracker {
                 (
                     addr,
                     fee - self
-                        .buffer_fee
+                        .buffer_window_fee
                         .get_mut(addr)
-                        .map(|expiring| expiring.get_sum(&self.buffer_duration))
+                        .map(|expiring| expiring.get_sum(&self.buffer_window_duration))
                         .unwrap_or_default(),
                 )
             })
+            .filter(|(_, fee)| *fee > 0)
             .fold(None, |acc: Option<(&Address, u128)>, (addr, fee)| {
                 if let Some((_, max_fee)) = acc {
                     if fee > max_fee {
@@ -129,30 +138,29 @@ impl SenderFeeTracker {
     }
 
     pub fn get_total_fee_outsite_buffer(&mut self) -> u128 {
-        self.total_fee - self.get_buffer_fee()
+        self.total_fee - self.get_buffer_fee().min(self.total_fee)
     }
 
     pub fn get_buffer_fee(&mut self) -> u128 {
-        self.buffer_fee.values_mut().fold(0u128, |acc, expiring| {
-            acc + expiring.get_sum(&self.buffer_duration)
-        })
+        self.buffer_window_fee
+            .values_mut()
+            .fold(0u128, |acc, expiring| {
+                acc + expiring.get_sum(&self.buffer_window_duration)
+            })
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::SenderFeeTracker;
-    use std::str::FromStr;
-    use thegraph_core::Address;
+    use alloy::primitives::address;
+    use std::{thread::sleep, time::Duration};
 
     #[test]
     fn test_allocation_id_tracker() {
-        let allocation_id_0: Address =
-            Address::from_str("0xabababababababababababababababababababab").unwrap();
-        let allocation_id_1: Address =
-            Address::from_str("0xbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbc").unwrap();
-        let allocation_id_2: Address =
-            Address::from_str("0xcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcd").unwrap();
+        let allocation_id_0 = address!("abababababababababababababababababababab");
+        let allocation_id_1 = address!("bcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbc");
+        let allocation_id_2 = address!("cdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcd");
 
         let mut tracker = SenderFeeTracker::default();
         assert_eq!(tracker.get_heaviest_allocation_id(), None);
@@ -204,4 +212,80 @@ mod tests {
         assert_eq!(tracker.get_heaviest_allocation_id(), None);
         assert_eq!(tracker.get_total_fee(), 0);
     }
+
+    #[test]
+    fn test_buffer_tracker_window() {
+        let allocation_id_0 = address!("abababababababababababababababababababab");
+        let allocation_id_1 = address!("bcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbcbc");
+        let allocation_id_2 = address!("cdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcdcd");
+
+        const BUFFER_WINDOW: Duration = Duration::from_millis(20);
+        let mut tracker = SenderFeeTracker::new(BUFFER_WINDOW);
+        assert_eq!(tracker.get_heaviest_allocation_id(), None);
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 0);
+        assert_eq!(tracker.get_total_fee(), 0);
+
+        tracker.add(allocation_id_0, 10);
+        assert_eq!(tracker.get_heaviest_allocation_id(), None);
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 0);
+        assert_eq!(tracker.get_total_fee(), 10);
+
+        sleep(BUFFER_WINDOW);
+
+        assert_eq!(tracker.get_heaviest_allocation_id(), Some(allocation_id_0));
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 10);
+        assert_eq!(tracker.get_total_fee(), 10);
+
+        tracker.add(allocation_id_2, 20);
+        assert_eq!(tracker.get_heaviest_allocation_id(), Some(allocation_id_0));
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 10);
+        assert_eq!(tracker.get_total_fee(), 30);
+
+        sleep(BUFFER_WINDOW);
+
+        tracker.block_allocation_id(allocation_id_2);
+        assert_eq!(tracker.get_heaviest_allocation_id(), Some(allocation_id_0));
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 30);
+        assert_eq!(tracker.get_total_fee(), 30);
+
+        tracker.unblock_allocation_id(allocation_id_2);
+        assert_eq!(tracker.get_heaviest_allocation_id(), Some(allocation_id_2));
+
+        tracker.add(allocation_id_1, 30);
+        assert_eq!(tracker.get_heaviest_allocation_id(), Some(allocation_id_2));
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 30);
+        assert_eq!(tracker.get_total_fee(), 60);
+
+        sleep(BUFFER_WINDOW);
+
+        assert_eq!(tracker.get_heaviest_allocation_id(), Some(allocation_id_1));
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 60);
+        assert_eq!(tracker.get_total_fee(), 60);
+
+        tracker.add(allocation_id_2, 20);
+        tracker.update(allocation_id_2, 0);
+        assert_eq!(tracker.get_heaviest_allocation_id(), Some(allocation_id_1));
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 20);
+        assert_eq!(tracker.get_total_fee(), 40);
+
+        sleep(BUFFER_WINDOW);
+
+        tracker.add(allocation_id_2, 100);
+        tracker.update(allocation_id_2, 0);
+        assert_eq!(tracker.get_heaviest_allocation_id(), Some(allocation_id_1));
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 0);
+        assert_eq!(tracker.get_total_fee(), 40);
+
+        sleep(BUFFER_WINDOW);
+
+        tracker.update(allocation_id_1, 0);
+        assert_eq!(tracker.get_heaviest_allocation_id(), Some(allocation_id_0));
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 10);
+        assert_eq!(tracker.get_total_fee(), 10);
+
+        tracker.update(allocation_id_0, 0);
+        assert_eq!(tracker.get_heaviest_allocation_id(), None);
+        assert_eq!(tracker.get_total_fee_outsite_buffer(), 0);
+        assert_eq!(tracker.get_total_fee(), 0);
+    }
 }