fix: unaggregated fees update race condition with tokio mutex guards

Signed-off-by: Alexis Asseman <[email protected]>

Author: aasseman

tap-agent/src/tap/sender_account.rs

@@ -1,10 +1,11 @@
 // Copyright 2023-, GraphOps and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
+use std::sync::Mutex as StdMutex;
 use std::{
     cmp::max,
     collections::{HashMap, HashSet},
-    sync::{Arc, Mutex},
+    sync::Arc,
     time::Duration,
 };
 
@@ -15,6 +16,7 @@ use eventuals::Eventual;
 use indexer_common::{escrow_accounts::EscrowAccounts, prelude::SubgraphClient};
 use sqlx::PgPool;
 use thegraph::types::Address;
+use tokio::sync::Mutex as TokioMutex;
 use tokio::{select, sync::Notify, time};
 use tracing::{error, warn};
 
@@ -34,10 +36,11 @@ enum AllocationState {
 pub struct Inner {
     config: &'static config::Cli,
     pgpool: PgPool,
-    allocations: Arc<Mutex<HashMap<Address, AllocationState>>>,
+    allocations: Arc<StdMutex<HashMap<Address, AllocationState>>>,
     sender: Address,
     sender_aggregator_endpoint: String,
-    unaggregated_fees: Arc<Mutex<UnaggregatedReceipts>>,
+    unaggregated_fees: Arc<StdMutex<UnaggregatedReceipts>>,
+    unaggregated_receipts_guard: Arc<TokioMutex<()>>,
 }
 
 impl Inner {
@@ -136,7 +139,7 @@ impl Inner {
             return;
         };
 
-        if let Err(e) = self.recompute_unaggregated_fees_static() {
+        if let Err(e) = self.recompute_unaggregated_fees().await {
             error!(
                 "Error while recomputing unaggregated fees for sender {}: {}",
                 self.sender, e
@@ -171,7 +174,11 @@ impl Inner {
     }
 
     /// Recompute the sender's total unaggregated fees value and last receipt ID.
-    fn recompute_unaggregated_fees_static(&self) -> Result<()> {
+    async fn recompute_unaggregated_fees(&self) -> Result<()> {
+        // Make sure to pause the handling of receipt notifications while we update the unaggregated
+        // fees.
+        let _guard = self.unaggregated_receipts_guard.lock().await;
+
         // Similar pattern to get_heaviest_allocation().
         let allocations: Vec<_> = self.allocations.lock().unwrap().values().cloned().collect();
 
@@ -236,6 +243,7 @@ pub struct SenderAccount {
     rav_requester_notify: Arc<Notify>,
     rav_requester_finalize_task: tokio::task::JoinHandle<()>,
     rav_requester_finalize_notify: Arc<Notify>,
+    unaggregated_receipts_guard: Arc<TokioMutex<()>>,
 }
 
 impl SenderAccount {
@@ -250,13 +258,16 @@ impl SenderAccount {
         tap_eip712_domain_separator: Eip712Domain,
         sender_aggregator_endpoint: String,
     ) -> Self {
+        let unaggregated_receipts_guard = Arc::new(TokioMutex::new(()));
+
         let inner = Arc::new(Inner {
             config,
             pgpool,
-            allocations: Arc::new(Mutex::new(HashMap::new())),
+            allocations: Arc::new(StdMutex::new(HashMap::new())),
             sender: sender_id,
             sender_aggregator_endpoint,
-            unaggregated_fees: Arc::new(Mutex::new(UnaggregatedReceipts::default())),
+            unaggregated_fees: Arc::new(StdMutex::new(UnaggregatedReceipts::default())),
+            unaggregated_receipts_guard: unaggregated_receipts_guard.clone(),
         });
 
         let rav_requester_notify = Arc::new(Notify::new());
@@ -289,6 +300,7 @@ impl SenderAccount {
             rav_requester_notify,
             rav_requester_finalize_task,
             rav_requester_finalize_notify,
+            unaggregated_receipts_guard,
         }
     }
 
@@ -346,21 +358,32 @@ impl SenderAccount {
         &self,
         new_receipt_notification: NewReceiptNotification,
     ) {
-        let mut unaggregated_fees = self.inner.unaggregated_fees.lock().unwrap();
-
-        // Else we already processed that receipt, most likely from pulling the receipts
-        // from the database.
-        if new_receipt_notification.id > unaggregated_fees.last_id {
-            if let Some(AllocationState::Active(allocation)) = self
-                .inner
-                .allocations
-                .lock()
-                .unwrap()
-                .get(&new_receipt_notification.allocation_id)
+        // Make sure to pause the handling of receipt notifications while we update the unaggregated
+        // fees.
+        let _guard = self.unaggregated_receipts_guard.lock().await;
+
+        let allocation_state = self
+            .inner
+            .allocations
+            .lock()
+            .unwrap()
+            .get(&new_receipt_notification.allocation_id)
+            .cloned();
+
+        if let Some(AllocationState::Active(allocation)) = allocation_state {
+            // Try to add the receipt value to the allocation's unaggregated fees value.
+            // If the fees were not added, it means the receipt was already processed, so we
+            // don't need to do anything.
+            if allocation
+                .fees_add(new_receipt_notification.value, new_receipt_notification.id)
+                .await
             {
                 // Add the receipt value to the allocation's unaggregated fees value.
-                allocation.fees_add(new_receipt_notification.value);
+                allocation
+                    .fees_add(new_receipt_notification.value, new_receipt_notification.id)
+                    .await;
                 // Add the receipt value to the sender's unaggregated fees value.
+                let mut unaggregated_fees = self.inner.unaggregated_fees.lock().unwrap();
                 *unaggregated_fees = UnaggregatedReceipts {
                     value: self
                         .inner
@@ -373,18 +396,18 @@ impl SenderAccount {
                 {
                     self.rav_requester_notify.notify_waiters();
                 }
-            } else {
-                error!(
-                    "Received a new receipt notification for allocation {} that doesn't exist \
-                    or is ineligible for sender {}.",
-                    new_receipt_notification.allocation_id, self.inner.sender
-                );
             }
+        } else {
+            error!(
+                "Received a new receipt notification for allocation {} that doesn't exist \
+                or is ineligible for sender {}.",
+                new_receipt_notification.allocation_id, self.inner.sender
+            );
         }
     }
 
-    pub fn recompute_unaggregated_fees(&self) -> Result<()> {
-        self.inner.recompute_unaggregated_fees_static()
+    pub async fn recompute_unaggregated_fees(&self) -> Result<()> {
+        self.inner.recompute_unaggregated_fees().await
     }
 }
 
@@ -507,7 +530,7 @@ mod tests {
                 *ALLOCATION_ID_2,
             ]))
             .await;
-        sender.recompute_unaggregated_fees().unwrap();
+        sender.recompute_unaggregated_fees().await.unwrap();
 
         sender
     }

tap-agent/src/tap/sender_accounts_manager.rs

@@ -262,6 +262,7 @@ impl SenderAccountsManager {
 
             sender
                 .recompute_unaggregated_fees()
+                .await
                 .expect("should be able to recompute unaggregated fees");
         }
         drop(sender_accounts_write_lock);

tap-agent/src/tap/sender_allocation.rs

@@ -1,11 +1,8 @@
 // Copyright 2023-, GraphOps and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::{
-    str::FromStr,
-    sync::{Arc, Mutex},
-    time::Duration,
-};
+use std::sync::Mutex as StdMutex;
+use std::{str::FromStr, sync::Arc, time::Duration};
 
 use alloy_primitives::hex::ToHex;
 use alloy_sol_types::Eip712Domain;
@@ -22,6 +19,7 @@ use tap_core::{
     tap_receipt::{ReceiptCheck, ReceivedReceipt},
 };
 use thegraph::types::Address;
+use tokio::sync::Mutex as TokioMutex;
 use tracing::{error, warn};
 
 use crate::{
@@ -48,10 +46,11 @@ pub struct SenderAllocation {
     allocation_id: Address,
     sender: Address,
     sender_aggregator_endpoint: String,
-    unaggregated_fees: Arc<Mutex<UnaggregatedReceipts>>,
+    unaggregated_fees: Arc<StdMutex<UnaggregatedReceipts>>,
     config: &'static config::Cli,
     escrow_accounts: Eventual<EscrowAccounts>,
-    rav_request_guard: tokio::sync::Mutex<()>,
+    rav_request_guard: TokioMutex<()>,
+    unaggregated_receipts_guard: TokioMutex<()>,
 }
 
 impl SenderAllocation {
@@ -110,10 +109,11 @@ impl SenderAllocation {
             allocation_id,
             sender,
             sender_aggregator_endpoint,
-            unaggregated_fees: Arc::new(Mutex::new(UnaggregatedReceipts::default())),
+            unaggregated_fees: Arc::new(StdMutex::new(UnaggregatedReceipts::default())),
             config,
             escrow_accounts,
-            rav_request_guard: tokio::sync::Mutex::new(()),
+            rav_request_guard: TokioMutex::new(()),
+            unaggregated_receipts_guard: TokioMutex::new(()),
         };
 
         sender_allocation
@@ -133,6 +133,10 @@ impl SenderAllocation {
     /// Delete obsolete receipts in the DB w.r.t. the last RAV in DB, then update the tap manager
     /// with the latest unaggregated fees from the database.
     async fn update_unaggregated_fees(&self) -> Result<()> {
+        // Make sure to pause the handling of receipt notifications while we update the unaggregated
+        // fees.
+        let _guard = self.unaggregated_receipts_guard.lock().await;
+
         self.tap_manager.remove_obsolete_receipts().await?;
 
         let signers = signers_trimmed(&self.escrow_accounts, self.sender).await?;
@@ -176,19 +180,19 @@ impl SenderAllocation {
         .fetch_one(&self.pgpool)
         .await?;
 
-        let mut unaggregated_fees = self.unaggregated_fees.lock().unwrap();
-
         ensure!(
             res.sum.is_none() == res.max.is_none(),
             "Exactly one of SUM(value) and MAX(id) is null. This should not happen."
         );
 
-        unaggregated_fees.last_id = res.max.unwrap_or(0).try_into()?;
-        unaggregated_fees.value = res
-            .sum
-            .unwrap_or(BigDecimal::from(0))
-            .to_string()
-            .parse::<u128>()?;
+        *self.unaggregated_fees.lock().unwrap() = UnaggregatedReceipts {
+            last_id: res.max.unwrap_or(0).try_into()?,
+            value: res
+                .sum
+                .unwrap_or(BigDecimal::from(0))
+                .to_string()
+                .parse::<u128>()?,
+        };
 
         // TODO: check if we need to run a RAV request here.
 
@@ -360,22 +364,40 @@ impl SenderAllocation {
         Ok(())
     }
 
-    /// Safe add the fees to the unaggregated fees value, log an error if there is an overflow and
-    /// set the unaggregated fees value to u128::MAX.
-    pub fn fees_add(&self, fees: u128) {
+    /// Safe add the fees to the unaggregated fees value if the receipt_id is greater than the
+    /// last_id. If the addition would overflow u128, log an error and set the unaggregated fees
+    /// value to u128::MAX.
+    ///
+    /// Returns true if the fees were added, false otherwise.
+    pub async fn fees_add(&self, fees: u128, receipt_id: u64) -> bool {
+        // Make sure to pause the handling of receipt notifications while we update the unaggregated
+        // fees.
+        let _guard = self.unaggregated_receipts_guard.lock().await;
+
+        let mut fees_added = false;
         let mut unaggregated_fees = self.unaggregated_fees.lock().unwrap();
-        unaggregated_fees.value = unaggregated_fees
-            .value
-            .checked_add(fees)
-            .unwrap_or_else(|| {
-                // This should never happen, but if it does, we want to know about it.
-                error!(
-                    "Overflow when adding receipt value {} to total unaggregated fees {} for \
-                    allocation {} and sender {}. Setting total unaggregated fees to u128::MAX.",
-                    fees, unaggregated_fees.value, self.allocation_id, self.sender
-                );
-                u128::MAX
-            });
+
+        if receipt_id > unaggregated_fees.last_id {
+            *unaggregated_fees = UnaggregatedReceipts {
+                last_id: receipt_id,
+                value: unaggregated_fees
+                    .value
+                    .checked_add(fees)
+                    .unwrap_or_else(|| {
+                        // This should never happen, but if it does, we want to know about it.
+                        error!(
+                            "Overflow when adding receipt value {} to total unaggregated fees {} \
+                            for allocation {} and sender {}. Setting total unaggregated fees to \
+                            u128::MAX.",
+                            fees, unaggregated_fees.value, self.allocation_id, self.sender
+                        );
+                        u128::MAX
+                    }),
+            };
+            fees_added = true;
+        }
+
+        fees_added
     }
 
     pub fn get_unaggregated_fees(&self) -> UnaggregatedReceipts {