dev: fix to use rwlock and use it properly

flashbots · Jul 7, 2024 · 819070a · 819070a
1 parent 882ce81
commit 819070a
Showing 1 changed file with 26 additions and 18 deletions.
diff --git a/crates/rbuilder/src/backtest/fetch/mod.rs b/crates/rbuilder/src/backtest/fetch/mod.rs
@@ -11,7 +11,11 @@ use eyre::WrapErr;
 use flashbots_db::RelayDB;
 use futures::TryStreamExt;
 use sqlx::PgPool;
-use std::{collections::HashMap, path::PathBuf, sync::Arc};
+use std::{
+    collections::HashMap,
+    path::PathBuf,
+    sync::{Arc, RwLock},
+};
 use time::{Duration, OffsetDateTime};
 use tokio::sync::Mutex;
 use tracing::{info, trace};
@@ -165,7 +169,7 @@ impl HistoricalDataFetcher {
 
         let parent_block = block_number - 1;
 
-        let nonce_cache = Arc::new(Mutex::new(HashMap::new()));
+        let nonce_cache = Arc::new(RwLock::new(HashMap::new()));
         let retain = Arc::new(Mutex::new(vec![false; nonces_to_check.len()]));
 
         let retain_clone = retain.clone();
@@ -176,29 +180,33 @@ impl HistoricalDataFetcher {
                 async move {
                     let mut all_nonces_failed = true;
                     for nonce in nonces {
-                        let onchain_nonce = {
-                            let mut nonce_cache = nonce_cache.lock().await;
+                        let mut res_onchain_nonce = 0_u64;
+                        if let Ok(nonce_cache) = nonce_cache.read() {
                             if let Some(onchain_nonce) = nonce_cache.get(&nonce.address) {
-                                *onchain_nonce
-                            } else {
-                                let address =
-                                    alloy_primitives::Address::from_slice(&nonce.address.0 .0);
-                                let onchain_nonce = self
-                                    .eth_provider
-                                    .get_transaction_count(address)
-                                    .block_id(BlockId::Number(parent_block.into()))
-                                    .await
-                                    .wrap_err("Failed to fetch onchain tx count")?;
+                                res_onchain_nonce = *onchain_nonce
+                            }
+                        }
+                        if res_onchain_nonce == 0 {
+                            let address = nonce.address;
+                            let onchain_nonce = self
+                                .eth_provider
+                                .get_transaction_count(address)
+                                .block_id(BlockId::Number(parent_block.into()))
+                                .await
+                                .wrap_err("Failed to fetch onchain tx count")?;
+
+                            if let Ok(mut nonce_cache) = nonce_cache.write() {
                                 nonce_cache.insert(nonce.address, onchain_nonce);
-                                onchain_nonce
                             }
-                        };
-                        if onchain_nonce > nonce.nonce && !nonce.optional {
+                            res_onchain_nonce = onchain_nonce;
+                        }
+
+                        if res_onchain_nonce > nonce.nonce && !nonce.optional {
                             trace!(
                                 "Order nonce too low, order: {:?}, nonce: {}, onchain tx count: {}",
                                 id,
                                 nonce.nonce,
-                                onchain_nonce,
+                                res_onchain_nonce,
                             );
                             return Ok(());
                         } else {