feat: update chain_head (#1836)

morph-dev · web-flow · commit 44a974d19a45 · 2025-05-19T23:36:51.000+03:00
diff --git a/bin/trin/src/run.rs b/bin/trin/src/run.rs
@@ -19,7 +19,7 @@ use trin_state::initialize_state_network;
 use trin_storage::{config::StorageCapacityConfig, PortalStorageConfigFactory};
 #[cfg(windows)]
 use trin_utils::cli::Web3TransportType;
-use trin_validation::oracle::HeaderOracle;
+use trin_validation::{chain_head::ChainHead, oracle::HeaderOracle};
 use utp_rs::socket::UtpSocket;
 
 use crate::{
@@ -121,6 +121,9 @@ async fn run_trin_internal(
     // Initialize validation oracle
     let header_oracle = Arc::new(RwLock::new(HeaderOracle::default()));
 
+    // Initialize head of the chain management.
+    let chain_head = ChainHead::new_pectra_defaults();
+
     // Initialize and spawn uTP socket
     let (utp_talk_reqs_tx, utp_talk_reqs_rx) = mpsc::unbounded_channel();
 
@@ -170,6 +173,7 @@ async fn run_trin_internal(
             portalnet_config.clone(),
             storage_config_factory.create(&Subnetwork::Beacon, Distance::MAX)?,
             header_oracle.clone(),
+            chain_head,
         )
         .await?
     } else {
diff --git a/crates/ethportal-api/src/types/consensus/light_client/update.rs b/crates/ethportal-api/src/types/consensus/light_client/update.rs
@@ -10,6 +10,7 @@ use ssz_types::{
 use superstruct::superstruct;
 use tree_hash_derive::TreeHash;
 
+use super::header::LightClientHeader;
 use crate::{
     light_client::header::{LightClientHeaderDeneb, LightClientHeaderElectra},
     types::consensus::{
@@ -108,4 +109,24 @@ impl LightClientUpdate {
             ForkName::Electra => LightClientUpdateElectra::from_ssz_bytes(bytes).map(Self::Electra),
         }
     }
+
+    pub fn attested_header(&self) -> LightClientHeader {
+        match self {
+            Self::Bellatrix(update) => LightClientHeader::Bellatrix(update.attested_header.clone()),
+            Self::Capella(update) => LightClientHeader::Capella(update.attested_header.clone()),
+            Self::Deneb(update) => LightClientHeader::Deneb(update.attested_header.clone()),
+            Self::Electra(update) => LightClientHeader::Electra(update.attested_header.clone()),
+        }
+    }
+
+    pub fn finalized_header(&self) -> LightClientHeader {
+        match self {
+            Self::Bellatrix(update) => {
+                LightClientHeader::Bellatrix(update.finalized_header.clone())
+            }
+            Self::Capella(update) => LightClientHeader::Capella(update.finalized_header.clone()),
+            Self::Deneb(update) => LightClientHeader::Deneb(update.finalized_header.clone()),
+            Self::Electra(update) => LightClientHeader::Electra(update.finalized_header.clone()),
+        }
+    }
 }
diff --git a/crates/light-client/src/client.rs b/crates/light-client/src/client.rs
@@ -19,6 +19,7 @@ use crate::{
     errors::NodeError,
     node::Node,
     rpc::Rpc,
+    watch::LightClientWatchReceivers,
 };
 
 #[derive(Default)]
@@ -452,4 +453,8 @@ impl<DB: Database, R: ConsensusRpc + 'static> Client<DB, R> {
     pub async fn get_light_client_store(&self) -> Result<LightClientStore> {
         self.node.read().await.get_light_client_store()
     }
+
+    pub async fn get_light_client_watch_receivers(&self) -> LightClientWatchReceivers {
+        self.node.read().await.consensus.get_watch_receivers()
+    }
 }
diff --git a/crates/light-client/src/consensus/consensus_client.rs b/crates/light-client/src/consensus/consensus_client.rs
@@ -14,7 +14,7 @@ use ethportal_api::{
         finality_update::{LightClientFinalityUpdate, LightClientFinalityUpdateElectra},
         optimistic_update::{LightClientOptimisticUpdate, LightClientOptimisticUpdateElectra},
         store::LightClientStore,
-        update::{FinalizedRootProofLenElectra, LightClientUpdateElectra},
+        update::{FinalizedRootProofLenElectra, LightClientUpdate, LightClientUpdateElectra},
     },
 };
 use milagro_bls::PublicKey;
@@ -28,6 +28,7 @@ use crate::{
     consensus::{
         constants::MAX_REQUEST_LIGHT_CLIENT_UPDATES, rpc::portal_rpc::expected_current_slot,
     },
+    watch::{light_client_watch_channels, LightClientWatchReceivers, LightClientWatchSenders},
 };
 
 // https://github.com/ethereum/consensus-specs/blob/dev/specs/altair/light-client/sync-protocol.md
@@ -40,6 +41,7 @@ pub struct ConsensusLightClient<R: ConsensusRpc> {
     initial_checkpoint: B256,
     pub last_checkpoint: Option<B256>,
     pub config: Arc<Config>,
+    watch_senders: LightClientWatchSenders,
 }
 
 impl<R: ConsensusRpc> ConsensusLightClient<R> {
@@ -49,23 +51,27 @@ impl<R: ConsensusRpc> ConsensusLightClient<R> {
         config: Arc<Config>,
     ) -> Result<ConsensusLightClient<R>> {
         let rpc = R::new(rpc);
+        let (watch_senders, _watch_receivers) = light_client_watch_channels();
 
         Ok(ConsensusLightClient {
             rpc,
             store: LightClientStore::default(),
             last_checkpoint: None,
             config,
             initial_checkpoint: checkpoint_block_root,
+            watch_senders,
         })
     }
 
     pub fn with_custom_rpc(rpc: R, checkpoint_block_root: B256, config: Arc<Config>) -> Self {
+        let (watch_senders, _watch_receivers) = light_client_watch_channels();
         ConsensusLightClient {
             rpc,
             store: LightClientStore::default(),
             last_checkpoint: None,
             config,
             initial_checkpoint: checkpoint_block_root,
+            watch_senders,
         }
     }
 
@@ -105,6 +111,10 @@ impl<R: ConsensusRpc> ConsensusLightClient<R> {
         &self.store
     }
 
+    pub fn get_watch_receivers(&self) -> LightClientWatchReceivers {
+        self.watch_senders.subscribe()
+    }
+
     pub async fn sync(&mut self) -> Result<()> {
         self.bootstrap().await?;
 
@@ -366,13 +376,35 @@ impl<R: ConsensusRpc> ConsensusLightClient<R> {
     }
 
     fn apply_update(&mut self, update: &LightClientUpdateElectra) {
-        let update = GenericUpdate::from(update);
-        self.apply_generic_update(&update);
+        let optimistic_header_slot = self.store.optimistic_header.slot;
+        let finalized_header_slot = self.store.finalized_header.slot;
+
+        let generic_update = GenericUpdate::from(update);
+        self.apply_generic_update(&generic_update);
+
+        if optimistic_header_slot < self.store.optimistic_header.slot
+            || finalized_header_slot < self.store.finalized_header.slot
+        {
+            self.watch_senders
+                .update
+                .send_replace(Some(LightClientUpdate::Electra(update.clone())));
+        }
     }
 
     fn apply_finality_update(&mut self, update: &LightClientFinalityUpdateElectra) {
-        let update = GenericUpdate::from(update);
-        self.apply_generic_update(&update);
+        let optimistic_header_slot = self.store.optimistic_header.slot;
+        let finalized_header_slot = self.store.finalized_header.slot;
+
+        let generic_update = GenericUpdate::from(update);
+        self.apply_generic_update(&generic_update);
+
+        if optimistic_header_slot < self.store.optimistic_header.slot
+            || finalized_header_slot < self.store.finalized_header.slot
+        {
+            self.watch_senders
+                .finality_update
+                .send_replace(Some(LightClientFinalityUpdate::Electra(update.clone())));
+        }
     }
 
     fn log_finality_update(&self, update: &GenericUpdate) {
@@ -393,8 +425,16 @@ impl<R: ConsensusRpc> ConsensusLightClient<R> {
     }
 
     fn apply_optimistic_update(&mut self, update: &LightClientOptimisticUpdateElectra) {
-        let update = GenericUpdate::from(update);
-        self.apply_generic_update(&update);
+        let optimistic_header_slot = self.store.optimistic_header.slot;
+
+        let generic_update = GenericUpdate::from(update);
+        self.apply_generic_update(&generic_update);
+
+        if optimistic_header_slot < self.store.optimistic_header.slot {
+            self.watch_senders
+                .optimistic_update
+                .send_replace(Some(LightClientOptimisticUpdate::Electra(update.clone())));
+        }
     }
 
     fn log_optimistic_update(&self, update: &GenericUpdate) {
diff --git a/crates/light-client/src/lib.rs b/crates/light-client/src/lib.rs
@@ -10,3 +10,4 @@ pub mod errors;
 pub mod node;
 pub mod rpc;
 pub mod utils;
+pub mod watch;
diff --git a/crates/light-client/src/watch.rs b/crates/light-client/src/watch.rs
@@ -0,0 +1,50 @@
+use ethportal_api::light_client::{
+    finality_update::LightClientFinalityUpdate, optimistic_update::LightClientOptimisticUpdate,
+    update::LightClientUpdate,
+};
+use tokio::sync::watch;
+
+/// The [tokio::sync::watch::channel] Receivers for light client update types.
+///
+/// Value will be `None` until first update.
+#[derive(Clone)]
+pub struct LightClientWatchReceivers {
+    pub update: watch::Receiver<Option<LightClientUpdate>>,
+    pub optimistic_update: watch::Receiver<Option<LightClientOptimisticUpdate>>,
+    pub finality_update: watch::Receiver<Option<LightClientFinalityUpdate>>,
+}
+
+/// The [tokio::sync::watch::channel] Senders for light client update types.
+#[derive(Debug, Clone)]
+pub struct LightClientWatchSenders {
+    pub update: watch::Sender<Option<LightClientUpdate>>,
+    pub optimistic_update: watch::Sender<Option<LightClientOptimisticUpdate>>,
+    pub finality_update: watch::Sender<Option<LightClientFinalityUpdate>>,
+}
+
+impl LightClientWatchSenders {
+    pub fn subscribe(&self) -> LightClientWatchReceivers {
+        LightClientWatchReceivers {
+            update: self.update.subscribe(),
+            optimistic_update: self.optimistic_update.subscribe(),
+            finality_update: self.finality_update.subscribe(),
+        }
+    }
+}
+
+pub fn light_client_watch_channels() -> (LightClientWatchSenders, LightClientWatchReceivers) {
+    let (update_sender, update_receiver) = watch::channel(None);
+    let (optimistic_update_sender, optimistic_update_receiver) = watch::channel(None);
+    let (finality_update_sender, finality_update_receiver) = watch::channel(None);
+    let senders = LightClientWatchSenders {
+        update: update_sender,
+        optimistic_update: optimistic_update_sender,
+        finality_update: finality_update_sender,
+    };
+    let receivers = LightClientWatchReceivers {
+        update: update_receiver,
+        optimistic_update: optimistic_update_receiver,
+        finality_update: finality_update_receiver,
+    };
+    (senders, receivers)
+}
diff --git a/crates/subnetworks/beacon/src/lib.rs b/crates/subnetworks/beacon/src/lib.rs
@@ -26,7 +26,7 @@ use tokio::{
 };
 use tracing::info;
 use trin_storage::PortalStorageConfig;
-use trin_validation::oracle::HeaderOracle;
+use trin_validation::{chain_head::ChainHead, oracle::HeaderOracle};
 use utp_rs::socket::UtpSocket;
 
 use crate::{events::BeaconEvents, jsonrpc::BeaconRequestHandler, network::BeaconNetwork};
@@ -46,6 +46,7 @@ pub async fn initialize_beacon_network(
     portalnet_config: PortalnetConfig,
     storage_config: PortalStorageConfig,
     header_oracle: Arc<RwLock<HeaderOracle>>,
+    chain_head: ChainHead,
 ) -> anyhow::Result<(
     BeaconHandler,
     BeaconNetworkTask,
@@ -62,6 +63,7 @@ pub async fn initialize_beacon_network(
         storage_config,
         portalnet_config.clone(),
         header_oracle,
+        chain_head,
     )
     .await?;
     let beacon_event_stream = beacon_network.overlay.event_stream().await?;
diff --git a/crates/subnetworks/beacon/src/network.rs b/crates/subnetworks/beacon/src/network.rs
@@ -16,7 +16,7 @@ use portalnet::{
 use tokio::sync::{Mutex, RwLock};
 use tracing::{error, info};
 use trin_storage::PortalStorageConfig;
-use trin_validation::oracle::HeaderOracle;
+use trin_validation::{chain_head::ChainHead, oracle::HeaderOracle};
 use utp_rs::socket::UtpSocket;
 
 use crate::{
@@ -51,14 +51,18 @@ impl BeaconNetwork {
         storage_config: PortalStorageConfig,
         portal_config: PortalnetConfig,
         header_oracle: Arc<RwLock<HeaderOracle>>,
+        chain_head: ChainHead,
     ) -> anyhow::Result<Self> {
         let config = OverlayConfig {
             bootnode_enrs: portal_config.bootnodes.clone(),
             utp_transfer_limit: portal_config.utp_transfer_limit,
             gossip_dropped: GOSSIP_DROPPED,
             ..Default::default()
         };
-        let storage = Arc::new(PLMutex::new(BeaconStorage::new(storage_config)?));
+        let storage = Arc::new(PLMutex::new(BeaconStorage::new(
+            storage_config,
+            chain_head.clone(),
+        )?));
         storage.lock().spawn_pruning_task(); // Spawn pruning task to clean up expired content.
         let storage_clone = Arc::clone(&storage);
         let validator = Arc::new(BeaconValidator::new(header_oracle));
@@ -81,17 +85,49 @@ impl BeaconNetwork {
         // Get the trusted block root to start syncing from.
         let trusted_block_root = get_trusted_block_root(&portal_config, storage_clone)?;
 
-        // Spawn the beacon sync task.
+        // Spawn light client sync and update watch task
         tokio::spawn(async move {
-            let beacon_sync = BeaconSync::new(overlay_tx);
-            let beacon_sync = beacon_sync.start(trusted_block_root).await;
-            match beacon_sync {
-                Ok(client) => {
-                    let mut beacon_client = beacon_client_clone.lock().await;
-                    *beacon_client = Some(client);
-                }
+            // Sync LightClient
+            let client = match BeaconSync::new(overlay_tx).start(trusted_block_root).await {
+                Ok(client) => client,
                 Err(err) => {
                     error!(error = %err, "Failed to start beacon sync.");
+                    return;
+                }
+            };
+            let mut watch_receivers = client.get_light_client_watch_receivers().await;
+            *beacon_client_clone.lock().await = Some(client);
+
+            // Watch for light clients updates and update ChainHead
+            loop {
+                tokio::select! {
+                    Ok(()) = watch_receivers.update.changed() => {
+                        let update = watch_receivers
+                            .update
+                            .borrow_and_update()
+                            .as_ref()
+                            .expect("Updated LightClientUpdate must be present")
+                            .clone();
+                        chain_head.process_update(update);
+                    }
+                    Ok(()) = watch_receivers.optimistic_update.changed() => {
+                        let update = watch_receivers
+                            .optimistic_update
+                            .borrow_and_update()
+                            .as_ref()
+                            .expect("Updated LightClientOptimisticUpdate must be present")
+                            .clone();
+                        chain_head.process_optimistic_update(update);
+                    }
+                    Ok(()) = watch_receivers.finality_update.changed() => {
+                        let update = watch_receivers
+                            .finality_update
+                            .borrow_and_update()
+                            .as_ref()
+                            .expect("Updated LightClientFinalityUpdate must be present")
+                            .clone();
+                        chain_head.process_finality_update(update);
+                    }
                 }
             }
         });
@@ -176,12 +212,14 @@ mod tests {
             temp_dir.path().to_path_buf(),
         )
         .unwrap();
-
         let storage_config = storage_cfg_factory
             .create(&Subnetwork::Beacon, Distance::MAX)
             .unwrap();
+
         // A mock storage that always returns None
-        let storage_clone = Arc::new(PLMutex::new(BeaconStorage::new(storage_config).unwrap()));
+        let storage_clone = Arc::new(PLMutex::new(
+            BeaconStorage::new(storage_config, ChainHead::new_pectra_defaults()).unwrap(),
+        ));
 
         let result = get_trusted_block_root(&portal_config, storage_clone)
             .expect("Function should not fail with an Err");
diff --git a/crates/subnetworks/beacon/src/storage.rs b/crates/subnetworks/beacon/src/storage.rs
diff --git a/crates/validation/src/chain_head/mod.rs b/crates/validation/src/chain_head/mod.rs
diff --git a/crates/validation/src/chain_head/store.rs b/crates/validation/src/chain_head/store.rs
