feat: a task that updates the final/safe/canonical when ingesting

Author: prestwich

crates/db/src/journal/ingestor.rs

@@ -0,0 +1,74 @@
+use crate::{SignetDbRw, journal::JournalDb};
+use futures_util::{Stream, StreamExt};
+use reth::providers::ProviderResult;
+use signet_node_types::NodeTypesDbTrait;
+use signet_types::primitives::SealedHeader;
+use std::sync::Arc;
+use tokio::task::JoinHandle;
+use trevm::journal::BlockUpdate;
+
+/// A task that ingests journals into a reth database.
+#[derive(Debug)]
+pub struct JournalIngestor<Db: NodeTypesDbTrait> {
+    db: Arc<SignetDbRw<Db>>,
+}
+
+impl<Db: NodeTypesDbTrait> From<SignetDbRw<Db>> for JournalIngestor<Db> {
+    fn from(value: SignetDbRw<Db>) -> Self {
+        Self::new(value.into())
+    }
+}
+
+impl<Db: NodeTypesDbTrait> From<Arc<SignetDbRw<Db>>> for JournalIngestor<Db> {
+    fn from(value: Arc<SignetDbRw<Db>>) -> Self {
+        Self::new(value)
+    }
+}
+
+impl<Db: NodeTypesDbTrait> JournalIngestor<Db> {
+    /// Create a new `JournalIngestor` with the given database provider.
+    pub const fn new(db: Arc<SignetDbRw<Db>>) -> Self {
+        Self { db }
+    }
+
+    async fn task_future<S>(self, mut stream: S) -> ProviderResult<()>
+    where
+        S: Stream<Item = (SealedHeader, BlockUpdate<'static>)> + Send + Unpin + 'static,
+    {
+        while let Some(item) = stream.next().await {
+            // FUTURE: Sanity check that the header height matches the update
+            // height. Sanity check that both heights are 1 greater than the
+            // last height in the database.
+
+            let db = self.db.clone();
+            let (header, block_update) = item;
+
+            // DB interaction is sync, so we spawn a blocking task for it. We
+            // immediately await that task. This prevents blocking the worker
+            // thread
+            tokio::task::spawn_blocking(move || db.ingest(header, block_update))
+                .await
+                .expect("ingestion should not panic")?;
+        }
+        // Stream has ended, return Ok
+        Ok(())
+    }
+
+    /// Spawn a task to ingest journals from the provided stream.
+    pub fn spawn<S>(self, stream: S) -> JoinHandle<ProviderResult<()>>
+    where
+        S: Stream<Item = (SealedHeader, BlockUpdate<'static>)> + Send + Unpin + 'static,
+    {
+        tokio::spawn(self.task_future(stream))
+    }
+}
+
+/// Ingest journals from a stream into a reth database.
+pub async fn ingest_journals<Db, S>(db: Arc<SignetDbRw<Db>>, stream: S) -> ProviderResult<()>
+where
+    Db: NodeTypesDbTrait,
+    S: Stream<Item = (SealedHeader, BlockUpdate<'static>)> + Send + Unpin + 'static,
+{
+    let ingestor = JournalIngestor::new(db);
+    ingestor.task_future(stream).await
+}

crates/db/src/journal/mod.rs

@@ -3,76 +3,8 @@
 mod r#trait;
 pub use r#trait::JournalDb;
 
-use crate::SignetDbRw;
-use futures_util::{Stream, StreamExt};
-use reth::providers::ProviderResult;
-use signet_node_types::NodeTypesDbTrait;
-use std::sync::Arc;
-use tokio::task::JoinHandle;
-use trevm::journal::BlockUpdate;
+mod provider;
+pub use provider::JournalProviderTask;
 
-/// A task that ingests journals into a reth database.
-#[derive(Debug)]
-pub struct JournalIngestor<Db: NodeTypesDbTrait> {
-    db: Arc<SignetDbRw<Db>>,
-}
-
-impl<Db: NodeTypesDbTrait> From<SignetDbRw<Db>> for JournalIngestor<Db> {
-    fn from(value: SignetDbRw<Db>) -> Self {
-        Self::new(value.into())
-    }
-}
-
-impl<Db: NodeTypesDbTrait> From<Arc<SignetDbRw<Db>>> for JournalIngestor<Db> {
-    fn from(value: Arc<SignetDbRw<Db>>) -> Self {
-        Self::new(value)
-    }
-}
-
-impl<Db: NodeTypesDbTrait> JournalIngestor<Db> {
-    /// Create a new `JournalIngestor` with the given database provider.
-    pub const fn new(db: Arc<SignetDbRw<Db>>) -> Self {
-        Self { db }
-    }
-
-    async fn task_future<S>(self, mut stream: S) -> ProviderResult<()>
-    where
-        S: Stream<Item = (alloy::consensus::Header, BlockUpdate<'static>)> + Send + Unpin + 'static,
-    {
-        while let Some(item) = stream.next().await {
-            // FUTURE: Sanity check that the header height matches the update
-            // height. Sanity check that both heights are 1 greater than the
-            // last height in the database.
-
-            let db = self.db.clone();
-            let (header, block_update) = item;
-
-            // DB interaction is sync, so we spawn a blocking task for it. We
-            // immediately await that task. This prevents blocking the worker
-            // thread
-            tokio::task::spawn_blocking(move || db.ingest(&header, block_update))
-                .await
-                .expect("ingestion should not panic")?;
-        }
-        // Stream has ended, return Ok
-        Ok(())
-    }
-
-    /// Spawn a task to ingest journals from the provided stream.
-    pub fn spawn<S>(self, stream: S) -> JoinHandle<ProviderResult<()>>
-    where
-        S: Stream<Item = (alloy::consensus::Header, BlockUpdate<'static>)> + Send + Unpin + 'static,
-    {
-        tokio::spawn(self.task_future(stream))
-    }
-}
-
-/// Ingest journals from a stream into a reth database.
-pub async fn ingest_journals<Db, S>(db: Arc<SignetDbRw<Db>>, stream: S) -> ProviderResult<()>
-where
-    Db: NodeTypesDbTrait,
-    S: Stream<Item = (alloy::consensus::Header, BlockUpdate<'static>)> + Send + Unpin + 'static,
-{
-    let ingestor = JournalIngestor::new(db);
-    ingestor.task_future(stream).await
-}
+mod ingestor;
+pub use ingestor::{JournalIngestor, ingest_journals};

crates/db/src/journal/provider.rs

@@ -0,0 +1,81 @@
+use crate::{DataCompat, journal::JournalDb};
+use futures_util::{Stream, StreamExt};
+use reth::{
+    providers::{
+        CanonChainTracker, DatabaseProviderFactory, DatabaseProviderRW, ProviderResult,
+        providers::BlockchainProvider,
+    },
+    rpc::types::engine::ForkchoiceState,
+};
+use signet_node_types::{NodeTypesDbTrait, SignetNodeTypes};
+use signet_types::primitives::SealedHeader;
+use tokio::task::JoinHandle;
+use trevm::journal::BlockUpdate;
+
+/// A task that processes journal updates for a specific database, and calls
+/// the appropriate methods on a [`BlockchainProvider`] to update the in-memory
+/// chain view.
+#[derive(Debug, Clone)]
+pub struct JournalProviderTask<Db: NodeTypesDbTrait> {
+    provider: BlockchainProvider<SignetNodeTypes<Db>>,
+}
+
+impl<Db: NodeTypesDbTrait> JournalProviderTask<Db> {
+    /// Instantiate a new task.
+    pub const fn new(provider: BlockchainProvider<SignetNodeTypes<Db>>) -> Self {
+        Self { provider }
+    }
+
+    /// Get a reference to the provider.
+    pub const fn provider(&self) -> &BlockchainProvider<SignetNodeTypes<Db>> {
+        &self.provider
+    }
+
+    /// Deconstruct the task into its provider.
+    pub fn into_inner(self) -> BlockchainProvider<SignetNodeTypes<Db>> {
+        self.provider
+    }
+
+    /// Create a future for the task, suitable for [`tokio::spawn`] or another
+    /// task-spawning system.
+    pub async fn task_future<S>(self, mut journals: S) -> ProviderResult<()>
+    where
+        S: Stream<Item = (SealedHeader, BlockUpdate<'static>)> + Send + Unpin + 'static,
+    {
+        loop {
+            let Some((header, block_update)) = journals.next().await else { break };
+
+            let block_hash = header.hash();
+
+            let rw = self.provider.database_provider_rw().map(DatabaseProviderRW);
+
+            let r_header = header.clone_convert();
+
+            // DB interaction is sync, so we spawn a blocking task for it. We
+            // immediately await that task. This prevents blocking the worker
+            // thread
+            tokio::task::spawn_blocking(move || rw?.ingest(header, block_update))
+                .await
+                .expect("ingestion should not panic")?;
+
+            self.provider.set_canonical_head(r_header.clone());
+            self.provider.set_safe(r_header.clone());
+            self.provider.set_finalized(r_header);
+            self.provider.on_forkchoice_update_received(&ForkchoiceState {
+                head_block_hash: block_hash,
+                safe_block_hash: block_hash,
+                finalized_block_hash: block_hash,
+            });
+        }
+
+        Ok(())
+    }
+
+    /// Spawn the journal provider task.
+    pub fn spawn<S>(self, journals: S) -> JoinHandle<ProviderResult<()>>
+    where
+        S: Stream<Item = (SealedHeader, BlockUpdate<'static>)> + Send + Unpin + 'static,
+    {
+        tokio::spawn(self.task_future(journals))
+    }
+}

crates/db/src/journal/trait.rs

@@ -18,7 +18,7 @@ pub trait JournalDb: RuWriter {
     ///
     /// This is intended to be used for tx simulation, and other purposes that
     /// need fast state access WITHTOUT needing to retrieve historical data.
-    fn ingest(&self, header: &Header, update: BlockUpdate<'_>) -> ProviderResult<()> {
+    fn ingest(&self, header: SealedHeader, update: BlockUpdate<'_>) -> ProviderResult<()> {
         let journal_hash = update.journal_hash();
 
         // TODO: remove the clone in future versions. This can be achieved by
@@ -29,7 +29,7 @@ pub trait JournalDb: RuWriter {
         let execution_outcome = ExecutionOutcome::new(bundle_state, vec![], header.number());
 
         let block: SealedBlock<TransactionSigned, Header> =
-            SealedBlock { header: SealedHeader::new(header.to_owned()), body: Default::default() };
+            SealedBlock { header, body: Default::default() };
         let block_result =
             BlockResult { sealed_block: RecoveredBlock::new(block, vec![]), execution_outcome };
 
@@ -40,7 +40,9 @@ pub trait JournalDb: RuWriter {
             std::iter::empty(),
             &block_result,
             journal_hash,
-        )
+        )?;
+
+        Ok(())
     }
 }