test(tap-agent): buffer window fee

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

Author: gusinacio

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(), 0);
+        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);
+    }
 }