feat(engine, snapshot): snapshot hook

Cargo.toml

@@ -32,6 +32,7 @@ members = [
     "crates/rpc/rpc-engine-api",
     "crates/rpc/rpc-types",
     "crates/rpc/rpc-testing-util",
+    "crates/snapshot",
     "crates/stages",
     "crates/storage/codecs",
     "crates/storage/db",

crates/consensus/beacon/Cargo.toml

@@ -19,6 +19,7 @@ reth-rpc-types.workspace = true
 reth-tasks.workspace = true
 reth-payload-builder.workspace = true
 reth-prune = { path = "../../prune" }
+reth-snapshot = { path = "../../snapshot" }
 reth-rpc-types-compat.workspace = true
 # async
 tokio = { workspace = true, features = ["sync"] }

crates/consensus/beacon/src/engine/hooks/controller.rs

@@ -54,7 +54,7 @@ impl EngineHooksController {
     ) -> Poll<Result<PolledHook, EngineHookError>> {
         let Some(mut hook) = self.running_hook_with_db_write.take() else { return Poll::Pending };
 
-        match hook.poll(cx, args) {
+        match hook.poll(cx, args)? {
             Poll::Ready((event, action)) => {
                 let result = PolledHook { event, action, db_access_level: hook.db_access_level() };
 
@@ -109,7 +109,7 @@ impl EngineHooksController {
             return Poll::Pending
         }
 
-        if let Poll::Ready((event, action)) = hook.poll(cx, args) {
+        if let Poll::Ready((event, action)) = hook.poll(cx, args)? {
             let result = PolledHook { event, action, db_access_level: hook.db_access_level() };
 
             debug!(

crates/consensus/beacon/src/engine/hooks/mod.rs

@@ -1,4 +1,4 @@
-use reth_interfaces::RethError;
+use reth_interfaces::{RethError, RethResult};
 use reth_primitives::BlockNumber;
 use std::{
     fmt,
@@ -11,6 +11,9 @@ pub(crate) use controller::{EngineHooksController, PolledHook};
 mod prune;
 pub use prune::PruneHook;
 
+mod snapshot;
+pub use snapshot::SnapshotHook;
+
 /// Collection of [engine hooks][`EngineHook`].
 #[derive(Default)]
 pub struct EngineHooks {
@@ -47,7 +50,7 @@ pub trait EngineHook: Send + Sync + 'static {
         &mut self,
         cx: &mut Context<'_>,
         ctx: EngineContext,
-    ) -> Poll<(EngineHookEvent, Option<EngineHookAction>)>;
+    ) -> Poll<RethResult<(EngineHookEvent, Option<EngineHookAction>)>>;
 
     /// Returns [db access level][`EngineHookDBAccessLevel`] the hook needs.
     fn db_access_level(&self) -> EngineHookDBAccessLevel;
@@ -58,6 +61,8 @@ pub trait EngineHook: Send + Sync + 'static {
 pub struct EngineContext {
     /// Tip block number.
     pub tip_block_number: BlockNumber,
+    /// Finalized block number, if known.
+    pub finalized_block_number: Option<BlockNumber>,
 }
 
 /// An event emitted when [hook][`EngineHook`] is polled.

crates/consensus/beacon/src/engine/hooks/prune.rs

@@ -9,7 +9,7 @@ use crate::{
 use futures::FutureExt;
 use metrics::Counter;
 use reth_db::database::Database;
-use reth_interfaces::RethError;
+use reth_interfaces::{RethError, RethResult};
 use reth_primitives::BlockNumber;
 use reth_prune::{Pruner, PrunerError, PrunerWithResult};
 use reth_tasks::TaskSpawner;
@@ -55,7 +55,7 @@ impl<DB: Database + 'static> PruneHook<DB> {
     fn poll_pruner(
         &mut self,
         cx: &mut Context<'_>,
-    ) -> Poll<(EngineHookEvent, Option<EngineHookAction>)> {
+    ) -> Poll<RethResult<(EngineHookEvent, Option<EngineHookAction>)>> {
         let result = match self.pruner_state {
             PrunerState::Idle(_) => return Poll::Pending,
             PrunerState::Running(ref mut fut) => {
@@ -69,14 +69,7 @@ impl<DB: Database + 'static> PruneHook<DB> {
 
                 match result {
                     Ok(_) => EngineHookEvent::Finished(Ok(())),
-                    Err(err) => EngineHookEvent::Finished(Err(match err {
-                        PrunerError::PrunePart(_) | PrunerError::InconsistentData(_) => {
-                            EngineHookError::Internal(Box::new(err))
-                        }
-                        PrunerError::Interface(err) => err.into(),
-                        PrunerError::Database(err) => RethError::Database(err).into(),
-                        PrunerError::Provider(err) => RethError::Provider(err).into(),
-                    })),
+                    Err(err) => EngineHookEvent::Finished(Err(err.into())),
                 }
             }
             Err(_) => {
@@ -85,14 +78,15 @@ impl<DB: Database + 'static> PruneHook<DB> {
             }
         };
 
-        Poll::Ready((event, None))
+        Poll::Ready(Ok((event, None)))
     }
 
     /// This will try to spawn the pruner if it is idle:
     /// 1. Check if pruning is needed through [Pruner::is_pruning_needed].
-    /// 2a. If pruning is needed, pass tip block number to the [Pruner::run] and spawn it in a
+    /// 2.
+    ///     1. If pruning is needed, pass tip block number to the [Pruner::run] and spawn it in a
     /// separate task. Set pruner state to [PrunerState::Running].
-    /// 2b. If pruning is not needed, set pruner state back to [PrunerState::Idle].
+    ///     2. If pruning is not needed, set pruner state back to [PrunerState::Idle].
     ///
     /// If pruner is already running, do nothing.
     fn try_spawn_pruner(
@@ -136,11 +130,11 @@ impl<DB: Database + 'static> EngineHook for PruneHook<DB> {
         &mut self,
         cx: &mut Context<'_>,
         ctx: EngineContext,
-    ) -> Poll<(EngineHookEvent, Option<EngineHookAction>)> {
+    ) -> Poll<RethResult<(EngineHookEvent, Option<EngineHookAction>)>> {
         // Try to spawn a pruner
         match self.try_spawn_pruner(ctx.tip_block_number) {
             Some((EngineHookEvent::NotReady, _)) => return Poll::Pending,
-            Some((event, action)) => return Poll::Ready((event, action)),
+            Some((event, action)) => return Poll::Ready(Ok((event, action))),
             None => (),
         }
 
@@ -176,3 +170,16 @@ struct Metrics {
     /// The number of times the pruner was run.
     runs: Counter,
 }
+
+impl From<PrunerError> for EngineHookError {
+    fn from(err: PrunerError) -> Self {
+        match err {
+            PrunerError::PrunePart(_) | PrunerError::InconsistentData(_) => {
+                EngineHookError::Internal(Box::new(err))
+            }
+            PrunerError::Interface(err) => err.into(),
+            PrunerError::Database(err) => RethError::Database(err).into(),
+            PrunerError::Provider(err) => RethError::Provider(err).into(),
+        }
+    }
+}

crates/consensus/beacon/src/engine/hooks/snapshot.rs

@@ -0,0 +1,161 @@
+//! Snapshot hook for the engine implementation.
+
+use crate::{
+    engine::hooks::{
+        EngineContext, EngineHook, EngineHookAction, EngineHookError, EngineHookEvent,
+    },
+    hooks::EngineHookDBAccessLevel,
+};
+use futures::FutureExt;
+use reth_db::database::Database;
+use reth_interfaces::{RethError, RethResult};
+use reth_primitives::BlockNumber;
+use reth_snapshot::{Snapshotter, SnapshotterError, SnapshotterWithResult};
+use reth_tasks::TaskSpawner;
+use std::task::{ready, Context, Poll};
+use tokio::sync::oneshot;
+
+/// Manages snapshotting under the control of the engine.
+///
+/// This type controls the [Snapshotter].
+#[derive(Debug)]
+pub struct SnapshotHook<DB> {
+    /// The current state of the snapshotter.
+    state: SnapshotterState<DB>,
+    /// The type that can spawn the snapshotter task.
+    task_spawner: Box<dyn TaskSpawner>,
+}
+
+impl<DB: Database + 'static> SnapshotHook<DB> {
+    /// Create a new instance
+    pub fn new(snapshotter: Snapshotter<DB>, task_spawner: Box<dyn TaskSpawner>) -> Self {
+        Self { state: SnapshotterState::Idle(Some(snapshotter)), task_spawner }
+    }
+
+    /// Advances the snapshotter state.
+    ///
+    /// This checks for the result in the channel, or returns pending if the snapshotter is idle.
+    fn poll_snapshotter(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<RethResult<(EngineHookEvent, Option<EngineHookAction>)>> {
+        let result = match self.state {
+            SnapshotterState::Idle(_) => return Poll::Pending,
+            SnapshotterState::Running(ref mut fut) => {
+                ready!(fut.poll_unpin(cx))
+            }
+        };
+
+        let event = match result {
+            Ok((snapshotter, result)) => {
+                self.state = SnapshotterState::Idle(Some(snapshotter));
+
+                match result {
+                    Ok(_) => EngineHookEvent::Finished(Ok(())),
+                    Err(err) => EngineHookEvent::Finished(Err(err.into())),
+                }
+            }
+            Err(_) => {
+                // failed to receive the snapshotter
+                EngineHookEvent::Finished(Err(EngineHookError::ChannelClosed))
+            }
+        };
+
+        Poll::Ready(Ok((event, None)))
+    }
+
+    /// This will try to spawn the snapshotter if it is idle:
+    /// 1. Check if snapshotting is needed through [Snapshotter::get_snapshot_targets] and then
+    ///    [SnapshotTargets::any](reth_snapshot::SnapshotTargets::any).
+    /// 2.
+    ///     1. If snapshotting is needed, pass snapshot request to the [Snapshotter::run] and spawn
+    ///        it in a separate task. Set snapshotter state to [SnapshotterState::Running].
+    ///     2. If snapshotting is not needed, set snapshotter state back to
+    ///        [SnapshotterState::Idle].
+    ///
+    /// If snapshotter is already running, do nothing.
+    fn try_spawn_snapshotter(
+        &mut self,
+        finalized_block_number: BlockNumber,
+    ) -> RethResult<Option<(EngineHookEvent, Option<EngineHookAction>)>> {
+        Ok(match &mut self.state {
+            SnapshotterState::Idle(snapshotter) => {
+                let Some(mut snapshotter) = snapshotter.take() else { return Ok(None) };
+
+                let targets = snapshotter.get_snapshot_targets(finalized_block_number)?;
+
+                // Check if the snapshotting of any parts has been requested.
+                if targets.any() {
+                    let (tx, rx) = oneshot::channel();
+                    self.task_spawner.spawn_critical_blocking(
+                        "snapshotter task",
+                        Box::pin(async move {
+                            let result = snapshotter.run(targets);
+                            let _ = tx.send((snapshotter, result));
+                        }),
+                    );
+                    self.state = SnapshotterState::Running(rx);
+
+                    Some((EngineHookEvent::Started, None))
+                } else {
+                    self.state = SnapshotterState::Idle(Some(snapshotter));
+                    Some((EngineHookEvent::NotReady, None))
+                }
+            }
+            SnapshotterState::Running(_) => None,
+        })
+    }
+}
+
+impl<DB: Database + 'static> EngineHook for SnapshotHook<DB> {
+    fn name(&self) -> &'static str {
+        "Snapshot"
+    }
+
+    fn poll(
+        &mut self,
+        cx: &mut Context<'_>,
+        ctx: EngineContext,
+    ) -> Poll<RethResult<(EngineHookEvent, Option<EngineHookAction>)>> {
+        let Some(finalized_block_number) = ctx.finalized_block_number else {
+            return Poll::Ready(Ok((EngineHookEvent::NotReady, None)))
+        };
+
+        // Try to spawn a snapshotter
+        match self.try_spawn_snapshotter(finalized_block_number)? {
+            Some((EngineHookEvent::NotReady, _)) => return Poll::Pending,
+            Some((event, action)) => return Poll::Ready(Ok((event, action))),
+            None => (),
+        }
+
+        // Poll snapshotter and check its status
+        self.poll_snapshotter(cx)
+    }
+
+    fn db_access_level(&self) -> EngineHookDBAccessLevel {
+        EngineHookDBAccessLevel::ReadOnly
+    }
+}
+
+/// The possible snapshotter states within the sync controller.
+///
+/// [SnapshotterState::Idle] means that the snapshotter is currently idle.
+/// [SnapshotterState::Running] means that the snapshotter is currently running.
+#[derive(Debug)]
+enum SnapshotterState<DB> {
+    /// Snapshotter is idle.
+    Idle(Option<Snapshotter<DB>>),
+    /// Snapshotter is running and waiting for a response
+    Running(oneshot::Receiver<SnapshotterWithResult<DB>>),
+}
+
+impl From<SnapshotterError> for EngineHookError {
+    fn from(err: SnapshotterError) -> Self {
+        match err {
+            SnapshotterError::InconsistentData(_) => EngineHookError::Internal(Box::new(err)),
+            SnapshotterError::Interface(err) => err.into(),
+            SnapshotterError::Database(err) => RethError::Database(err).into(),
+            SnapshotterError::Provider(err) => RethError::Provider(err).into(),
+        }
+    }
+}

crates/consensus/beacon/src/engine/mod.rs

@@ -1796,7 +1796,10 @@ where
             // any engine messages until it's finished.
             if let Poll::Ready(result) = this.hooks.poll_running_hook_with_db_write(
                 cx,
-                EngineContext { tip_block_number: this.blockchain.canonical_tip().number },
+                EngineContext {
+                    tip_block_number: this.blockchain.canonical_tip().number,
+                    finalized_block_number: this.blockchain.finalized_block_number()?,
+                },
             )? {
                 this.on_hook_result(result)?;
             }
@@ -1856,7 +1859,10 @@ where
             if !this.forkchoice_state_tracker.is_latest_invalid() {
                 if let Poll::Ready(result) = this.hooks.poll_next_hook(
                     cx,
-                    EngineContext { tip_block_number: this.blockchain.canonical_tip().number },
+                    EngineContext {
+                        tip_block_number: this.blockchain.canonical_tip().number,
+                        finalized_block_number: this.blockchain.finalized_block_number()?,
+                    },
                     this.sync.is_pipeline_active(),
                 )? {
                     this.on_hook_result(result)?;

crates/snapshot/Cargo.toml

@@ -0,0 +1,30 @@
+[package]
+name = "reth-snapshot"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+homepage.workspace = true
+repository.workspace = true
+description = """
+Snapshotting implementation
+"""
+
+[dependencies]
+# reth
+reth-primitives.workspace = true
+reth-db.workspace = true
+reth-provider.workspace = true
+reth-interfaces.workspace = true
+
+# misc
+thiserror.workspace = true
+
+[dev-dependencies]
+# reth
+reth-db = { workspace = true, features = ["test-utils"] }
+reth-stages = { path = "../stages", features = ["test-utils"] }
+
+# misc
+
+assert_matches.workspace = true