Skip to content

Commit

Permalink
refactor: ArbiterDB and Environment
Browse files Browse the repository at this point in the history
  • Loading branch information
@Autoparallel
Autoparallel committed Feb 29, 2024
1 parent 93ff285 commit 72bff89
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 56 deletions.
59 changes: 42 additions & 17 deletions core/src/database/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,25 +21,41 @@ use super::*;
pub mod fork;
pub mod inspector;

/// A [`ArbiterDB`] is a wrapper around a [`CacheDB`] that is used to provide
/// access to the [`environment::Environment`]'s database to multiple
/// A [`ArbiterDB`] is contains both a [`CacheDB`] that is used to provide
/// state for the [`environment::Environment`]'s as well as for multiple
/// [`coprocessor::Coprocessor`]s.
/// The `logs` field is a [`BTreeMap`] to store [`ethers::types::Log`]s that can be queried from at any point.
#[derive(Debug, Serialize, Deserialize)]
pub struct ArbiterDB(pub Arc<RwLock<CacheDB<EmptyDB>>>);
pub struct ArbiterDB {
/// The state of the `ArbiterDB`. This is a `CacheDB` that is used to provide
/// a db for the `Environment` to use.
pub state: Arc<RwLock<CacheDB<EmptyDB>>>,

/// The logs of the `ArbiterDB`. This is a `BTreeMap` that is used to store
/// logs that can be queried from at any point.
pub logs: Arc<RwLock<BTreeMap<U256, Vec<eLog>>>>,
}

// Implement `Clone` by hand so we utilize the `Arc`'s `Clone` implementation.
impl Clone for ArbiterDB {
fn clone(&self) -> Self {
Self(self.0.clone())
Self {
state: self.state.clone(),
logs: self.logs.clone(),
}
}
}

impl ArbiterDB {
/// Create a new `ArbiterDB`.
pub fn new() -> Self {
Self(Arc::new(RwLock::new(CacheDB::new(EmptyDB::new()))))
Self {
state: Arc::new(RwLock::new(CacheDB::new(EmptyDB::new()))),
logs: Arc::new(RwLock::new(BTreeMap::new())),
}
}

/// Write the `ArbiterDB` to a file at the given path.
/// Write the `ArbiterDB` to a file at the given path.``
pub fn write_to_file(&self, path: &str) -> io::Result<()> {
// Serialize the ArbiterDB
let serialized = serde_json::to_string(self)?;
Expand All @@ -55,9 +71,18 @@ impl ArbiterDB {
let mut file = fs::File::open(path)?;
let mut contents = String::new();
file.read_to_string(&mut contents)?;

// Deserialize the content into ArbiterDB
let cache_db = serde_json::from_str(&contents)?;
Ok(Self(Arc::new(RwLock::new(cache_db))))
#[derive(Deserialize)]
struct TempDB {
state: Option<CacheDB<EmptyDB>>,
logs: Option<BTreeMap<U256, Vec<eLog>>>,
}
let temp_db: TempDB = serde_json::from_str(&contents)?;
Ok(Self {
state: Arc::new(RwLock::new(temp_db.state.unwrap_or_default())),
logs: Arc::new(RwLock::new(temp_db.logs.unwrap_or_default())),
})
}
}

Expand All @@ -83,23 +108,23 @@ impl Database for ArbiterDB {
&mut self,
address: revm::primitives::Address,
) -> Result<Option<AccountInfo>, Self::Error> {
self.0.write().unwrap().basic(address)
self.state.write().unwrap().basic(address)
}

fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
self.0.write().unwrap().code_by_hash(code_hash)
self.state.write().unwrap().code_by_hash(code_hash)
}

fn storage(
&mut self,
address: revm::primitives::Address,
index: U256,
) -> Result<U256, Self::Error> {
self.0.write().unwrap().storage(address, index)
self.state.write().unwrap().storage(address, index)
}

fn block_hash(&mut self, number: U256) -> Result<B256, Self::Error> {
self.0.write().unwrap().block_hash(number)
self.state.write().unwrap().block_hash(number)
}
}

Expand All @@ -110,23 +135,23 @@ impl DatabaseRef for ArbiterDB {
&self,
address: revm::primitives::Address,
) -> Result<Option<AccountInfo>, Self::Error> {
self.0.read().unwrap().basic_ref(address)
self.state.read().unwrap().basic_ref(address)
}

fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
self.0.read().unwrap().code_by_hash_ref(code_hash)
self.state.read().unwrap().code_by_hash_ref(code_hash)
}

fn storage_ref(
&self,
address: revm::primitives::Address,
index: U256,
) -> Result<U256, Self::Error> {
self.0.read().unwrap().storage_ref(address, index)
self.state.read().unwrap().storage_ref(address, index)
}

fn block_hash_ref(&self, number: U256) -> Result<B256, Self::Error> {
self.0.read().unwrap().block_hash_ref(number)
self.state.read().unwrap().block_hash_ref(number)
}
}

Expand All @@ -135,7 +160,7 @@ impl DatabaseCommit for ArbiterDB {
&mut self,
changes: hashbrown::HashMap<revm::primitives::Address, revm::primitives::Account>,
) {
self.0.write().unwrap().commit(changes)
self.state.write().unwrap().commit(changes)
}
}

Expand Down
66 changes: 27 additions & 39 deletions core/src/environment/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,8 +77,6 @@ pub struct Environment {
/// calls and transactions.
pub(crate) db: ArbiterDB,

pub(crate) log_storage: Arc<RwLock<BTreeMap<U256, Vec<eLog>>>>,

inspector: Option<ArbiterInspector>,

/// This gives a means of letting the "outside world" connect to the
Expand Down Expand Up @@ -147,10 +145,22 @@ impl EnvironmentBuilder {
self
}

/// Sets the database for the [`Environment`]. This can come from a
/// [`fork::Fork`].
pub fn with_db(mut self, db: impl Into<CacheDB<EmptyDB>>) -> Self {
self.db = ArbiterDB(Arc::new(RwLock::new(db.into())));
/// Sets the state for the [`Environment`]. This can come from a saved state of a simulation
/// or a [`database::fork::Fork`].
pub fn with_state(mut self, state: impl Into<CacheDB<EmptyDB>>) -> Self {
self.db.state = Arc::new(RwLock::new(state.into()));
self
}

/// Sets the logs for the [`Environment`]. This can come from a saved state of a simulation and can be useful for doing analysis.
pub fn with_logs(mut self, logs: impl Into<BTreeMap<U256, Vec<eLog>>>) -> Self {
self.db.logs = Arc::new(RwLock::new(logs.into()));
self
}

/// Sets the entire database for the [`Environment`] including both the state and logs. This can come from the saved state of a simulation and can be useful for doing analysis.
pub fn with_arbiter_db(mut self, db: ArbiterDB) -> Self {
self.db = db;
self
}

Expand All @@ -175,7 +185,7 @@ impl Environment {
pub fn builder() -> EnvironmentBuilder {
EnvironmentBuilder {
parameters: EnvironmentParameters::default(),
db: ArbiterDB(Arc::new(RwLock::new(CacheDB::new(EmptyDB::new())))),
db: ArbiterDB::default(),
}
}

Expand All @@ -202,7 +212,6 @@ impl Environment {
inspector,
parameters,
db,
log_storage: Arc::new(RwLock::new(BTreeMap::new())),
handle: None,
}
}
Expand All @@ -215,7 +224,6 @@ impl Environment {

// Bring in the EVM db and log storage by cloning the interior Arc (lightweight).
let db = self.db.clone();
let log_storage = self.log_storage.clone();

// Bring in the EVM ENV
let mut env = Env::default();
Expand All @@ -232,7 +240,7 @@ impl Environment {
let handle = thread::spawn(move || {
// Create a new EVM builder
let mut evm = Evm::builder()
.with_db(db)
.with_db(db.clone())
.with_env(Box::new(env))
.with_external_context(inspector)
.append_handler_register(inspector_handle_register)
Expand Down Expand Up @@ -260,14 +268,7 @@ impl Environment {
account_state: AccountState::None,
storage: HashMap::new(),
};
let db = &mut evm.context.evm.db;
match db
.0
.write()
.unwrap()
.accounts
.insert(recast_address, account)
{
match db.state.write()?.accounts.insert(recast_address, account) {
None => outcome_sender.send(Ok(Outcome::AddAccountCompleted))?,
Some(_) => {
outcome_sender.send(Err(ArbiterCoreError::AccountCreationError))?;
Expand Down Expand Up @@ -307,13 +308,11 @@ impl Environment {
key,
block: _,
} => {
let db = &mut evm.context.evm.db;

let recast_address = Address::from(account.as_fixed_bytes());
let recast_key = B256::from(key.as_fixed_bytes()).into();

// Get the account storage value at the key in the db.
match db.0.write().unwrap().accounts.get_mut(&recast_address) {
match db.state.write().unwrap().accounts.get_mut(&recast_address) {
Some(account) => {
// Returns zero if the account is missing.
let value: U256 = match account.storage.get::<U256>(&recast_key)
Expand All @@ -336,16 +335,13 @@ impl Environment {
key,
value,
} => {
// Get the underlying database
let db = &mut evm.context.evm.db;

let recast_address = Address::from(account.as_fixed_bytes());
let recast_key = B256::from(key.as_fixed_bytes());
let recast_value = B256::from(value.as_fixed_bytes());

// Mutate the db by inserting the new key-value pair into the account's
// storage and send the successful CheatcodeCompleted outcome.
match db.0.write().unwrap().accounts.get_mut(&recast_address) {
match db.state.write().unwrap().accounts.get_mut(&recast_address) {
Some(account) => {
account
.storage
Expand All @@ -362,9 +358,8 @@ impl Environment {
};
}
Cheatcodes::Deal { address, amount } => {
let db = &mut evm.context.evm.db;
let recast_address = Address::from(address.as_fixed_bytes());
match db.0.write().unwrap().accounts.get_mut(&recast_address) {
match db.state.write().unwrap().accounts.get_mut(&recast_address) {
Some(account) => {
account.info.balance += U256::from_limbs(amount.0);
outcome_sender.send(Ok(Outcome::CheatcodeReturn(
Expand All @@ -378,10 +373,8 @@ impl Environment {
};
}
Cheatcodes::Access { address } => {
let db = &mut evm.context.evm.db;
let recast_address = Address::from(address.as_fixed_bytes());

match db.0.write().unwrap().accounts.get(&recast_address) {
match db.state.write().unwrap().accounts.get(&recast_address) {
Some(account) => {
let account_state = match account.account_state {
AccountState::None => AccountStateSerializable::None,
Expand Down Expand Up @@ -472,8 +465,7 @@ impl Environment {
cumulative_gas_per_block,
};

// TODO: Don't unwrap this.
let mut logs = log_storage.write().unwrap();
let mut logs = db.logs.write()?;
match logs.get_mut(&evm.block().number) {
Some(log_vec) => {
log_vec.extend(revm_logs_to_ethers_logs(
Expand Down Expand Up @@ -525,10 +517,8 @@ impl Environment {
Ok(Outcome::QueryReturn(evm.tx().gas_price.to_string()))
}
EnvironmentData::Balance(address) => {
// This unwrap should never fail.
let db = &mut evm.context.evm.db;
match db
.0
.state
.read()
.unwrap()
.accounts
Expand All @@ -541,9 +531,8 @@ impl Environment {
}
}
EnvironmentData::TransactionCount(address) => {
let db = &mut evm.context.evm.db;
match db
.0
.state
.read()
.unwrap()
.accounts
Expand All @@ -556,7 +545,7 @@ impl Environment {
}
}
EnvironmentData::Logs { filter } => {
let logs = log_storage.read().unwrap();
let logs = db.logs.read().unwrap();
let from_block = U256::from(
filter
.block_option
Expand Down Expand Up @@ -623,7 +612,6 @@ impl Environment {
warn!("Stop signal was not sent to any listeners. Are there any listeners?")
}
}
let (db, _) = evm.into_db_and_env_with_handler_cfg();
outcome_sender.send(Ok(Outcome::StopCompleted(db)))?;
break;
}
Expand Down
11 changes: 11 additions & 0 deletions core/src/errors.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
//! Errors that can occur when managing or interfacing with Arbiter's sandboxed
//! Ethereum environment.
use std::sync::{PoisonError, RwLockWriteGuard};

// use crossbeam_channel::SendError;
use crossbeam_channel::{RecvError, SendError};
use ethers::{
Expand Down Expand Up @@ -101,6 +103,9 @@ pub enum ArbiterCoreError {
/// Failed to reply to instruction.
#[error("{0}")]
ReplyError(String),

#[error("{0}")]
RwLockError(String),
}

impl From<SendError<Result<Outcome, ArbiterCoreError>>> for ArbiterCoreError {
Expand All @@ -109,6 +114,12 @@ impl From<SendError<Result<Outcome, ArbiterCoreError>>> for ArbiterCoreError {
}
}

impl<T> From<PoisonError<RwLockWriteGuard<'_, T>>> for ArbiterCoreError {
fn from(e: PoisonError<RwLockWriteGuard<'_, T>>) -> Self {
ArbiterCoreError::RwLockError(e.to_string())
}
}

impl MiddlewareError for ArbiterCoreError {
type Inner = ProviderError;

Expand Down

0 comments on commit 72bff89

Please sign in to comment.