feat: only bridge newly created contracts

portal-bridge/src/bridge/state.rs

@@ -11,7 +11,6 @@ use ethportal_api::{
     ContentValue, Enr, OverlayContentKey, StateContentKey, StateContentValue,
     StateNetworkApiClient,
 };
-use revm::Database;
 use revm_primitives::{keccak256, Bytecode, SpecId, B256};
 use tokio::{
     sync::{OwnedSemaphorePermit, Semaphore},
@@ -104,7 +103,7 @@ impl StateBridge {
 
         // Enable contract storage changes caching required for gossiping the storage trie
         let state_config = StateConfig {
-            cache_contract_storage_changes: true,
+            cache_contract_changes: true,
             block_to_trace: BlockToTrace::None,
         };
         let mut trin_execution = TrinExecution::new(temp_directory.path(), state_config).await?;
@@ -115,7 +114,7 @@ impl StateBridge {
                 .process_range_of_blocks(start_block - 1, None)
                 .await?;
             // flush the database cache
-            trin_execution.database.storage_cache.lock().clear();
+            trin_execution.database.clear_contract_cache();
         }
 
         info!("Gossiping state data from block {start_block} to {end_block} (inclusively)");
@@ -194,17 +193,21 @@ impl StateBridge {
             // if the code_hash is empty then then don't try to gossip the contract bytecode
             if account.code_hash != keccak256([]) {
                 // gossip contract bytecode
-                let code = trin_execution.database.code_by_hash(account.code_hash)?;
-                self.gossip_contract_bytecode(
-                    address_hash,
-                    &account_proof,
-                    block_hash,
-                    account.code_hash,
-                    code,
-                    content_idx,
-                )
-                .await?;
-                content_idx += 1;
+                if let Some(code) = trin_execution
+                    .database
+                    .get_newly_created_contract_if_available(account.code_hash)
+                {
+                    self.gossip_contract_bytecode(
+                        address_hash,
+                        &account_proof,
+                        block_hash,
+                        account.code_hash,
+                        code,
+                        content_idx,
+                    )
+                    .await?;
+                    content_idx += 1;
+                }
             }
 
             // gossip contract storage
@@ -227,8 +230,8 @@ impl StateBridge {
         }
 
         // flush the database cache
-        // This is used for gossiping storage trie diffs
-        trin_execution.database.storage_cache.lock().clear();
+        // This is used for gossiping storage trie diffs and newly created contracts
+        trin_execution.database.clear_contract_cache();
 
         Ok(())
     }

trin-execution/src/config.rs

@@ -3,16 +3,17 @@ use crate::{cli::TrinExecutionConfig, types::block_to_trace::BlockToTrace};
 #[derive(Debug, Clone)]
 pub struct StateConfig {
     /// This flag when enabled will storage all the trie keys from block execution in a cache, this
-    /// is needed for gossiping the storage trie's changes.
-    pub cache_contract_storage_changes: bool,
+    /// is needed for gossiping the storage trie's changes. It is also needed for gossiping newly
+    /// created contracts.
+    pub cache_contract_changes: bool,
     pub block_to_trace: BlockToTrace,
 }
 
 #[allow(clippy::derivable_impls)]
 impl Default for StateConfig {
     fn default() -> Self {
         Self {
-            cache_contract_storage_changes: false,
+            cache_contract_changes: false,
             block_to_trace: BlockToTrace::None,
         }
     }
@@ -21,7 +22,7 @@ impl Default for StateConfig {
 impl From<TrinExecutionConfig> for StateConfig {
     fn from(cli_config: TrinExecutionConfig) -> Self {
         Self {
-            cache_contract_storage_changes: false,
+            cache_contract_changes: false,
             block_to_trace: cli_config.block_to_trace,
         }
     }

trin-execution/src/storage/evm_db.rs

@@ -39,8 +39,10 @@ fn start_processing_timer(name: &str) -> HistogramTimer {
 pub struct EvmDB {
     /// State config
     pub config: StateConfig,
-    /// Storage cache for the accounts used optionally for gossiping, keyed by address hash.
-    pub storage_cache: Arc<Mutex<FbHashMap<32, HashSet<B256>>>>,
+    /// Storage cache for the accounts required for gossiping stat diffs, keyed by address hash.
+    storage_cache: Arc<Mutex<FbHashMap<32, HashSet<B256>>>>,
+    /// Cache for newly created contracts required for gossiping stat diffs, keyed by code hash.
+    newly_created_contracts: Arc<Mutex<FbHashMap<32, Bytecode>>>,
     /// The underlying database.
     pub db: Arc<RocksDB>,
     /// To get proofs and to verify trie state.
@@ -68,9 +70,11 @@ impl EvmDB {
         ));
 
         let storage_cache = Arc::new(Mutex::new(FbHashMap::default()));
+        let newly_created_contracts = Arc::new(Mutex::new(FbHashMap::default()));
         Ok(Self {
             config,
             storage_cache,
+            newly_created_contracts,
             db,
             trie,
         })
@@ -102,6 +106,15 @@ impl EvmDB {
         trie_diff
     }
 
+    pub fn get_newly_created_contract_if_available(&self, code_hash: B256) -> Option<Bytecode> {
+        self.newly_created_contracts.lock().get(&code_hash).cloned()
+    }
+
+    pub fn clear_contract_cache(&self) {
+        self.storage_cache.lock().clear();
+        self.newly_created_contracts.lock().clear();
+    }
+
     fn commit_account(
         &mut self,
         address_hash: B256,
@@ -228,7 +241,7 @@ impl EvmDB {
             trie_diff,
         } = trie.root_hash_with_changed_nodes()?;
 
-        if self.config.cache_contract_storage_changes {
+        if self.config.cache_contract_changes {
             let mut storage_cache_guard = self.storage_cache.lock();
             let account_storage_cache = storage_cache_guard.entry(address_hash).or_default();
             for key in trie_diff.keys() {
@@ -295,6 +308,12 @@ impl EvmDB {
         // TODO: Delete contract code if no accounts point to it: https://github.com/ethereum/trin/issues/1428
         let timer = start_commit_timer("contract:committing_contracts_total");
         for (hash, bytecode) in plain_state.contracts {
+            // Cache contract code for gossiping if flag is set
+            if self.config.cache_contract_changes {
+                self.newly_created_contracts
+                    .lock()
+                    .insert(hash, bytecode.clone());
+            }
             let timer = start_commit_timer("committing_contract");
             self.db
                 .put(hash, bytecode.original_bytes().as_ref())

trin-execution/tests/content_generation.rs

@@ -8,7 +8,6 @@ use ethportal_api::{
     },
     ContentValue, OverlayContentKey, StateContentKey, StateContentValue,
 };
-use revm::DatabaseRef;
 use revm_primitives::keccak256;
 use tracing::info;
 use tracing_test::traced_test;
@@ -96,7 +95,7 @@ async fn test_we_can_generate_content_key_values_up_to_x() -> Result<()> {
     let mut trin_execution = TrinExecution::new(
         temp_directory.path(),
         StateConfig {
-            cache_contract_storage_changes: true,
+            cache_contract_changes: true,
             block_to_trace: BlockToTrace::None,
         },
     )
@@ -158,16 +157,18 @@ async fn test_we_can_generate_content_key_values_up_to_x() -> Result<()> {
 
             // check contract code content key/value
             if account.code_hash != keccak256([]) {
-                let code = trin_execution
+                if let Some(code) = trin_execution
                     .database
-                    .code_by_hash_ref(account.code_hash)?;
-
-                let content_key = create_contract_content_key(address_hash, account.code_hash)
-                    .expect("Content key should be present");
-                let content_value = create_contract_content_value(block_hash, &account_proof, code)
-                    .expect("Content key should be present");
-                block_stats.check_content(&content_key, &content_value);
-                stats.check_content(&content_key, &content_value);
+                    .get_newly_created_contract_if_available(account.code_hash)
+                {
+                    let content_key = create_contract_content_key(address_hash, account.code_hash)
+                        .expect("Content key should be present");
+                    let content_value =
+                        create_contract_content_value(block_hash, &account_proof, code)
+                            .expect("Content key should be present");
+                    block_stats.check_content(&content_key, &content_value);
+                    stats.check_content(&content_key, &content_value);
+                }
             }
 
             // check contract storage content key/value
@@ -188,7 +189,7 @@ async fn test_we_can_generate_content_key_values_up_to_x() -> Result<()> {
 
         // flush the database cache
         // This is used for gossiping storage trie diffs
-        trin_execution.database.storage_cache.lock().clear();
+        trin_execution.database.clear_contract_cache();
         info!("Block {block_number} finished: {block_stats:?}");
     }
     temp_directory.close()?;