Cached Trie Updates For In-Memory Chains With >1 Block

[rkrasiuk]

Context

Trie updates are in-memory state trie diffs that represent trie changes that would occur after committing the block.

reth/crates/trie/src/updates.rs

Lines 45 to 49 in 0231492

	/// The aggregation of trie updates.
	#[derive(Debug, Default, Clone, PartialEq, Eq, Deref)]
	pub struct TrieUpdates {
	trie_operations: HashMap<TrieKey, TrieOp>,
	}

Cached trie updates were enabled for in-memory canonical chains for any length since alpha.14 (#5871). Recently, we discovered that extended trie updates for in-memory chains with >1 block sometimes resulted in an invalid state of the trie (#7559) and disabled it for chains with more than 1 block (#7753). This was in part surfaced by the degraded performance of Engine API message processing with the release of the beta version (issue to be linked).

Test Case

Consider the following test:

use reth_db::{cursor::DbCursorRW, tables, transaction::{DbTx, DbTxMut}};
use reth_primitives::{keccak256, trie::Nibbles, Address, StorageEntry, B256, U256};
use reth_provider::test_utils::create_test_provider_factory;
use reth_trie::{hashed_cursor::HashedPostStateCursorFactory, test_utils::storage_root_prehashed, updates::TrieUpdates, HashedPostState, HashedStorage, StorageRoot};
use std::{collections::BTreeMap, iter};

#[test]
fn trie_updates_across_multiple_iterations() {
    let address = Address::random();
    let hashed_address = keccak256(address);

    let factory = create_test_provider_factory();

    let mut hashed_storage = BTreeMap::default();
    let mut post_state = HashedPostState::default();

    // Block #1
    // Update specific storage slots
    let mut modified_storage = BTreeMap::default();

    // 0x0f..
    let modified_key_prefix = Nibbles::from_nibbles([0x0, 0xf].into_iter().chain(iter::repeat(0).take(62)).collect::<Vec<_>>());

    // 0x0faa0..
    let mut modified_entry1 = modified_key_prefix.clone();
    modified_entry1.set_at(2, 0xa);
    modified_entry1.set_at(3, 0xa);

    // 0x0faaa..
    let mut modified_entry2 = modified_key_prefix.clone();
    modified_entry2.set_at(2, 0xa);
    modified_entry2.set_at(3, 0xa);
    modified_entry2.set_at(4, 0xa);

    // 0x0fab0..
    let mut modified_entry3 = modified_key_prefix.clone();
    modified_entry3.set_at(2, 0xa);
    modified_entry3.set_at(3, 0xb);

    // 0x0fba0..
    let mut modified_entry4 = modified_key_prefix.clone();
    modified_entry4.set_at(2, 0xb);
    modified_entry4.set_at(3, 0xa);

    [modified_entry1, modified_entry2, modified_entry3.clone(), modified_entry4]
        .into_iter()
        .for_each(|key| {
            modified_storage.insert(B256::from_slice(&key.pack()), U256::from(1));
        });

    // Update main hashed storage.
    hashed_storage.extend(modified_storage.clone());
    post_state.extend(HashedPostState::default().with_storages([(
        hashed_address,
        HashedStorage::from_iter(false, modified_storage.clone()),
    )]));

    let (storage_root, block1_updates) = compute_storage_root(address, factory.provider().unwrap().tx_ref(), &post_state);
    assert_eq!(storage_root, storage_root_prehashed(hashed_storage.clone()));

    // Block #2
    // Set 0x0fab0.. hashed slot to 0
    modified_storage.insert(B256::from_slice(&modified_entry3.pack()), U256::ZERO);

    // Update main hashed storage.
    hashed_storage.remove(&B256::from_slice(&modified_entry3.pack()));
    post_state.extend(HashedPostState::default().with_storages([(
        hashed_address,
        HashedStorage::from_iter(false, modified_storage.clone()),
    )]));

    let (storage_root, block2_updates) = compute_storage_root(address, factory.provider().unwrap().tx_ref(), &post_state);
    assert_eq!(storage_root, storage_root_prehashed(hashed_storage.clone()));

    // Commit trie updates
    {
        let mut updates = block1_updates.clone();
        updates.extend(block2_updates);

        let provider_rw = factory.provider_rw().unwrap();
        let mut hashed_storage_cursor = provider_rw.tx_ref().cursor_dup_write::<tables::HashedStorages>().unwrap();
        for (hashed_slot, value) in &hashed_storage {
            hashed_storage_cursor.upsert(hashed_address, StorageEntry { key: *hashed_slot, value: *value }).unwrap();
        }
        updates.flush(provider_rw.tx_ref()).unwrap();
        provider_rw.commit().unwrap();
    }

    // Recompute storage root for block #3
    let storage_root = StorageRoot::from_tx(factory.provider().unwrap().tx_ref(), address).root().unwrap();
    assert_eq!(storage_root, storage_root_prehashed(hashed_storage.clone()));
}

fn compute_storage_root<TX: DbTx>(
    address: Address,
    tx: &TX,
    post_state: &HashedPostState,
) -> (B256, TrieUpdates) {
    let mut prefix_sets = post_state.construct_prefix_sets();
    let (root, _, updates) = StorageRoot::from_tx(tx, address)
        .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(tx, &post_state.clone().into_sorted()))
        .with_prefix_set(prefix_sets.storage_prefix_sets.remove(&keccak256(address)).unwrap())
        .root_with_updates()
        .unwrap();
    (root, updates)
}

In the test above account starts off with an empty storage.
In block 1 four account storage slots get modified to a non-zero value which results in the following intermediate nodes:
format: <trie node key>: <trie node value>

0x0f: BranchNodeCompact { state_mask: TrieMask(0000110000000000), tree_mask: TrieMask(0000010000000000), hash_mask: TrieMask(0000010000000000), hashes: [0x6e3e74ff5fc7ca11f62f34c59e9fd8d8a736a67ba7b96be41042b1285fcd1f81], .. }
0x0fa: BranchNodeCompact { state_mask: TrieMask(0000110000000000), tree_mask: TrieMask(0000000000000000), hash_mask: TrieMask(0000010000000000), hashes: [0xb34f450ed88ff0b44dedcaa3b45173b58d37be7ac83e30b6aa33791ee1b51198], .. }

In block 2 storage slot with hash 0x0fab0.. gets zeroed out resulting in no trie updates.

Upon extension, trie updates from block 1 will remain in the cache resulting in them being committed to the database.

Final test result - storage root mismatch:

assertion `left == right` failed
  left: 0x262cd61ad4821504938e6554e08db1188394f1767414aa75f907926c04053c30
  right: 0x301ce07aeae98a9a438bf74031e181d2fee7406aacf52e82e8d0940731752519

Root Cause

The root cause for this is that root computation does not consider updated nodes from previous computations and only produces trie diffs against the current trie state in the database.

Solutions

Solution 1: Consider only the latest trie updates

Considering that at the moment we rely only on database state trie, the latest (and only the latest) produced trie updates should theoretically be sufficient for updating the trie. This solution needs validation.

Solution 2: Overlay database trie nodes with in-memory ones

Solution 1 still does not allow us to reuse the results of previous state root computations to speed it up for higher in-memory blocks. Another solution would be to create a data structure that would act as an in-memory overlay of the current database trie state akin to how HashedPostState is an overlay for database hashed state

reth/crates/trie/src/state.rs

Lines 25 to 32 in 0231492

	/// Representation of in-memory hashed state.
	#[derive(PartialEq, Eq, Clone, Default, Debug)]
	pub struct HashedPostState {
	/// Mapping of hashed address to account info, `None` if destroyed.
	pub accounts: HashMap<B256, Option<Account>>,
	/// Mapping of hashed address to hashed storage.
	pub storages: HashMap<B256, HashedStorage>,
	}

and implement TrieCursorFactory for this new structure.

reth/crates/trie/src/trie_cursor/mod.rs

Lines 19 to 29 in 0231492

	/// Factory for creating trie cursors.
	pub trait TrieCursorFactory {
	/// Create an account trie cursor.
	fn account_trie_cursor(&self) -> Result<Box<dyn TrieCursor + '_>, DatabaseError>;
	/// Create a storage tries cursor.
	fn storage_tries_cursor(
	&self,
	hashed_address: B256,
	) -> Result<Box<dyn TrieCursor + '_>, DatabaseError>;
	}

[gakonst]

Fixed in #7753.

[onbjerg]

It's not, this is for the long term solution/re-enabling of what we disabled in #7753

[gakonst]

Oh, OK, right nvm.

[lakshya-sky]

I had a pass at this and I made changes to as best as I could understand the problem. I have added in-memory cursor which successfully passes the test @rkrasiuk provided.

here is the update trie cursor https://github.com/lakshya-sky/reth/blob/testing-trie-updates/crates/trie/src/trie_cursor/update.rs

and the test where I use it:

https://github.com/lakshya-sky/reth/blob/47b30889bfeb4121e2a6db8b40346e6133d3d4bb/crates/trie/src/tests.rs#L127-L144

I don't know where should I use this new code to test on a live node!

[lakshya-sky]

I've managed to sync full node Holesky with latest changes, which is running for a day now. But I would like to keep it running to test against reorgs.

[github-actions]

This issue is stale because it has been open for 21 days with no activity.

[rkrasiuk]

closed via transition to new engine