Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: store tokio::runtime::Handle in ethers/alloyDB #1557

Merged
merged 5 commits into from
Jun 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 22 additions & 33 deletions crates/revm/src/db/alloydb.rs
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
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
Loading