Skip to content

Commit

Permalink
chore: store tokio::runtime::Handle in ethers/alloyDB (#1557)
Browse files Browse the repository at this point in the history 
* store tokio::runtime::Handle in ethers/alloyDB

* update docs

* fmt

* modify constructor to accept &self in block_on

* fix
  • Loading branch information
@chirag-bgh
chirag-bgh authored Jun 28, 2024
1 parent 64dc252 commit b40e4e3
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 73 deletions.
55 changes: 22 additions & 33 deletions crates/revm/src/db/alloydb.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use alloy_eips::BlockId;
use alloy_provider::{Network, Provider};
use alloy_transport::{Transport, TransportError};
use std::future::IntoFuture;
use tokio::runtime::{Builder, Handle};
use tokio::runtime::Handle;

/// An alloy-powered REVM [Database].
///
Expand All @@ -17,50 +17,39 @@ pub struct AlloyDB<T: Transport + Clone, N: Network, P: Provider<T, N>> {
provider: P,
/// The block number on which the queries will be based on.
block_number: BlockId,
/// handle to the tokio runtime
handle: Handle,
_marker: std::marker::PhantomData<fn() -> (T, N)>,
}

impl<T: Transport + Clone, N: Network, P: Provider<T, N>> AlloyDB<T, N, P> {
/// Create a new AlloyDB instance, with a [Provider] and a block (Use None for latest).
pub fn new(provider: P, block_number: BlockId) -> Self {
Self {
///
/// Returns `None` if no tokio runtime is available or if the current runtime is a current-thread runtime.
pub fn new(provider: P, block_number: BlockId) -> Option<Self> {
let handle = match Handle::try_current() {
Ok(handle) => match handle.runtime_flavor() {
tokio::runtime::RuntimeFlavor::CurrentThread => return None,
_ => handle,
},
Err(_) => return None,
};
Some(Self {
provider,
block_number,
handle,
_marker: std::marker::PhantomData,
}
})
}

/// Internal utility function that allows us to block on a future regardless of the runtime flavor.
#[inline]
fn block_on<F>(f: F) -> F::Output
fn block_on<F>(&self, f: F) -> F::Output
where
F: std::future::Future + Send,
F::Output: Send,
{
match Handle::try_current() {
Ok(handle) => match handle.runtime_flavor() {
// This is essentially equal to tokio::task::spawn_blocking because tokio doesn't
// allow the current_thread runtime to block_in_place.
// See <https://docs.rs/tokio/latest/tokio/task/fn.block_in_place.html> for more info.
tokio::runtime::RuntimeFlavor::CurrentThread => std::thread::scope(move |s| {
s.spawn(move || {
Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(f)
})
.join()
.unwrap()
}),
_ => tokio::task::block_in_place(move || handle.block_on(f)),
},
Err(_) => Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(f),
}
tokio::task::block_in_place(move || self.handle.block_on(f))
}

/// Set the block number on which the queries will be based on.
Expand Down Expand Up @@ -93,7 +82,7 @@ impl<T: Transport + Clone, N: Network, P: Provider<T, N>> DatabaseRef for AlloyD
)
};

let (nonce, balance, code) = Self::block_on(f);
let (nonce, balance, code) = self.block_on(f);

let balance = balance?;
let code = Bytecode::new_raw(code?.0.into());
Expand All @@ -104,7 +93,7 @@ impl<T: Transport + Clone, N: Network, P: Provider<T, N>> DatabaseRef for AlloyD
}

fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
let block = Self::block_on(
let block = self.block_on(
self.provider
// SAFETY: We know number <= u64::MAX, so we can safely convert it to u64
.get_block_by_number(number.into(), false),
Expand All @@ -123,7 +112,7 @@ impl<T: Transport + Clone, N: Network, P: Provider<T, N>> DatabaseRef for AlloyD
.provider
.get_storage_at(address, index)
.block_id(self.block_number);
let slot_val = Self::block_on(f.into_future())?;
let slot_val = self.block_on(f.into_future())?;
Ok(slot_val)
}
}
Expand Down Expand Up @@ -172,7 +161,7 @@ mod tests {
.parse()
.unwrap();

let acc_info = alloydb.basic_ref(address).unwrap().unwrap();
let acc_info = alloydb.unwrap().basic_ref(address).unwrap().unwrap();
assert!(acc_info.exists());
}
}
74 changes: 34 additions & 40 deletions crates/revm/src/db/ethersdb.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use std::sync::Arc;

use ethers_core::types::{Block, BlockId, TxHash, H160 as eH160, H256, U64 as eU64};
use ethers_providers::Middleware;
use tokio::runtime::{Builder, Handle, RuntimeFlavor};
use tokio::runtime::Handle;

use crate::primitives::{AccountInfo, Address, Bytecode, B256, U256};
use crate::{Database, DatabaseRef};
Expand All @@ -11,55 +11,49 @@ use crate::{Database, DatabaseRef};
pub struct EthersDB<M: Middleware> {
client: Arc<M>,
block_number: Option<BlockId>,
handle: Handle,
}

impl<M: Middleware> EthersDB<M> {
/// create ethers db connector inputs are url and block on what we are basing our database (None for latest)
/// Create ethers db connector inputs are url and block on what we are basing our database (None for latest).
///
/// Returns `None` if no tokio runtime is available or if the current runtime is a current-thread runtime.
pub fn new(client: Arc<M>, block_number: Option<BlockId>) -> Option<Self> {
let block_number: Option<BlockId> = if block_number.is_some() {
block_number
} else {
Some(BlockId::from(
Self::block_on(client.get_block_number()).ok()?,
))
let handle = match Handle::try_current() {
Ok(handle) => match handle.runtime_flavor() {
tokio::runtime::RuntimeFlavor::CurrentThread => return None,
_ => handle,
},
Err(_) => return None,
};

Some(Self {
client,
block_number,
})
if block_number.is_some() {
Some(Self {
client,
block_number,
handle,
})
} else {
let mut instance = Self {
client: client.clone(),
block_number: None,
handle,
};
instance.block_number = Some(BlockId::from(
instance.block_on(client.get_block_number()).ok()?,
));
Some(instance)
}
}

/// internal utility function to call tokio feature and wait for output
/// Internal utility function to call tokio feature and wait for output
#[inline]
fn block_on<F>(f: F) -> F::Output
fn block_on<F>(&self, f: F) -> F::Output
where
F: core::future::Future + Send,
F::Output: Send,
{
match Handle::try_current() {
Ok(handle) => match handle.runtime_flavor() {
// This essentially equals to tokio::task::spawn_blocking because tokio doesn't
// allow current_thread runtime to block_in_place
RuntimeFlavor::CurrentThread => std::thread::scope(move |s| {
s.spawn(move || {
Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(f)
})
.join()
.unwrap()
}),
_ => tokio::task::block_in_place(move || handle.block_on(f)),
},
Err(_) => Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(f),
}
tokio::task::block_in_place(move || self.handle.block_on(f))
}

/// set block number on which upcoming queries will be based
Expand All @@ -81,7 +75,7 @@ impl<M: Middleware> DatabaseRef for EthersDB<M> {
let code = self.client.get_code(add, self.block_number);
tokio::join!(nonce, balance, code)
};
let (nonce, balance, code) = Self::block_on(f);
let (nonce, balance, code) = self.block_on(f);

let balance = U256::from_limbs(balance?.0);
let nonce = nonce?.as_u64();
Expand All @@ -99,14 +93,14 @@ impl<M: Middleware> DatabaseRef for EthersDB<M> {
let add = eH160::from(address.0 .0);
let index = H256::from(index.to_be_bytes());
let slot_value: H256 =
Self::block_on(self.client.get_storage_at(add, index, self.block_number))?;
self.block_on(self.client.get_storage_at(add, index, self.block_number))?;
Ok(U256::from_be_bytes(slot_value.to_fixed_bytes()))
}

fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
let number = eU64::from(number);
let block: Option<Block<TxHash>> =
Self::block_on(self.client.get_block(BlockId::from(number)))?;
self.block_on(self.client.get_block(BlockId::from(number)))?;
// If number is given, the block is supposed to be finalized so unwrap is safe too.
Ok(B256::new(block.unwrap().hash.unwrap().0))
}
Expand Down

0 comments on commit b40e4e3

Please sign in to comment.