Use HandleOrRuntime to allow alloydb/ethersdb to hold a custom runtime

crates/revm/src/db.rs

@@ -1,5 +1,8 @@
 //! [Database] implementations.
 
+#[cfg(any(feature = "alloydb", feature = "ethersdb"))]
+mod utils;
+
 #[cfg(feature = "alloydb")]
 mod alloydb;
 pub mod emptydb;

crates/revm/src/db/alloydb.rs

@@ -6,50 +6,80 @@ use alloy_eips::BlockId;
 use alloy_provider::{Network, Provider};
 use alloy_transport::{Transport, TransportError};
 use std::future::IntoFuture;
-use tokio::runtime::Handle;
+use tokio::runtime::{Handle, Runtime};
+
+use super::utils::HandleOrRuntime;
 
 /// An alloy-powered REVM [Database].
 ///
 /// When accessing the database, it'll use the given provider to fetch the corresponding account's data.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct AlloyDB<T: Transport + Clone, N: Network, P: Provider<T, N>> {
     /// The provider to fetch the data from.
     provider: P,
     /// The block number on which the queries will be based on.
     block_number: BlockId,
     /// handle to the tokio runtime
-    handle: Handle,
+    rt: HandleOrRuntime,
     _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).
+    /// Create a new AlloyDB instance, with a [Provider] and a block.
     ///
     /// 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() {
+        let rt = match Handle::try_current() {
             Ok(handle) => match handle.runtime_flavor() {
                 tokio::runtime::RuntimeFlavor::CurrentThread => return None,
-                _ => handle,
+                _ => HandleOrRuntime::Handle(handle),
             },
             Err(_) => return None,
         };
         Some(Self {
             provider,
             block_number,
-            handle,
+            rt,
             _marker: std::marker::PhantomData,
         })
     }
 
+    // Create a new AlloyDB instance, with a provider and a block and a runtime.
+    //
+    // Refer to [tokio::runtime::Builder] on how to create a runtime if you are in synchronous world.
+    // If you are already using something like [tokio::main], call AlloyDB::new instead.
+    pub fn with_runtime(provider: P, block_number: BlockId, runtime: Runtime) -> Self {
+        let rt = HandleOrRuntime::Runtime(runtime);
+        Self {
+            provider,
+            block_number,
+            rt,
+            _marker: std::marker::PhantomData,
+        }
+    }
+
+    // Create a new AlloyDB instance, with a provider and a block and a runtime handle.
+    //
+    // Refer to [tokio::runtime::Builder] how to create a runtime if you are in synchronous world.
+    // If you are already using something like [tokio::main], call AlloyDB::new instead.
+    pub fn with_handle(provider: P, block_number: BlockId, handle: Handle) -> Self {
+        let rt = HandleOrRuntime::Handle(handle);
+        Self {
+            provider,
+            block_number,
+            rt,
+            _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>(&self, f: F) -> F::Output
     where
         F: std::future::Future + Send,
         F::Output: Send,
     {
-        tokio::task::block_in_place(move || self.handle.block_on(f))
+        self.rt.block_on(f)
     }
 
     /// Set the block number on which the queries will be based on.

crates/revm/src/db/ethersdb.rs

@@ -2,27 +2,29 @@ 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::Handle;
+use tokio::runtime::{Handle, Runtime};
 
 use crate::primitives::{AccountInfo, Address, Bytecode, B256, U256};
 use crate::{Database, DatabaseRef};
 
-#[derive(Debug, Clone)]
+use super::utils::HandleOrRuntime;
+
+#[derive(Debug)]
 pub struct EthersDB<M: Middleware> {
     client: Arc<M>,
     block_number: Option<BlockId>,
-    handle: Handle,
+    rt: HandleOrRuntime,
 }
 
 impl<M: Middleware> EthersDB<M> {
     /// 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 handle = match Handle::try_current() {
+        let rt = match Handle::try_current() {
             Ok(handle) => match handle.runtime_flavor() {
                 tokio::runtime::RuntimeFlavor::CurrentThread => return None,
-                _ => handle,
+                _ => HandleOrRuntime::Handle(handle),
             },
             Err(_) => return None,
         };
@@ -31,29 +33,73 @@ impl<M: Middleware> EthersDB<M> {
             Some(Self {
                 client,
                 block_number,
-                handle,
+                rt,
             })
         } else {
             let mut instance = Self {
-                client: client.clone(),
+                client,
                 block_number: None,
-                handle,
+                rt,
             };
             instance.block_number = Some(BlockId::from(
-                instance.block_on(client.get_block_number()).ok()?,
+                instance.block_on(instance.client.get_block_number()).ok()?,
             ));
             Some(instance)
         }
     }
 
+    // Create a new EthersDB instance, with a provider and a block (None for latest) and a runtime.
+    //
+    // Refer to [tokio::runtime::Builder] how to create a runtime if you are in synchronous world.
+    // If you are already using something like [tokio::main], call EthersDB::new instead.
+    pub fn with_runtime(
+        client: Arc<M>,
+        block_number: Option<BlockId>,
+        runtime: Runtime,
+    ) -> Option<Self> {
+        let rt = HandleOrRuntime::Runtime(runtime);
+        let mut instance = Self {
+            client,
+            block_number,
+            rt,
+        };
+
+        instance.block_number = Some(BlockId::from(
+            instance.block_on(instance.client.get_block_number()).ok()?,
+        ));
+        Some(instance)
+    }
+
+    // Create a new EthersDB instance, with a provider and a block (None for latest) and a handle.
+    //
+    // Refer to [tokio::runtime::Builder] how to create a runtime if you are in synchronous world.
+    // If you are already using something like [tokio::main], call EthersDB::new instead.
+    pub fn with_handle(
+        client: Arc<M>,
+        block_number: Option<BlockId>,
+        handle: Handle,
+    ) -> Option<Self> {
+        let rt = HandleOrRuntime::Handle(handle);
+        let mut instance = Self {
+            client,
+            block_number,
+            rt,
+        };
+
+        instance.block_number = Some(BlockId::from(
+            instance.block_on(instance.client.get_block_number()).ok()?,
+        ));
+        Some(instance)
+    }
+
     /// Internal utility function to call tokio feature and wait for output
     #[inline]
     fn block_on<F>(&self, f: F) -> F::Output
     where
         F: core::future::Future + Send,
         F::Output: Send,
     {
-        tokio::task::block_in_place(move || self.handle.block_on(f))
+        self.rt.block_on(f)
     }
 
     /// set block number on which upcoming queries will be based

crates/revm/src/db/utils.rs

@@ -0,0 +1,22 @@
+use tokio::runtime::{Handle, Runtime};
+
+// Hold a tokio runtime handle or full runtime
+#[derive(Debug)]
+pub(crate) enum HandleOrRuntime {
+    Handle(Handle),
+    Runtime(Runtime),
+}
+
+impl HandleOrRuntime {
+    #[inline]
+    pub(crate) fn block_on<F>(&self, f: F) -> F::Output
+    where
+        F: std::future::Future + Send,
+        F::Output: Send,
+    {
+        match self {
+            Self::Handle(handle) => tokio::task::block_in_place(move || handle.block_on(f)),
+            Self::Runtime(rt) => rt.block_on(f),
+        }
+    }
+}