feat: add Finish stage (paradigmxyz#1279)

Author: onbjerg

bin/reth/src/args/network_args.rs

@@ -2,7 +2,11 @@
 
 use crate::dirs::{KnownPeersPath, PlatformPath};
 use clap::Args;
-use reth_primitives::NodeRecord;
+use reth_discv4::bootnodes::mainnet_nodes;
+use reth_net_nat::NatResolver;
+use reth_network::NetworkConfigBuilder;
+use reth_primitives::{ChainSpec, NodeRecord};
+use reth_staged_sync::Config;
 use std::path::PathBuf;
 
 /// Parameters for configuring the network more granularity via CLI
@@ -22,7 +26,7 @@ pub struct NetworkArgs {
     #[arg(long)]
     pub trusted_only: bool,
 
-    /// Nodes to bootstrap network discovery.
+    /// Bootnodes to connect to initially.
     ///
     /// Will fall back to a network-specific default if not specified.
     #[arg(long, value_delimiter = ',')]
@@ -37,6 +41,23 @@ pub struct NetworkArgs {
     /// Do not persist peers. Cannot be used with --peers-file
     #[arg(long, verbatim_doc_comment, conflicts_with = "peers_file")]
     pub no_persist_peers: bool,
+
+    /// NAT resolution method.
+    #[arg(long, default_value = "any")]
+    pub nat: NatResolver,
+}
+
+impl NetworkArgs {
+    /// Build a [`NetworkConfigBuilder`] from a [`Config`] and a [`ChainSpec`], in addition to the
+    /// values in this option struct.
+    pub fn network_config(&self, config: &Config, chain_spec: ChainSpec) -> NetworkConfigBuilder {
+        let peers_file = (!self.no_persist_peers).then_some(&self.peers_file);
+        config
+            .network_config(self.nat, peers_file.map(|f| f.as_ref().to_path_buf()))
+            .boot_nodes(self.bootnodes.clone().unwrap_or_else(mainnet_nodes))
+            .chain_spec(chain_spec)
+            .set_discovery(self.disable_discovery)
+    }
 }
 
 // === impl NetworkArgs ===

bin/reth/src/chain/import.rs

@@ -13,6 +13,7 @@ use reth_downloaders::{
 };
 use reth_interfaces::{
     consensus::{Consensus, ForkchoiceState},
+    p2p::headers::client::NoopStatusUpdater,
     sync::SyncStateUpdater,
 };
 use reth_primitives::ChainSpec;
@@ -142,22 +143,23 @@ impl ImportCommand {
         let mut pipeline = Pipeline::builder()
             .with_sync_state_updater(file_client)
             .add_stages(
-                OnlineStages::new(consensus.clone(), header_downloader, body_downloader).set(
-                    TotalDifficultyStage {
-                        chain_spec: self.chain.clone(),
-                        commit_threshold: config.stages.total_difficulty.commit_threshold,
-                    },
-                ),
-            )
-            .add_stages(
-                OfflineStages::default()
-                    .set(SenderRecoveryStage {
-                        commit_threshold: config.stages.sender_recovery.commit_threshold,
-                    })
-                    .set(ExecutionStage {
-                        chain_spec: self.chain.clone(),
-                        commit_threshold: config.stages.execution.commit_threshold,
-                    }),
+                DefaultStages::new(
+                    consensus.clone(),
+                    header_downloader,
+                    body_downloader,
+                    NoopStatusUpdater::default(),
+                )
+                .set(TotalDifficultyStage {
+                    chain_spec: self.chain.clone(),
+                    commit_threshold: config.stages.total_difficulty.commit_threshold,
+                })
+                .set(SenderRecoveryStage {
+                    commit_threshold: config.stages.sender_recovery.commit_threshold,
+                })
+                .set(ExecutionStage {
+                    chain_spec: self.chain.clone(),
+                    commit_threshold: config.stages.execution.commit_threshold,
+                }),
             )
             .with_max_block(0)
             .build();

bin/reth/src/node/mod.rs

@@ -12,22 +12,29 @@ use eyre::Context;
 use fdlimit::raise_fd_limit;
 use futures::{pin_mut, stream::select as stream_select, Stream, StreamExt};
 use reth_consensus::beacon::BeaconConsensus;
-use reth_db::mdbx::{Env, WriteMap};
+use reth_db::{
+    database::Database,
+    mdbx::{Env, WriteMap},
+    tables,
+    transaction::DbTx,
+};
 use reth_downloaders::{
     bodies::bodies::BodiesDownloaderBuilder,
     headers::reverse_headers::ReverseHeadersDownloaderBuilder,
 };
 use reth_interfaces::{
     consensus::{Consensus, ForkchoiceState},
-    p2p::{bodies::downloader::BodyDownloader, headers::downloader::HeaderDownloader},
+    p2p::{
+        bodies::downloader::BodyDownloader,
+        headers::{client::StatusUpdater, downloader::HeaderDownloader},
+    },
     sync::SyncStateUpdater,
 };
-use reth_net_nat::NatResolver;
 use reth_network::{
     error::NetworkError, NetworkConfig, NetworkEvent, NetworkHandle, NetworkManager,
 };
 use reth_network_api::NetworkInfo;
-use reth_primitives::{BlockNumber, ChainSpec, H256};
+use reth_primitives::{BlockNumber, ChainSpec, Head, H256};
 use reth_provider::{BlockProvider, HeaderProvider, ShareableDatabase};
 use reth_rpc_builder::{RethRpcModule, RpcServerConfig, TransportRpcModuleConfig};
 use reth_staged_sync::{
@@ -40,7 +47,7 @@ use reth_staged_sync::{
 };
 use reth_stages::{
     prelude::*,
-    stages::{ExecutionStage, SenderRecoveryStage, TotalDifficultyStage},
+    stages::{ExecutionStage, SenderRecoveryStage, TotalDifficultyStage, FINISH},
 };
 use reth_tasks::TaskExecutor;
 use std::{net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
@@ -89,9 +96,6 @@ pub struct Command {
     #[clap(flatten)]
     network: NetworkArgs,
 
-    #[arg(long, default_value = "any")]
-    nat: NatResolver,
-
     /// Set the chain tip manually for testing purposes.
     ///
     /// NOTE: This is a temporary flag
@@ -135,8 +139,9 @@ impl Command {
         self.init_trusted_nodes(&mut config);
 
         info!(target: "reth::cli", "Connecting to P2P network");
-        let netconf = self.load_network_config(&config, Arc::clone(&db), ctx.task_executor.clone());
-        let network = self.start_network(netconf, &ctx.task_executor, ()).await?;
+        let network_config =
+            self.load_network_config(&config, Arc::clone(&db), ctx.task_executor.clone());
+        let network = self.start_network(network_config, &ctx.task_executor, ()).await?;
         info!(target: "reth::cli", peer_id = %network.peer_id(), local_addr = %network.local_addr(), "Connected to P2P network");
 
         // TODO: Use the resolved secret to spawn the Engine API server
@@ -277,22 +282,42 @@ impl Command {
         Ok(handle)
     }
 
+    fn fetch_head(&self, db: Arc<Env<WriteMap>>) -> Result<Head, reth_interfaces::db::Error> {
+        db.view(|tx| {
+            let head = FINISH.get_progress(tx)?.unwrap_or_default();
+            let header = tx
+                .get::<tables::Headers>(head)?
+                .expect("the header for the latest block is missing, database is corrupt");
+            let total_difficulty = tx.get::<tables::HeaderTD>(head)?.expect(
+                "the total difficulty for the latest block is missing, database is corrupt",
+            );
+            let hash = tx
+                .get::<tables::CanonicalHeaders>(head)?
+                .expect("the hash for the latest block is missing, database is corrupt");
+            Ok::<Head, reth_interfaces::db::Error>(Head {
+                number: head,
+                hash,
+                difficulty: header.difficulty,
+                total_difficulty: total_difficulty.into(),
+                timestamp: header.timestamp,
+            })
+        })?
+        .map_err(Into::into)
+    }
+
     fn load_network_config(
         &self,
         config: &Config,
         db: Arc<Env<WriteMap>>,
         executor: TaskExecutor,
     ) -> NetworkConfig<ShareableDatabase<Arc<Env<WriteMap>>>> {
-        let peers_file = self.network.persistent_peers_file();
-        config.network_config(
-            db,
-            self.chain.clone(),
-            self.network.disable_discovery,
-            self.network.bootnodes.clone(),
-            self.nat,
-            peers_file,
-            Some(executor),
-        )
+        let head = self.fetch_head(Arc::clone(&db)).expect("the head block is missing");
+
+        self.network
+            .network_config(config, self.chain.clone())
+            .executor(Some(executor))
+            .set_head(head)
+            .build(Arc::new(ShareableDatabase::new(db)))
     }
 
     async fn build_pipeline<H, B, U>(
@@ -306,7 +331,7 @@ impl Command {
     where
         H: HeaderDownloader + 'static,
         B: BodyDownloader + 'static,
-        U: SyncStateUpdater,
+        U: SyncStateUpdater + StatusUpdater + Clone + 'static,
     {
         let stage_conf = &config.stages;
 
@@ -318,17 +343,13 @@ impl Command {
         }
 
         let pipeline = builder
-            .with_sync_state_updater(updater)
+            .with_sync_state_updater(updater.clone())
             .add_stages(
-                OnlineStages::new(consensus.clone(), header_downloader, body_downloader).set(
-                    TotalDifficultyStage {
+                DefaultStages::new(consensus.clone(), header_downloader, body_downloader, updater)
+                    .set(TotalDifficultyStage {
                         chain_spec: self.chain.clone(),
                         commit_threshold: stage_conf.total_difficulty.commit_threshold,
-                    },
-                ),
-            )
-            .add_stages(
-                OfflineStages::default()
+                    })
                     .set(SenderRecoveryStage {
                         commit_threshold: stage_conf.sender_recovery.commit_threshold,
                     })

bin/reth/src/p2p/mod.rs

@@ -10,6 +10,7 @@ use reth_interfaces::p2p::{
 };
 use reth_network::FetchClient;
 use reth_primitives::{BlockHashOrNumber, ChainSpec, NodeRecord, SealedHeader};
+use reth_provider::ShareableDatabase;
 use reth_staged_sync::{
     utils::{chainspec::chain_spec_value_parser, hash_or_num_value_parser},
     Config,
@@ -98,15 +99,10 @@ impl Command {
         config.peers.connect_trusted_nodes_only = self.trusted_only;
 
         let network = config
-            .network_config(
-                noop_db,
-                self.chain.clone(),
-                self.disable_discovery,
-                None,
-                self.nat,
-                None,
-                None,
-            )
+            .network_config(self.nat, None)
+            .set_discovery(self.disable_discovery)
+            .chain_spec(self.chain.clone())
+            .build(Arc::new(ShareableDatabase::new(noop_db)))
             .start_network()
             .await?;
 

bin/reth/src/stage/mod.rs

@@ -9,9 +9,8 @@ use crate::{
 use clap::{Parser, ValueEnum};
 use reth_consensus::beacon::BeaconConsensus;
 use reth_downloaders::bodies::bodies::BodiesDownloaderBuilder;
-use reth_net_nat::NatResolver;
 use reth_primitives::ChainSpec;
-use reth_provider::Transaction;
+use reth_provider::{ShareableDatabase, Transaction};
 use reth_staged_sync::{
     utils::{chainspec::chain_spec_value_parser, init::init_db},
     Config,
@@ -85,9 +84,6 @@ pub struct Command {
 
     #[clap(flatten)]
     network: NetworkArgs,
-
-    #[arg(long, default_value = "any")]
-    nat: NatResolver,
 }
 
 #[derive(Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Ord, ValueEnum)]
@@ -137,16 +133,10 @@ impl Command {
                     });
                 }
 
-                let network = config
-                    .network_config(
-                        db.clone(),
-                        self.chain.clone(),
-                        self.network.disable_discovery,
-                        None,
-                        self.nat,
-                        None,
-                        None,
-                    )
+                let network = self
+                    .network
+                    .network_config(&config, self.chain.clone())
+                    .build(Arc::new(ShareableDatabase::new(db.clone())))
                     .start_network()
                     .await?;
                 let fetch_client = Arc::new(network.fetch_client().await?);

crates/interfaces/src/p2p/headers/client.rs

@@ -45,3 +45,12 @@ pub trait StatusUpdater: Send + Sync {
     /// Updates the status of the p2p node
     fn update_status(&self, head: Head);
 }
+
+/// A [StatusUpdater] implementation that does nothing.
+#[derive(Debug, Clone, Default)]
+#[non_exhaustive]
+pub struct NoopStatusUpdater;
+
+impl StatusUpdater for NoopStatusUpdater {
+    fn update_status(&self, _: Head) {}
+}

crates/interfaces/src/test_utils/headers.rs

@@ -298,12 +298,27 @@ impl TestConsensus {
     }
 }
 
-/// Nil status updater for testing
-#[derive(Debug, Clone, Default)]
-pub struct TestStatusUpdater;
+/// Status updater for testing.
+///
+/// [`TestStatusUpdater::new()`] creates a new [`TestStatusUpdater`] that is **not** shareable. This
+/// struct wraps the sender side of a [`tokio::sync::watch`] channel. The receiving side of the
+/// channel (which is shareable by cloning it) is also returned.
+#[derive(Debug)]
+pub struct TestStatusUpdater(tokio::sync::watch::Sender<Head>);
+
+impl TestStatusUpdater {
+    /// Create a new test status updater and a receiver to listen to status updates on.
+    pub fn new() -> (Self, tokio::sync::watch::Receiver<Head>) {
+        let (tx, rx) = tokio::sync::watch::channel(Head::default());
+
+        (Self(tx), rx)
+    }
+}
 
 impl StatusUpdater for TestStatusUpdater {
-    fn update_status(&self, _: Head) {}
+    fn update_status(&self, head: Head) {
+        self.0.send(head).expect("could not send status update");
+    }
 }
 
 #[async_trait::async_trait]

crates/net/network/src/config.rs

@@ -185,6 +185,16 @@ impl NetworkConfigBuilder {
         self
     }
 
+    /// Sets the highest synced block.
+    ///
+    /// This is used to construct the appropriate [`ForkFilter`] and [`Status`] message.
+    ///
+    /// If not set, this defaults to the genesis specified by the current chain specification.
+    pub fn set_head(mut self, head: Head) -> Self {
+        self.head = Some(head);
+        self
+    }
+
     /// Sets the `HelloMessage` to send when connecting to peers.
     ///
     /// ```
@@ -265,6 +275,7 @@ impl NetworkConfigBuilder {
     }
 
     /// Sets the discovery service off on true.
+    // TODO(onbjerg): This name does not imply `true` = disable
     pub fn set_discovery(mut self, disable_discovery: bool) -> Self {
         if disable_discovery {
             self.disable_discovery();
@@ -309,7 +320,7 @@ impl NetworkConfigBuilder {
         let head = head.unwrap_or(Head {
             hash: chain_spec.genesis_hash(),
             number: 0,
-            timestamp: 0,
+            timestamp: chain_spec.genesis.timestamp,
             difficulty: chain_spec.genesis.difficulty,
             total_difficulty: chain_spec.genesis.difficulty,
         });