Skip to content

Commit 1e3b017

Browse files
karen-sarkisyankaren-sarkisyan
committed
refactor get_account method in cache provider
1 parent a0d4b2c commit 1e3b017

File tree

2 files changed

+107
-83
lines changed

2 files changed

+107
-83
lines changed

‎core/src/execution/cache/provider.rs‎

Lines changed: 48 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use alloy::{
44
eips::BlockId,
55
network::{primitives::HeaderResponse, BlockResponse},
66
primitives::{Address, B256},
7-
rpc::types::{EIP1186AccountProofResponse, Filter, Log},
7+
rpc::types::{Filter, Log},
88
};
99
use async_trait::async_trait;
1010
use eyre::{eyre, Result};
@@ -56,95 +56,64 @@ where
5656
block_id: BlockId,
5757
) -> Result<Account> {
5858
let block_hash = match block_id {
59-
BlockId::Hash(hash) => {
60-
let hash: B256 = hash.into();
61-
hash
62-
}
63-
block_id => {
64-
let block = self
65-
.inner
66-
.get_block(block_id, false)
67-
.await?
68-
.ok_or(eyre!("block not found"))?;
69-
block.header().hash()
70-
}
59+
BlockId::Hash(hash) => hash.into(),
60+
_ => self
61+
.inner
62+
.get_block(block_id, false)
63+
.await?
64+
.ok_or_else(|| eyre!("block not found"))?
65+
.header()
66+
.hash(),
7167
};
7268

73-
let slots_left_to_fetch = match self.cache.get_account(address, slots, block_hash) {
74-
Some((cached, missing_slots)) => {
75-
if missing_slots.is_empty() {
76-
if with_code && cached.code.is_none() {
77-
// All account data is cached, but code is missing - fetch code directly
78-
let expected = cached.account.code_hash;
79-
let acc_with_code = self
69+
let cached = self.cache.get_account(address, slots, block_hash);
70+
71+
match cached {
72+
None => {
73+
if with_code {
74+
if let Some((cached_code_hash, cached_code)) =
75+
self.cache.get_code_optimistically(address)
76+
{
77+
let fetched = self
8078
.inner
81-
.get_account(address, &[], true, block_hash.into())
79+
.get_account(address, slots, false, block_hash.into())
8280
.await?;
83-
let code = acc_with_code
84-
.code
85-
.ok_or_else(|| eyre!("provider did not return code"))?;
86-
if acc_with_code.account.code_hash != expected {
87-
return Err(eyre!("code_hash changed while fetching code"));
81+
if fetched.account.code_hash == cached_code_hash {
82+
let mut account = fetched;
83+
account.code = Some(cached_code);
84+
self.cache.insert_account(address, &account, block_hash);
85+
return Ok(account);
8886
}
89-
self.cache.insert_code(expected, code.clone());
90-
return Ok(Account {
91-
code: Some(code),
92-
..cached
93-
});
9487
}
95-
return Ok(cached);
9688
}
97-
missing_slots
89+
let fetched = self
90+
.inner
91+
.get_account(address, slots, with_code, block_hash.into())
92+
.await?;
93+
self.cache.insert_account(address, &fetched, block_hash);
94+
Ok(fetched)
9895
}
99-
None => slots.to_vec(),
100-
};
101-
102-
// Always fetch account proof without full code because it might be already in cache
103-
let account = self
104-
.inner
105-
.get_account(address, &slots_left_to_fetch, false, block_hash.into())
106-
.await?;
107-
108-
let response = EIP1186AccountProofResponse {
109-
address,
110-
balance: account.account.balance,
111-
code_hash: account.account.code_hash,
112-
nonce: account.account.nonce,
113-
storage_hash: account.account.storage_root,
114-
account_proof: account.account_proof.clone(),
115-
storage_proof: account.storage_proof.clone(),
116-
};
117-
118-
self.cache.insert(response, None, block_hash);
119-
120-
// Retrieve from cache again to get the full merged result
121-
let (mut full_account, missing_after) = self
122-
.cache
123-
.get_account(address, slots, block_hash)
124-
.ok_or(eyre!("Cache inconsistency after write"))?;
125-
126-
if !missing_after.is_empty() {
127-
return Err(eyre!("Failed to fetch all slots"));
128-
}
96+
Some((mut account, missing_slots)) => {
97+
let need_code = with_code && account.code.is_none();
98+
99+
if !missing_slots.is_empty() || need_code {
100+
let fetched = self
101+
.inner
102+
.get_account(address, &missing_slots, need_code, block_hash.into())
103+
.await?;
104+
self.cache.insert_account(address, &fetched, block_hash);
105+
106+
account.account = fetched.account;
107+
account.account_proof = fetched.account_proof;
108+
account.storage_proof.extend(fetched.storage_proof);
109+
if let Some(code) = fetched.code {
110+
account.code = Some(code);
111+
}
112+
}
129113

130-
// If code is still missing, fetch via get_account with with_code=true
131-
if with_code && full_account.code.is_none() {
132-
let expected = full_account.account.code_hash;
133-
let acc_with_code = self
134-
.inner
135-
.get_account(address, &[], true, block_hash.into())
136-
.await?;
137-
let code = acc_with_code
138-
.code
139-
.ok_or_else(|| eyre!("provider did not return code"))?;
140-
if acc_with_code.account.code_hash != expected {
141-
return Err(eyre!("code_hash changed while fetching code"));
114+
Ok(account)
142115
}
143-
self.cache.insert_code(expected, code.clone());
144-
full_account.code = Some(code);
145116
}
146-
147-
Ok(full_account)
148117
}
149118
}
150119

‎core/src/execution/cache/state.rs‎

Lines changed: 59 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -105,10 +105,38 @@ impl Cache {
105105
code.get(&code_hash).cloned()
106106
}
107107

108-
/// Insert code into the cache by code hash. Hash MUST come from a verified account proof.
109-
pub fn insert_code(&self, code_hash: B256, code: Bytes) {
110-
let mut code_cache = self.code.write().unwrap_or_else(|e| e.into_inner());
111-
code_cache.insert(code_hash, code);
108+
/// Try to recover code for an address from any cached account entry.
109+
///
110+
/// This is intentionally "optimistic": we only use the code if a freshly
111+
/// fetched proof confirms the same code hash.
112+
pub fn get_code_optimistically(&self, address: Address) -> Option<(B256, Bytes)> {
113+
let code_hash = {
114+
let accounts = self.accounts.read().unwrap_or_else(|e| e.into_inner());
115+
accounts.iter().find_map(|((cached_address, _), account)| {
116+
(*cached_address == address).then_some(account.code_hash)
117+
})?
118+
};
119+
120+
let code = self.get_code(code_hash)?;
121+
Some((code_hash, code))
122+
}
123+
124+
/// Insert a verified [`Account`] into the cache, distributing its data to
125+
/// the account-proof, storage-proof, and code sub-caches.
126+
pub(crate) fn insert_account(&self, address: Address, account: &Account, block_hash: B256) {
127+
self.insert(
128+
EIP1186AccountProofResponse {
129+
address,
130+
balance: account.account.balance,
131+
code_hash: account.account.code_hash,
132+
nonce: account.account.nonce,
133+
storage_hash: account.account.storage_root,
134+
account_proof: account.account_proof.clone(),
135+
storage_proof: account.storage_proof.clone(),
136+
},
137+
account.code.clone(),
138+
block_hash,
139+
);
112140
}
113141

114142
/// Get an account proof response with requested storage slots.
@@ -350,4 +378,31 @@ mod tests {
350378
let (account, _) = cache.get_account(address, &[], block_hash).unwrap();
351379
assert_eq!(account.code, Some(code));
352380
}
381+
382+
#[test]
383+
fn test_get_code_optimistically() {
384+
let cache = Cache::new();
385+
386+
let address = address!("0000000000000000000000000000000000000001");
387+
let other_address = address!("0000000000000000000000000000000000000002");
388+
let block_hash1 = b256!("0000000000000000000000000000000000000000000000000000000000000001");
389+
let block_hash2 = b256!("0000000000000000000000000000000000000000000000000000000000000002");
390+
let storage_hash =
391+
b256!("0000000000000000000000000000000000000000000000000000000000000003");
392+
let code_hash = b256!("0000000000000000000000000000000000000000000000000000000000000004");
393+
let code = Bytes::from_static(&[0x60, 0x80, 0x60, 0x40]);
394+
395+
let response = mock_account_proof_response(address, storage_hash, code_hash, vec![]);
396+
cache.insert(response, Some(code.clone()), block_hash1);
397+
398+
// Same address at another block keeps the optimistic lookup valid.
399+
let response2 = mock_account_proof_response(address, storage_hash, code_hash, vec![]);
400+
cache.insert(response2, None, block_hash2);
401+
402+
assert_eq!(
403+
cache.get_code_optimistically(address),
404+
Some((code_hash, code))
405+
);
406+
assert_eq!(cache.get_code_optimistically(other_address), None);
407+
}
353408
}

0 commit comments

Comments
 (0)