diff --git a/src/backend.rs b/src/backend.rs index fb42353..c7db1f2 100644 --- a/src/backend.rs +++ b/src/backend.rs @@ -32,6 +32,7 @@ use std::{ path::Path, pin::Pin, sync::{ + atomic::{AtomicU8, Ordering}, mpsc::{channel as oneshot_channel, Sender as OneshotSender}, Arc, }, @@ -64,6 +65,11 @@ type AddressData = AddressHashMap; type StorageData = AddressHashMap; type BlockHashData = HashMap; +/// Constants for account fetching modes +const ACCOUNT_FETCH_UNCHECKED: u8 = 0; +const ACCOUNT_FETCH_SUPPORTS_ACC_INFO: u8 = 1; +const ACCOUNT_FETCH_SEPARATE_REQUESTS: u8 = 2; + struct AnyRequestFuture { sender: OneshotSender>, future: Pin> + Send>>, @@ -163,6 +169,8 @@ pub struct BackendHandler

{ /// The block to fetch data from. // This is an `Option` so that we can have less code churn in the functions below block_id: Option, + /// The mode for fetching account data + account_fetch_mode: Arc, } impl

BackendHandler

@@ -185,6 +193,7 @@ where queued_requests: Default::default(), incoming: rx, block_id, + account_fetch_mode: Arc::new(AtomicU8::new(ACCOUNT_FETCH_UNCHECKED)), } } @@ -281,16 +290,72 @@ where /// returns the future that fetches the account data fn get_account_req(&self, address: Address) -> ProviderRequest { trace!(target: "backendhandler", "preparing account request, address={:?}", address); + let provider = self.provider.clone(); let block_id = self.block_id.unwrap_or_default(); + let mode = Arc::clone(&self.account_fetch_mode); + let fut = Box::pin(async move { - let balance = provider.get_balance(address).block_id(block_id).into_future(); - let nonce = provider.get_transaction_count(address).block_id(block_id).into_future(); - let code = provider.get_code_at(address).block_id(block_id).into_future(); - let resp = tokio::try_join!(balance, nonce, code).map_err(Into::into); - (resp, address) + let initial_mode = mode.load(Ordering::Relaxed); + match initial_mode { + ACCOUNT_FETCH_UNCHECKED => { + let acc_info_fut = provider.get_account_info(address).block_id(block_id); + let balance_fut = + provider.get_balance(address).block_id(block_id).into_future(); + let nonce_fut = + provider.get_transaction_count(address).block_id(block_id).into_future(); + let code_fut = provider.get_code_at(address).block_id(block_id).into_future(); + + tokio::select! { + acc_info = acc_info_fut => { + match acc_info { + Ok(info) => { + mode.store(ACCOUNT_FETCH_SUPPORTS_ACC_INFO, Ordering::Relaxed); + Ok((info.balance, info.nonce, info.code)) + } + Err(_) => { + mode.store(ACCOUNT_FETCH_SEPARATE_REQUESTS, Ordering::Relaxed); + tokio::try_join!(balance_fut, nonce_fut, code_fut).map_err(Into::into) + } + } + } + triple = async { tokio::try_join!(balance_fut, nonce_fut, code_fut) } => { + match triple { + Ok((balance, nonce, code)) => { + mode.store(ACCOUNT_FETCH_SEPARATE_REQUESTS, Ordering::Relaxed); + Ok((balance, nonce, code)) + } + Err(err) => Err(err.into()) + } + } + } + } + + ACCOUNT_FETCH_SUPPORTS_ACC_INFO => provider + .get_account_info(address) + .block_id(block_id) + .await + .map(|info| (info.balance, info.nonce, info.code)) + .map_err(Into::into), + + ACCOUNT_FETCH_SEPARATE_REQUESTS => { + let (balance, nonce, code) = tokio::try_join!( + provider.get_balance(address).block_id(block_id), + provider.get_transaction_count(address).block_id(block_id), + provider.get_code_at(address).block_id(block_id) + ) + .map_err(Into::into)?; + Ok((balance, nonce, code)) + } + + _ => unreachable!("Invalid account fetch mode"), + } }); - ProviderRequest::Account(fut) + + ProviderRequest::Account(Box::pin(async move { + let result = fut.await; + (result, address) + })) } /// process a request for an account @@ -301,7 +366,8 @@ where } Entry::Vacant(entry) => { entry.insert(vec![listener]); - self.pending_requests.push(self.get_account_req(address)); + let account_req = self.get_account_req(address); + self.pending_requests.push(account_req); } } }