Cache org client and remove solana cache invalidation for invalids

iot_packet_verifier/src/balances.rs

@@ -5,7 +5,10 @@ use crate::{
 use futures_util::StreamExt;
 use helium_crypto::PublicKeyBinary;
 use solana::SolanaNetwork;
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::{hash_map::Entry, HashMap},
+    sync::Arc,
+};
 use tokio::sync::Mutex;
 
 /// Caches balances fetched from the solana chain and debits made by the
@@ -73,30 +76,30 @@ where
         &self,
         payer: &PublicKeyBinary,
         amount: u64,
+        trigger_balance_check_threshold: u64,
     ) -> Result<Option<u64>, S::Error> {
         let mut balances = self.balances.lock().await;
 
-        let balance = if !balances.contains_key(payer) {
-            let new_balance = self.solana.payer_balance(payer).await?;
-            balances.insert(payer.clone(), Balance::new(new_balance));
-            balances.get_mut(payer).unwrap()
-        } else {
-            let balance = balances.get_mut(payer).unwrap();
+        // Fetch the balance if we haven't seen the payer before
+        if let Entry::Vacant(balance) = balances.entry(payer.clone()) {
+            let balance = balance.insert(Balance::new(self.solana.payer_balance(payer).await?));
+            return Ok((balance.balance >= amount).then(|| {
+                balance.burned += amount;
+                balance.balance - amount
+            }));
+        }
 
-            // If the balance is not sufficient, check to see if it has been increased
-            if balance.balance < amount + balance.burned {
-                balance.balance = self.solana.payer_balance(payer).await?;
+        let balance = balances.get_mut(payer).unwrap();
+        match balance.balance.checked_sub(amount + balance.burned) {
+            Some(remaining_balance) => {
+                if remaining_balance < trigger_balance_check_threshold {
+                    balance.balance = self.solana.payer_balance(payer).await?;
+                }
+                balance.burned += amount;
+                Ok(Some(balance.balance - balance.burned))
             }
-
-            balance
-        };
-
-        Ok(if balance.balance >= amount + balance.burned {
-            balance.burned += amount;
-            Some(balance.balance - balance.burned)
-        } else {
-            None
-        })
+            None => Ok(None),
+        }
     }
 }
 

iot_packet_verifier/src/daemon.rs

@@ -2,7 +2,7 @@ use crate::{
     balances::BalanceCache,
     burner::Burner,
     settings::Settings,
-    verifier::{ConfigServer, Verifier},
+    verifier::{CachedOrgClient, ConfigServer, Verifier},
 };
 use anyhow::{bail, Error, Result};
 use file_store::{
@@ -21,11 +21,10 @@ use tokio::{
     signal,
     sync::{mpsc::Receiver, Mutex},
 };
-use tracing::debug;
 
 struct Daemon {
     pool: Pool<Postgres>,
-    verifier: Verifier<BalanceCache<Option<Arc<SolanaRpc>>>, Arc<Mutex<OrgClient>>>,
+    verifier: Verifier<BalanceCache<Option<Arc<SolanaRpc>>>, Arc<Mutex<CachedOrgClient>>>,
     report_files: Receiver<FileInfoStream<PacketRouterPacketReport>>,
     valid_packets: FileSinkClient,
     invalid_packets: FileSinkClient,
@@ -69,9 +68,7 @@ impl Daemon {
                 &self.invalid_packets,
             )
             .await?;
-        debug!("Committing transaction");
         transaction.commit().await?;
-        debug!("Committing files");
         self.valid_packets.commit().await?;
         self.invalid_packets.commit().await?;
 
@@ -159,9 +156,9 @@ impl Cmd {
         .create()
         .await?;
 
-        let org_client = Arc::new(Mutex::new(OrgClient::from_settings(
+        let org_client = Arc::new(Mutex::new(CachedOrgClient::new(OrgClient::from_settings(
             &settings.iot_config_client,
-        )?));
+        )?)));
 
         let file_store = FileStore::from_settings(&settings.ingest).await?;
 

iot_packet_verifier/src/verifier.rs

@@ -19,7 +19,6 @@ use tokio::{
     task::JoinError,
     time::{sleep_until, Duration, Instant},
 };
-use tracing::debug;
 
 pub struct Verifier<D, C> {
     pub debiter: D,
@@ -65,32 +64,25 @@ where
         tokio::pin!(reports);
 
         while let Some(report) = reports.next().await {
-            debug!(%report.received_timestamp, "Processing packet report");
-
             let debit_amount = payload_size_to_dc(report.payload_size as u64);
 
-            debug!(%report.oui, "Fetching payer");
             let payer = self
                 .config_server
                 .fetch_org(report.oui, &mut org_cache)
                 .await
                 .map_err(VerificationError::ConfigError)?;
-            debug!(%payer, "Debiting payer");
-            let remaining_balance = self
+
+            if let Some(remaining_balance) = self
                 .debiter
-                .debit_if_sufficient(&payer, debit_amount)
+                .debit_if_sufficient(&payer, debit_amount, minimum_allowed_balance)
                 .await
-                .map_err(VerificationError::DebitError)?;
-
-            if let Some(remaining_balance) = remaining_balance {
-                debug!(%debit_amount, "Adding debit amount to pending burns");
-
+                .map_err(VerificationError::DebitError)?
+            {
                 pending_burns
                     .add_burned_amount(&payer, debit_amount)
                     .await
                     .map_err(VerificationError::BurnError)?;
 
-                debug!("Writing valid packet report");
                 valid_packets
                     .write(ValidPacket {
                         packet_timestamp: report.timestamp(),
@@ -103,14 +95,12 @@ where
                     .map_err(VerificationError::ValidPacketWriterError)?;
 
                 if remaining_balance < minimum_allowed_balance {
-                    debug!(%report.oui, "Disabling org");
                     self.config_server
                         .disable_org(report.oui)
                         .await
                         .map_err(VerificationError::ConfigError)?;
                 }
             } else {
-                debug!("Writing invalid packet report");
                 invalid_packets
                     .write(InvalidPacket {
                         payload_size: report.payload_size,
@@ -145,6 +135,7 @@ pub trait Debiter {
         &self,
         payer: &PublicKeyBinary,
         amount: u64,
+        trigger_balance_check_threshold: u64,
     ) -> Result<Option<u64>, Self::Error>;
 }
 
@@ -156,6 +147,7 @@ impl Debiter for Arc<Mutex<HashMap<PublicKeyBinary, u64>>> {
         &self,
         payer: &PublicKeyBinary,
         amount: u64,
+        _trigger_balance_check_threshold: u64,
     ) -> Result<Option<u64>, Infallible> {
         let map = self.lock().await;
         let balance = map.get(payer).unwrap();
@@ -276,8 +268,22 @@ pub enum ConfigServerError {
     NotFound(u64),
 }
 
+pub struct CachedOrgClient {
+    client: OrgClient,
+    cache: HashMap<u64, bool>,
+}
+
+impl CachedOrgClient {
+    pub fn new(client: OrgClient) -> Self {
+        Self {
+            client,
+            cache: HashMap::new(),
+        }
+    }
+}
+
 #[async_trait]
-impl ConfigServer for Arc<Mutex<OrgClient>> {
+impl ConfigServer for Arc<Mutex<CachedOrgClient>> {
     type Error = ConfigServerError;
 
     async fn fetch_org(
@@ -289,6 +295,7 @@ impl ConfigServer for Arc<Mutex<OrgClient>> {
             let pubkey = PublicKeyBinary::from(
                 self.lock()
                     .await
+                    .client
                     .get(oui)
                     .await?
                     .org
@@ -301,19 +308,28 @@ impl ConfigServer for Arc<Mutex<OrgClient>> {
     }
 
     async fn disable_org(&self, oui: u64) -> Result<(), Self::Error> {
-        self.lock().await.disable(oui).await?;
+        let mut cached_client = self.lock().await;
+        if *cached_client.cache.entry(oui).or_insert(true) {
+            cached_client.client.disable(oui).await?;
+            *cached_client.cache.get_mut(&oui).unwrap() = false;
+        }
         Ok(())
     }
 
     async fn enable_org(&self, oui: u64) -> Result<(), Self::Error> {
-        self.lock().await.enable(oui).await?;
+        let mut cached_client = self.lock().await;
+        if !*cached_client.cache.entry(oui).or_insert(false) {
+            cached_client.client.enable(oui).await?;
+            *cached_client.cache.get_mut(&oui).unwrap() = true;
+        }
         Ok(())
     }
 
     async fn list_orgs(&self) -> Result<Vec<Org>, Self::Error> {
         Ok(self
             .lock()
             .await
+            .client
             .list()
             .await?
             .into_iter()

iot_packet_verifier/tests/integration_tests.rs

@@ -87,6 +87,7 @@ impl Debiter for InstantBurnedBalance {
         &self,
         payer: &PublicKeyBinary,
         amount: u64,
+        _trigger_balance_check_threshold: u64,
     ) -> Result<Option<u64>, ()> {
         let map = self.0.lock().await;
         let balance = map.get(payer).unwrap();