fix: recalculate all allocations

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

Author: gusinacio

tap-agent/src/tracker.rs

@@ -24,5 +24,6 @@ pub trait AllocationStats<U> {
     fn update(&mut self, v: U);
     fn is_allowed_to_trigger_rav_request(&self) -> bool;
     fn get_stats(&self) -> U;
+    fn get_total_fee(&self) -> u128;
     fn get_valid_fee(&mut self) -> u128;
 }

tap-agent/src/tracker/generic_tracker.rs

@@ -10,7 +10,6 @@ use alloy::primitives::Address;
 use std::{
     collections::{HashMap, HashSet},
     marker::PhantomData,
-    ops::{AddAssign, SubAssign},
     time::Duration,
 };
 
@@ -26,18 +25,6 @@ impl PartialOrd for GlobalFeeTracker {
     }
 }
 
-impl SubAssign<UnaggregatedReceipts> for GlobalFeeTracker {
-    fn sub_assign(&mut self, rhs: UnaggregatedReceipts) {
-        self.total_fee -= rhs.value;
-    }
-}
-
-impl AddAssign<UnaggregatedReceipts> for GlobalFeeTracker {
-    fn add_assign(&mut self, rhs: UnaggregatedReceipts) {
-        self.total_fee += rhs.value;
-    }
-}
-
 #[derive(Debug, Clone, Default)]
 pub struct GenericTracker<G, F, E, U>
 where
@@ -67,14 +54,20 @@ where
             .id_to_fee
             .entry(id)
             .or_insert(F::default_from_extra(&self.extra_data));
-        self.global -= fee.get_stats();
+
         fee.update(value);
-        self.global += value;
+
+        self.recalculate_global();
+    }
+
+    fn recalculate_global(&mut self) {
+        self.global
+            .update(self.id_to_fee.values().map(|fee| fee.get_total_fee()).sum());
     }
 
     pub fn remove(&mut self, id: Address) {
-        if let Some(old_fee) = self.id_to_fee.remove(&id) {
-            self.global -= old_fee.get_stats();
+        if self.id_to_fee.remove(&id).is_some() {
+            self.recalculate_global();
         }
     }
 
@@ -242,4 +235,8 @@ impl AllocationStats<u128> for u128 {
     fn get_valid_fee(&mut self) -> u128 {
         self.get_stats()
     }
+
+    fn get_total_fee(&self) -> u128 {
+        *self
+    }
 }

tap-agent/src/tracker/global_tracker.rs

@@ -1,24 +1,31 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::ops::{AddAssign, SubAssign};
-
 use crate::agent::unaggregated_receipts::UnaggregatedReceipts;
 
 use super::GlobalFeeTracker;
 
-pub trait GlobalTracker<T>: SubAssign<T> + AddAssign<T> + Sized {
+pub trait GlobalTracker<T>: Sized {
     fn get_total_fee(&self) -> u128;
+    fn update(&mut self, new_fee: u128);
 }
 
 impl GlobalTracker<u128> for u128 {
     fn get_total_fee(&self) -> u128 {
         *self
     }
+
+    fn update(&mut self, new_fee: u128) {
+        *self = new_fee;
+    }
 }
 
 impl GlobalTracker<UnaggregatedReceipts> for GlobalFeeTracker {
     fn get_total_fee(&self) -> u128 {
         self.total_fee - self.requesting
     }
+
+    fn update(&mut self, new_fee: u128) {
+        self.total_fee = new_fee;
+    }
 }

tap-agent/src/tracker/sender_fee_stats.rs

@@ -119,4 +119,8 @@ impl AllocationStats<UnaggregatedReceipts> for SenderFeeStats {
     fn get_valid_fee(&mut self) -> u128 {
         self.total_fee - self.buffer_info.get_sum().min(self.total_fee)
     }
+
+    fn get_total_fee(&self) -> u128 {
+        self.total_fee
+    }
 }