refactor: use arc<rwlock<hashmap<...>>> instead of dashmap

Author: JuArce

crates/batcher/src/lib.rs

@@ -17,7 +17,6 @@ use types::batch_state::BatchState;
 use types::user_state::UserState;
 
 use batch_queue::calculate_batch_size;
-use dashmap::DashMap;
 use std::collections::HashMap;
 use std::env;
 use std::net::SocketAddr;
@@ -111,7 +110,7 @@ pub struct Batcher {
     /// It's used a way to "lock" all the user_states at the same time
     /// If one needed is taken in the handle message it will timeout
     is_recovering_from_submission_failure: RwLock<bool>,
-    user_states: DashMap<Address, Arc<Mutex<UserState>>>,
+    user_states: Arc<tokio::sync::RwLock<HashMap<Address, Arc<Mutex<UserState>>>>>,
 
     last_uploaded_batch_block: Mutex<u64>,
 
@@ -262,7 +261,7 @@ impl Batcher {
         .await
         .expect("Failed to get fallback Service Manager contract");
 
-        let user_states = DashMap::new();
+        let user_states = Arc::new(tokio::sync::RwLock::new(HashMap::new()));
         let batch_state = BatchState::new(config.batcher.max_queue_size);
         let non_paying_config = if let Some(non_paying_config) = config.batcher.non_paying {
             warn!("Non-paying address configuration detected. Will replace non-paying address {} with configured address.",
@@ -276,7 +275,7 @@ impl Batcher {
                 .expect("Could not get non-paying nonce from Ethereum");
 
             let non_paying_user_state = UserState::new(nonpaying_nonce);
-            user_states.insert(
+            user_states.write().await.insert(
                 non_paying_config.replacement.address(),
                 Arc::new(Mutex::new(non_paying_user_state)),
             );
@@ -335,7 +334,7 @@ impl Batcher {
         }
     }
 
-    fn update_evicted_user_state_with_lock(
+    async fn update_evicted_user_state_with_lock(
         &self,
         removed_entry: &types::batch_queue::BatchQueueEntry,
         batch_queue: &types::batch_queue::BatchQueue,
@@ -350,7 +349,7 @@ impl Batcher {
         {
             Some((last_entry, _)) => last_entry.nonced_verification_data.max_fee,
             None => {
-                self.user_states.remove(&addr);
+                self.user_states.write().await.remove(&addr);
                 return;
             }
         };
@@ -376,12 +375,12 @@ impl Batcher {
         {
             Some((last_entry, _)) => last_entry.nonced_verification_data.max_fee,
             None => {
-                self.user_states.remove(&addr);
+                self.user_states.write().await.remove(&addr);
                 return Some(());
             }
         };
 
-        let user_state = self.user_states.get(&addr)?;
+        let user_state = self.user_states.read().await.get(&addr)?.clone();
         let mut user_state_guard = user_state.lock().await;
         user_state_guard.proofs_in_batch -= 1;
         user_state_guard.nonce -= U256::one();
@@ -432,7 +431,7 @@ impl Batcher {
     where
         F: std::future::Future<Output = T>,
     {
-        match timeout(Duration::from_secs(15), lock_future).await {
+        match timeout(Duration::from_secs(150), lock_future).await {
             Ok(result) => Some(result),
             Err(_) => {
                 warn!("Batch lock acquisition timed out");
@@ -735,7 +734,7 @@ impl Batcher {
         }
 
         let cached_user_nonce = {
-            let user_state_ref = self.user_states.get(&address);
+            let user_state_ref = self.user_states.read().await.get(&address).cloned();
             match user_state_ref {
                 Some(user_state_ref) => {
                     let Some(user_state_guard) = self
@@ -881,18 +880,18 @@ impl Batcher {
         // If it was not present, then the user nonce is queried to the Aligned contract.
         // Lastly, we get a lock of the batch state again and insert the user state if it was still missing.
 
-        let is_user_in_state = self.user_states.contains_key(&addr);
+        let is_user_in_state = self.user_states.read().await.contains_key(&addr);
 
         if !is_user_in_state {
             warn!("User state for address {addr:?} not found, creating a new one");
             // We add a dummy user state to grab a lock on the user state
             let dummy_user_state = UserState::new(U256::zero());
             self.user_states
-                .insert(addr, Arc::new(Mutex::new(dummy_user_state)));
+                .write().await.insert(addr, Arc::new(Mutex::new(dummy_user_state)));
             warn!("Dummy user state for address {addr:?} created");
         }
 
-        let Some(user_state_ref) = self.user_states.get(&addr) else {
+        let Some(user_state_ref) = self.user_states.read().await.get(&addr).cloned() else {
             error!("This should never happen, user state has previously been inserted if it didn't exist");
             send_message(
                 ws_conn_sink.clone(),
@@ -1060,7 +1059,7 @@ impl Batcher {
 
             // Try to find any candidate whose lock we can acquire and immediately process them
             for candidate_addr in eviction_candidates {
-                if let Some(user_state_arc) = self.user_states.get(&candidate_addr) {
+                if let Some(user_state_arc) = self.user_states.read().await.get(&candidate_addr).cloned() {
                     if let Ok(mut user_guard) = user_state_arc.try_lock() {
                         // Found someone whose lock we can get - now find and remove their entry
                         let entries_to_check: Vec<_> = batch_state_lock
@@ -1094,7 +1093,7 @@ impl Batcher {
                                     &removed,
                                     &batch_state_lock.batch_queue,
                                     &mut user_guard,
-                                );
+                                ).await;
 
                                 if let Some(ref removed_entry_ws) = removed.messaging_sink {
                                     let ws_sink = removed_entry_ws.clone();
@@ -1553,7 +1552,7 @@ impl Batcher {
     ) -> Result<(), BatcherError> {
         // Update each user's state with proper lock ordering
         for addr in affected_users {
-            if let Some(user_state) = self.user_states.get(&addr) {
+            if let Some(user_state) = self.user_states.read().await.get(&addr).cloned() {
                 let mut user_state_guard = user_state.lock().await; // First: user lock
                 let batch_state_lock = self.batch_state.lock().await; // Second: batch lock
 
@@ -1588,7 +1587,7 @@ impl Batcher {
     /// Cleans up user states after successful batch submission.
     /// Resets last_max_fee_limit to U256::MAX for users who had proofs in the submitted batch
     /// but now have no proofs left in the queue.
-    fn cleanup_user_states_after_successful_submission(&self, finalized_batch: &[BatchQueueEntry]) {
+    async fn cleanup_user_states_after_successful_submission(&self, finalized_batch: &[BatchQueueEntry]) {
         use std::collections::HashSet;
 
         // Get unique users from the submitted batch
@@ -1612,7 +1611,7 @@ impl Batcher {
         for user_addr in users_in_batch {
             if !current_user_states.contains_key(&user_addr) {
                 // User has no proofs left in queue - reset their max_fee_limit
-                if let Some(user_state_ref) = self.user_states.get(&user_addr) {
+                if let Some(user_state_ref) = self.user_states.read().await.get(&user_addr).cloned() {
                     if let Ok(mut user_state_guard) = user_state_ref.try_lock() {
                         user_state_guard.last_max_fee_limit = U256::max_value();
                     }
@@ -1833,7 +1832,7 @@ impl Batcher {
         }
 
         // Clean up user states for users who had proofs in this batch but now have no proofs left
-        self.cleanup_user_states_after_successful_submission(finalized_batch);
+        self.cleanup_user_states_after_successful_submission(finalized_batch).await;
 
         connection::send_batch_inclusion_data_responses(finalized_batch, &batch_merkle_tree).await
     }
@@ -1851,7 +1850,7 @@ impl Batcher {
 
         let Some(nonpaying_replacement_addr) = self.get_nonpaying_replacement_addr() else {
             batch_state_lock.batch_queue.clear();
-            self.user_states.clear();
+            self.user_states.write().await.clear();
             return;
         };
 
@@ -1863,13 +1862,13 @@ impl Batcher {
             .await
         else {
             batch_state_lock.batch_queue.clear();
-            self.user_states.clear();
+            self.user_states.write().await.clear();
             return;
         };
         batch_state_lock.batch_queue.clear();
-        self.user_states.clear();
+        self.user_states.write().await.clear();
         let nonpaying_user_state = UserState::new(nonpaying_replacement_addr_nonce);
-        self.user_states.insert(
+        self.user_states.write().await.insert(
             nonpaying_replacement_addr,
             Arc::new(Mutex::new(nonpaying_user_state)),
         );