opt: Concurrent build the output and reverts (#19)

Two optimization:
1. Concurrent build the output and reverts. According to the trace, currently building reverts does not consume time, but it takes time to transmit and insert `HashMap` through channels.
2. Use `AHashMap` instead of `HashMap`, because HashMap's performance is unstable.

Author: AshinGau

Cargo.toml

@@ -8,6 +8,8 @@ description = "Create Parallel EVM"
 revm = "14.0.0"
 fastrace = "0.7"
 tracing = "0.1.40"
+ahash = { version = "0.8.11", features = ["serde"] }
+rayon = "1.10.0"
 
 # Alloy
 alloy-chains = "0.1.18"

src/lib.rs

@@ -19,7 +19,6 @@ use revm::{
 use std::{
     cmp::min,
     fmt::{Display, Formatter},
-    sync::atomic::{AtomicUsize, Ordering},
     thread,
 };
 use tokio::runtime::{Builder, Runtime};
@@ -28,6 +27,7 @@ mod partition;
 mod scheduler;
 mod storage;
 mod tx_dependency;
+use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 
 lazy_static! {
     static ref CPU_CORES: usize = thread::available_parallelism().map(|n| n.get()).unwrap_or(8);
@@ -161,20 +161,14 @@ where
     F: Fn(usize, usize, usize) + Send + Sync + 'scope,
 {
     let parallel_cnt = num_partitions.unwrap_or(*CPU_CORES * 2 + 1);
-    let index = AtomicUsize::new(0);
     let remaining = num_elements % parallel_cnt;
     let chunk_size = num_elements / parallel_cnt;
-    thread::scope(|scope| {
-        for _ in 0..parallel_cnt {
-            scope.spawn(|| {
-                let index = index.fetch_add(1, Ordering::SeqCst);
-                let start_pos = chunk_size * index + min(index, remaining);
-                let mut end_pos = start_pos + chunk_size;
-                if index < remaining {
-                    end_pos += 1;
-                }
-                f(start_pos, end_pos, index);
-            });
+    (0..parallel_cnt).into_par_iter().for_each(|index| {
+        let start_pos = chunk_size * index + min(index, remaining);
+        let mut end_pos = start_pos + chunk_size;
+        if index < remaining {
+            end_pos += 1;
         }
+        f(start_pos, end_pos, index);
     });
 }

src/partition.rs

@@ -3,12 +3,13 @@ use crate::{
     LocationAndType, PartitionId, ResultAndTransition, SharedTxStates, TransactionStatus, TxId,
     TxState,
 };
+use ahash::{AHashMap as HashMap, AHashSet as HashSet};
 use revm::{
     primitives::{Address, EVMError, Env, ResultAndState, SpecId, TxEnv, TxKind},
     DatabaseRef, EvmBuilder,
 };
 use std::{
-    collections::{BTreeSet, HashMap, HashSet},
+    collections::BTreeSet,
     sync::Arc,
     time::{Duration, Instant},
 };

src/scheduler.rs

@@ -1,10 +1,3 @@
-use std::{
-    collections::{BTreeSet, HashMap, HashSet},
-    ops::DerefMut,
-    sync::{Arc, RwLock},
-    time::{Duration, Instant},
-};
-
 use crate::{
     fork_join_util,
     hint::ParallelExecutionHints,
@@ -14,7 +7,15 @@ use crate::{
     GrevmError, LocationAndType, ResultAndTransition, TransactionStatus, TxId, CPU_CORES,
     GREVM_RUNTIME, MAX_NUM_ROUND,
 };
+use fastrace::Span;
+use std::{
+    collections::BTreeSet,
+    ops::DerefMut,
+    sync::{Arc, RwLock},
+    time::{Duration, Instant},
+};
 
+use ahash::{AHashMap as HashMap, AHashSet as HashSet};
 use metrics::{counter, gauge};
 use revm::{
     db::{states::bundle_state::BundleRetention, BundleState},
@@ -238,6 +239,7 @@ where
     DB::Error: Send + Sync,
 {
     /// Creates a new GrevmScheduler instance.
+    #[fastrace::trace]
     pub fn new(spec_id: SpecId, env: Env, db: DB, txs: Arc<Vec<TxEnv>>) -> Self {
         let coinbase = env.block.coinbase;
         let num_partitions = *CPU_CORES * 2 + 1; // 2 * cpu + 1 for initial partition number
@@ -507,12 +509,16 @@ where
             unsafe { &mut *(&(*self.tx_states) as *const Vec<TxState> as *mut Vec<TxState>) };
         let mut miner_updates = Vec::with_capacity(finality_tx_cnt);
         let start_txid = self.num_finality_txs - finality_tx_cnt;
+
+        let span = Span::enter_with_local_parent("database commit transitions");
         for txid in start_txid..self.num_finality_txs {
             miner_updates.push(tx_states[txid].execute_result.miner_update.clone());
             database
                 .commit_transition(std::mem::take(&mut tx_states[txid].execute_result.transition));
             self.results.push(tx_states[txid].execute_result.result.clone().unwrap());
         }
+        drop(span);
+
         // Each transaction updates three accounts: from, to, and coinbase.
         // If every tx updates the coinbase account, it will cause conflicts across all txs.
         // Therefore, we handle miner rewards separately. We don't record miner’s address in r/w

src/storage.rs

@@ -1,16 +1,21 @@
-use crate::{LocationAndType, LocationSet};
+use crate::{fork_join_util, LocationAndType, LocationSet};
+use ahash::{AHashMap, AHashSet};
+use fastrace::Span;
 use revm::{
     db::{
         states::{bundle_state::BundleRetention, CacheAccount},
-        BundleState, PlainAccount,
+        AccountRevert, BundleAccount, BundleState, PlainAccount,
     },
     precompile::Address,
     primitives::{Account, AccountInfo, Bytecode, EvmState, B256, BLOCK_HASH_HISTORY, U256},
     CacheState, Database, DatabaseRef, TransitionAccount, TransitionState,
 };
 use std::{
-    collections::{btree_map, hash_map, BTreeMap, HashMap, HashSet},
-    sync::Arc,
+    collections::{btree_map, hash_map, BTreeMap, HashMap},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc, Mutex,
+    },
 };
 
 /// LazyUpdateValue is used to update the balance of the miner's account.
@@ -68,6 +73,87 @@ impl LazyUpdateValue {
     }
 }
 
+trait ParallelBundleState {
+    fn parallel_apply_transitions_and_create_reverts(
+        &mut self,
+        transitions: TransitionState,
+        retention: BundleRetention,
+    );
+}
+
+impl ParallelBundleState for BundleState {
+    #[fastrace::trace]
+    fn parallel_apply_transitions_and_create_reverts(
+        &mut self,
+        transitions: TransitionState,
+        retention: BundleRetention,
+    ) {
+        assert!(self.state.is_empty());
+        let include_reverts = retention.includes_reverts();
+        // pessimistically pre-allocate assuming _all_ accounts changed.
+        let reverts_capacity = if include_reverts { transitions.transitions.len() } else { 0 };
+        let transitions = transitions.transitions;
+        let addresses: Vec<Address> = transitions.keys().cloned().collect();
+        let reverts: Vec<Option<(Address, AccountRevert)>> = vec![None; reverts_capacity];
+        let bundle_state: Vec<Option<(Address, BundleAccount)>> = vec![None; transitions.len()];
+        let state_size = AtomicUsize::new(0);
+        let contracts = Mutex::new(HashMap::new());
+
+        let span = Span::enter_with_local_parent("parallel create reverts");
+        fork_join_util(transitions.len(), None, |start_pos, end_pos, _| {
+            #[allow(invalid_reference_casting)]
+            let reverts = unsafe {
+                &mut *(&reverts as *const Vec<Option<(Address, AccountRevert)>>
+                    as *mut Vec<Option<(Address, AccountRevert)>>)
+            };
+            #[allow(invalid_reference_casting)]
+            let addresses =
+                unsafe { &mut *(&addresses as *const Vec<Address> as *mut Vec<Address>) };
+            #[allow(invalid_reference_casting)]
+            let bundle_state = unsafe {
+                &mut *(&bundle_state as *const Vec<Option<(Address, BundleAccount)>>
+                    as *mut Vec<Option<(Address, BundleAccount)>>)
+            };
+
+            for pos in start_pos..end_pos {
+                let address = addresses[pos];
+                let transition = transitions.get(&address).cloned().unwrap();
+                // add new contract if it was created/changed.
+                if let Some((hash, new_bytecode)) = transition.has_new_contract() {
+                    contracts.lock().unwrap().insert(hash, new_bytecode.clone());
+                }
+                let present_bundle = transition.present_bundle_account();
+                let revert = transition.create_revert();
+                if let Some(revert) = revert {
+                    state_size.fetch_add(present_bundle.size_hint(), Ordering::Relaxed);
+                    bundle_state[pos] = Some((address, present_bundle));
+                    if include_reverts {
+                        reverts[pos] = Some((address, revert));
+                    }
+                }
+            }
+        });
+        self.state_size = state_size.load(Ordering::Acquire);
+        drop(span);
+
+        // much faster than bundle_state.into_iter().filter_map(|r| r).collect()
+        self.state.reserve(transitions.len());
+        for bundle in bundle_state {
+            if let Some((address, state)) = bundle {
+                self.state.insert(address, state);
+            }
+        }
+        let mut final_reverts = Vec::with_capacity(reverts_capacity);
+        for revert in reverts {
+            if let Some(r) = revert {
+                final_reverts.push(r);
+            }
+        }
+        self.reverts.push(final_reverts);
+        self.contracts = contracts.into_inner().unwrap();
+    }
+}
+
 /// SchedulerDB is a database wrapper that manages state transitions and caching for the EVM.
 /// It maintains a cache of committed data, a transition state for ongoing transactions, and a
 /// bundle state for finalizing block state changes. It also tracks block hashes for quick access.
@@ -143,7 +229,8 @@ impl<DB> SchedulerDB<DB> {
     #[fastrace::trace]
     pub(crate) fn merge_transitions(&mut self, retention: BundleRetention) {
         if let Some(transition_state) = self.transition_state.as_mut().map(TransitionState::take) {
-            self.bundle_state.apply_transitions_and_create_reverts(transition_state, retention);
+            self.bundle_state
+                .parallel_apply_transitions_and_create_reverts(transition_state, retention);
         }
     }
 }
@@ -345,7 +432,7 @@ pub(crate) struct PartitionDB<DB> {
     pub block_hashes: BTreeMap<u64, B256>,
 
     /// Record the read set of current tx, will be consumed after the execution of each tx
-    tx_read_set: HashMap<LocationAndType, Option<U256>>,
+    tx_read_set: AHashMap<LocationAndType, Option<U256>>,
 }
 
 impl<DB> PartitionDB<DB> {
@@ -355,12 +442,12 @@ impl<DB> PartitionDB<DB> {
             cache: CacheState::new(false),
             scheduler_db,
             block_hashes: BTreeMap::new(),
-            tx_read_set: HashMap::new(),
+            tx_read_set: AHashMap::new(),
         }
     }
 
     /// consume the read set after evm.transact() for each tx
-    pub(crate) fn take_read_set(&mut self) -> HashMap<LocationAndType, Option<U256>> {
+    pub(crate) fn take_read_set(&mut self) -> AHashMap<LocationAndType, Option<U256>> {
         core::mem::take(&mut self.tx_read_set)
     }
 
@@ -373,7 +460,7 @@ impl<DB> PartitionDB<DB> {
     ) -> (LocationSet, LazyUpdateValue, bool) {
         let mut miner_update = LazyUpdateValue::default();
         let mut remove_miner = true;
-        let mut write_set = HashSet::new();
+        let mut write_set = AHashSet::new();
         for (address, account) in &mut *changes {
             if account.is_selfdestructed() {
                 write_set.insert(LocationAndType::Code(*address));
@@ -483,9 +570,9 @@ where
     /// transaction.
     pub(crate) fn check_read_set(
         &mut self,
-        read_set: &HashMap<LocationAndType, Option<U256>>,
+        read_set: &AHashMap<LocationAndType, Option<U256>>,
     ) -> bool {
-        let mut visit_account = HashSet::new();
+        let mut visit_account = AHashSet::new();
         for (location, _) in read_set {
             match location {
                 LocationAndType::Basic(address) => {

src/tx_dependency.rs

@@ -1,13 +1,15 @@
 use smallvec::SmallVec;
 use std::{
     cmp::{min, Reverse},
-    collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, VecDeque},
+    collections::{BTreeMap, BTreeSet, BinaryHeap, VecDeque},
 };
 
 use crate::{fork_join_util, LocationAndType, SharedTxStates, TxId};
 
 pub(crate) type DependentTxsVec = SmallVec<[TxId; 1]>;
 
+use ahash::AHashMap as HashMap;
+
 const RAW_TRANSFER_WEIGHT: usize = 1;
 
 /// TxDependency is used to store the dependency relationship between transactions.