refactor: make EscrowAccounts attributes private

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

common/src/escrow_accounts.rs

@@ -1,7 +1,10 @@
 // Copyright 2023-, GraphOps and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::{collections::HashMap, time::Duration};
+use std::{
+    collections::{HashMap, HashSet},
+    time::Duration,
+};
 
 use alloy_primitives::Address;
 use anyhow::Result;
@@ -15,9 +18,9 @@ use crate::prelude::{Query, SubgraphClient};
 
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub struct EscrowAccounts {
-    pub senders_balances: HashMap<Address, U256>,
-    pub signers_to_senders: HashMap<Address, Address>,
-    pub senders_to_signers: HashMap<Address, Vec<Address>>,
+    senders_balances: HashMap<Address, U256>,
+    signers_to_senders: HashMap<Address, Address>,
+    senders_to_signers: HashMap<Address, Vec<Address>>,
 }
 
 impl EscrowAccounts {
@@ -36,6 +39,47 @@ impl EscrowAccounts {
             senders_to_signers,
         }
     }
+
+    pub fn get_signers_for_sender(&self, sender: &Address) -> Result<Vec<Address>> {
+        self.senders_to_signers
+            .get(sender)
+            .filter(|signers| !signers.is_empty())
+            .ok_or(anyhow::format_err!(
+                "No signers found for sender {}.",
+                sender
+            ))
+            .map(|signers| signers.to_owned())
+    }
+
+    pub fn get_sender_for_signer(&self, signer: &Address) -> Result<Address> {
+        self.signers_to_senders
+            .get(signer)
+            .ok_or(anyhow::format_err!(
+                "Sender not found for receipt signer {}.",
+                signer
+            ))
+            .copied()
+    }
+
+    pub fn get_balance_for_sender(&self, sender: &Address) -> Result<U256> {
+        self.senders_balances
+            .get(sender)
+            .ok_or(anyhow::format_err!(
+                "Balance not found for sender {}.",
+                sender
+            ))
+            .copied()
+    }
+
+    pub fn get_balance_for_signer(&self, signer: &Address) -> Result<U256> {
+        self.get_sender_for_signer(signer)
+            .and_then(|sender| self.get_balance_for_sender(&sender))
+            .map_err(|e| anyhow::format_err!("Could not get balance for signer {}: {}", signer, e))
+    }
+
+    pub fn get_senders(&self) -> HashSet<Address> {
+        self.senders_balances.keys().copied().collect()
+    }
 }
 
 pub fn escrow_accounts(

common/src/tap_manager.rs

@@ -69,25 +69,14 @@ impl TapManager {
 
         let escrow_accounts = self.escrow_accounts.value_immediate().unwrap_or_default();
 
-        let receipt_sender = escrow_accounts
-            .signers_to_senders
-            .get(&receipt_signer)
-            .ok_or_else(|| {
-                anyhow!(
-                    "Receipt signer `{}` is not eligible for this indexer",
-                    receipt_signer
-                )
-            })?;
-
         if !escrow_accounts
-            .senders_balances
-            .get(receipt_sender)
-            .map_or(false, |balance| balance > &U256::zero())
+            .get_balance_for_signer(&receipt_signer)
+            .map_or(false, |balance| balance > U256::zero())
         {
-            return Err(anyhow!(
+            anyhow::bail!(
                 "Receipt sender `{}` is not eligible for this indexer",
-                receipt_signer
-            ));
+                receipt_signer,
+            );
         }
 
         // TODO: consider doing this in another async task to avoid slowing down the paid query flow.

tap-agent/src/tap/escrow_adapter.rs

@@ -53,27 +53,16 @@ impl EscrowAdapterTrait for EscrowAdapter {
                     error: format!("Could not get escrow accounts from eventual: {:?}.", e),
                 })?;
 
-        let sender =
-            escrow_accounts
-                .signers_to_senders
-                .get(&sender)
-                .ok_or(AdapterError::AdapterError {
-                    error: format!(
-                        "Sender {} not found for receipt signer, could not get available escrow.",
-                        sender
-                    )
-                    .to_string(),
-                })?;
+        let sender = escrow_accounts
+            .get_sender_for_signer(&sender)
+            .map_err(|e| AdapterError::AdapterError {
+                error: format!("{}", e).to_string(),
+            })?;
 
         let balance = escrow_accounts
-            .senders_balances
-            .get(sender)
-            .ok_or(AdapterError::AdapterError {
-                error: format!(
-                    "Sender {} not found in escrow balances map, could not get available escrow.",
-                    sender
-                )
-                .to_string(),
+            .get_balance_for_sender(&sender)
+            .map_err(|e| AdapterError::AdapterError {
+                error: format!("Could not get available escrow: {}", e).to_string(),
             })?
             .to_owned();
         let balance: u128 = balance.try_into().map_err(|_| AdapterError::AdapterError {
@@ -89,7 +78,7 @@ impl EscrowAdapterTrait for EscrowAdapter {
             .sender_pending_fees
             .read()
             .await
-            .get(sender)
+            .get(&sender)
             .copied()
             .unwrap_or(0);
         Ok(balance - fees)
@@ -106,17 +95,15 @@ impl EscrowAdapterTrait for EscrowAdapter {
 
         let current_available_escrow = self.get_available_escrow(sender).await?;
 
-        let sender =
-            escrow_accounts
-                .signers_to_senders
-                .get(&sender)
-                .ok_or(AdapterError::AdapterError {
-                    error: format!(
-                        "Sender {} not found for receipt signer, could not get available escrow.",
-                        sender
-                    )
-                    .to_string(),
-                })?;
+        let sender = escrow_accounts
+            .get_sender_for_signer(&sender)
+            .map_err(|e| AdapterError::AdapterError {
+                error: format!(
+                    "Could not get available escrow for receipt signer {}: {}",
+                    sender, e
+                )
+                .to_string(),
+            })?;
 
         let mut fees_write = self.sender_pending_fees.write().await;
         let fees = fees_write.entry(sender.to_owned()).or_insert(0);

tap-agent/src/tap/mod.rs

@@ -24,9 +24,7 @@ async fn signers_trimmed(
         .value()
         .await
         .map_err(|e| anyhow!("Error while getting escrow accounts: {:?}", e))?
-        .senders_to_signers
-        .get(&sender)
-        .ok_or(anyhow!("No signers found for sender {}.", sender))?
+        .get_signers_for_sender(&sender)?
         .iter()
         .map(|s| s.to_string().trim_start_matches("0x").to_owned())
         .collect::<Vec<String>>();

tap-agent/src/tap/receipt_checks_adapter.rs

@@ -126,10 +126,8 @@ impl ReceiptChecksAdapterTrait for ReceiptChecksAdapter {
                 })?;
 
         Ok(escrow_accounts
-            .signers_to_senders
-            .get(&sender_id)
-            .and_then(|sender| escrow_accounts.senders_balances.get(sender))
-            .map(|balance| *balance > U256::from(0))
+            .get_balance_for_signer(&sender_id)
+            .map(|balance| balance > U256::from(0))
             .unwrap_or(false))
     }
 }

tap-agent/src/tap/receipt_storage_adapter.rs

@@ -217,12 +217,7 @@ mod test {
         received_receipt_vec: &[(u64, ReceivedReceipt)],
         range: R,
     ) -> Result<()> {
-        let signers_to_senders = escrow_accounts
-            .value()
-            .await
-            .unwrap()
-            .signers_to_senders
-            .to_owned();
+        let escrow_accounts_snapshot = escrow_accounts.value().await.unwrap();
 
         // Filtering the received receipts by timestamp range
         let received_receipt_vec: Vec<(u64, ReceivedReceipt)> = received_receipt_vec
@@ -231,14 +226,14 @@ mod test {
                 range.contains(&received_receipt.signed_receipt().message.timestamp_ns)
                     && (received_receipt.signed_receipt().message.allocation_id
                         == storage_adapter.allocation_id)
-                    && (signers_to_senders
-                        .get(
+                    && (escrow_accounts_snapshot
+                        .get_sender_for_signer(
                             &received_receipt
                                 .signed_receipt()
                                 .recover_signer(&TAP_EIP712_DOMAIN_SEPARATOR)
                                 .unwrap(),
                         )
-                        .map_or(false, |v| *v == storage_adapter.sender))
+                        .map_or(false, |v| v == storage_adapter.sender))
             })
             .cloned()
             .collect();
@@ -280,12 +275,7 @@ mod test {
         received_receipt_vec: &[ReceivedReceipt],
         range: R,
     ) -> Result<()> {
-        let signers_to_senders = escrow_accounts
-            .value()
-            .await
-            .unwrap()
-            .signers_to_senders
-            .to_owned();
+        let escrow_accounts_snapshot = escrow_accounts.value().await.unwrap();
 
         // Storing the receipts
         let mut received_receipt_id_vec = Vec::new();
@@ -310,14 +300,14 @@ mod test {
             .filter(|(_, received_receipt)| {
                 if (received_receipt.signed_receipt().message.allocation_id
                     == storage_adapter.allocation_id)
-                    && (signers_to_senders
-                        .get(
+                    && (escrow_accounts_snapshot
+                        .get_sender_for_signer(
                             &received_receipt
                                 .signed_receipt()
                                 .recover_signer(&TAP_EIP712_DOMAIN_SEPARATOR)
                                 .unwrap(),
                         )
-                        .map_or(false, |v| *v == storage_adapter.sender))
+                        .map_or(false, |v| v == storage_adapter.sender))
                 {
                     !range.contains(&received_receipt.signed_receipt().message.timestamp_ns)
                 } else {

tap-agent/src/tap/sender_allocation_relationships_manager.rs

@@ -87,11 +87,7 @@ impl SenderAllocationRelationshipsManager {
                 .value()
                 .await
                 .expect("Should get indexer allocations from Eventual"),
-            escrow_accounts_snapshot
-                .senders_balances
-                .keys()
-                .copied()
-                .collect(),
+            escrow_accounts_snapshot.get_senders(),
         )
         .await
         .expect("Should be able to update sender_allocation_relationships");
@@ -132,9 +128,7 @@ impl SenderAllocationRelationshipsManager {
             let signer = Address::from_str(&row.signer_address)
                 .expect("signer_address should be a valid address");
             let sender = escrow_accounts_snapshot
-                .signers_to_senders
-                .get(&signer)
-                .copied()
+                .get_sender_for_signer(&signer)
                 .expect("should be able to get sender from signer");
 
             // Only create a SenderAllocationRelationship if it doesn't exist yet.
@@ -190,7 +184,7 @@ impl SenderAllocationRelationshipsManager {
                 Self::update_sender_allocation_relationships(
                     &inner,
                     indexer_allocations,
-                    escrow_accounts.senders_balances.keys().copied().collect(),
+                    escrow_accounts.get_senders(),
                 )
                 .await
                 .unwrap_or_else(|e| {
@@ -234,13 +228,11 @@ impl SenderAllocationRelationshipsManager {
                 .value()
                 .await
                 .expect("should be able to get escrow accounts")
-                .signers_to_senders
-                .get(&new_receipt_notification.signer_address)
-                .copied();
+                .get_sender_for_signer(&new_receipt_notification.signer_address);
 
             let sender_address = match sender_address {
-                Some(sender_address) => sender_address,
-                None => {
+                Ok(sender_address) => sender_address,
+                Err(_) => {
                     error!(
                         "No sender address found for receipt signer address {}. \
                         This should not happen.",