fix: handle batch revert (#275)

* feat: update the chain orchestrator to return finalized batches on finalized events

* feat: update manager to only push finalized batches in the derivation pipeline

* feat: update engine to correctly set finalized head

* test: fix graceful shutdown test

* fix: clippy

* fix: codespell

Author: greged93

crates/chain-orchestrator/src/event.rs

@@ -2,7 +2,9 @@ use alloy_consensus::Header;
 use alloy_primitives::{Signature, B256};
 use reth_network_peers::PeerId;
 use reth_scroll_primitives::ScrollBlock;
-use rollup_node_primitives::{BatchInfo, BlockInfo, ChainImport, L2BlockInfoWithL1Messages};
+use rollup_node_primitives::{
+    BatchInfo, BlockInfo, ChainImport, L2BlockInfoWithL1Messages, WithFinalizedBlockNumber,
+};
 
 /// An event emitted by the `ChainOrchestrator`.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -38,12 +40,12 @@ pub enum ChainOrchestratorEvent {
         /// The safe L2 block info.
         safe_head: Option<BlockInfo>,
     },
-    /// A batch has been finalized returning the batch hash and new an optional finalized
-    /// L2 block.
-    BatchFinalized(B256, Option<BlockInfo>),
-    /// An L1 block has been finalized returning the L1 block number and an optional
-    /// finalized L2 block.
-    L1BlockFinalized(u64, Option<BlockInfo>),
+    /// A batch has been finalized returning an optional finalized L2 block. Also returns a
+    /// [`BatchInfo`] if the finalized event occurred in a finalized L1 block.
+    BatchFinalized(Option<WithFinalizedBlockNumber<BatchInfo>>, Option<BlockInfo>),
+    /// An L1 block has been finalized returning the L1 block number, the list of finalized batches
+    /// and an optional finalized L2 block.
+    L1BlockFinalized(u64, Vec<BatchInfo>, Option<BlockInfo>),
     /// A `L1Message` event has been committed returning the message queue index.
     L1MessageCommitted(u64),
     /// A reorg has occurred on L1, returning the L1 block number of the new L1 head,

crates/chain-orchestrator/src/lib.rs

@@ -11,7 +11,7 @@ use reth_network_p2p::{BlockClient, BodiesClient};
 use reth_scroll_primitives::ScrollBlock;
 use rollup_node_primitives::{
     BatchCommitData, BatchInfo, BlockInfo, BoundedVec, ChainImport, L1MessageEnvelope,
-    L2BlockInfoWithL1Messages,
+    L2BlockInfoWithL1Messages, WithBlockNumber,
 };
 use rollup_node_watcher::L1Notification;
 use scroll_alloy_consensus::TxL1Message;
@@ -21,6 +21,7 @@ use scroll_db::{Database, DatabaseError, DatabaseOperations, L1MessageStart, Unw
 use scroll_network::NewBlockWithPeer;
 use std::{
     collections::{HashMap, VecDeque},
+    ops::Add,
     pin::Pin,
     sync::{
         atomic::{AtomicU64, Ordering},
@@ -545,12 +546,13 @@ impl<
                     )),
                 ))
             }
-            L1Notification::BatchFinalization { hash, block_number, .. } => {
+            L1Notification::BatchFinalization { hash, index, block_number } => {
                 ChainOrchestratorFuture::HandleBatchFinalization(self.handle_metered(
                     ChainOrchestratorItem::BatchFinalization,
                     Box::pin(Self::handle_batch_finalization(
                         self.database.clone(),
                         hash,
+                        index,
                         block_number,
                         self.l1_finalized_block_number.clone(),
                         self.l2_finalized_block_number.clone(),
@@ -602,8 +604,8 @@ impl<
         }))
     }
 
-    /// Handles a finalized event by updating the chain orchestrator L1 finalized block and
-    /// returning the new finalized L2 chain block.
+    /// Handles a finalized event by updating the chain orchestrator L1 finalized block, returning
+    /// the new finalized L2 chain block and the list of finalized batches.
     async fn handle_finalized(
         database: Arc<Database>,
         block_number: u64,
@@ -613,20 +615,30 @@ impl<
         // Set the latest finalized L1 block in the database.
         database.set_latest_finalized_l1_block_number(block_number).await?;
 
-        // get the newest finalized batch.
-        let batch_hash = database.get_finalized_batch_hash_at_height(block_number).await?;
+        // get the finalized batch infos.
+        // we add 1 to the low finalized l1 block number to avoid fetching the last finalized batch
+        // a second time.
+        let low_finalized_l1_block_number =
+            l1_block_number.load(Ordering::Relaxed).add(1).max(block_number);
+        let finalized_batches = database
+            .get_batches_by_finalized_block_range(low_finalized_l1_block_number, block_number)
+            .await?;
 
         // get the finalized block for the batch.
-        let finalized_block = if let Some(hash) = batch_hash {
-            Self::fetch_highest_finalized_block(database, hash, l2_block_number).await?
+        let finalized_block = if let Some(info) = finalized_batches.last() {
+            Self::fetch_highest_finalized_block(database, info.hash, l2_block_number).await?
         } else {
             None
         };
 
         // update the chain orchestrator l1 block number.
         l1_block_number.store(block_number, Ordering::Relaxed);
 
-        Ok(Some(ChainOrchestratorEvent::L1BlockFinalized(block_number, finalized_block)))
+        Ok(Some(ChainOrchestratorEvent::L1BlockFinalized(
+            block_number,
+            finalized_batches,
+            finalized_block,
+        )))
     }
 
     /// Handles an L1 message by inserting it into the database.
@@ -699,23 +711,30 @@ impl<
     async fn handle_batch_finalization(
         database: Arc<Database>,
         batch_hash: B256,
+        batch_index: u64,
         block_number: u64,
         l1_block_number: Arc<AtomicU64>,
         l2_block_number: Arc<AtomicU64>,
     ) -> Result<Option<ChainOrchestratorEvent>, ChainOrchestratorError> {
         // finalized the batch.
         database.finalize_batch(batch_hash, block_number).await?;
 
-        // check if the block where the batch was finalized is finalized on L1.
         let mut finalized_block = None;
+        let mut finalized_batch = None;
+
+        // check if the block where the batch was finalized is finalized on L1.
         let l1_block_number_value = l1_block_number.load(Ordering::Relaxed);
-        if l1_block_number_value > block_number {
+        if l1_block_number_value >= block_number {
             // fetch the finalized block.
             finalized_block =
                 Self::fetch_highest_finalized_block(database, batch_hash, l2_block_number).await?;
+
+            // set the finalized batch info.
+            finalized_batch =
+                Some(WithBlockNumber::new(block_number, BatchInfo::new(batch_index, batch_hash)));
         }
 
-        let event = ChainOrchestratorEvent::BatchFinalized(batch_hash, finalized_block);
+        let event = ChainOrchestratorEvent::BatchFinalized(finalized_batch, finalized_block);
         Ok(Some(event))
     }
 

crates/database/db/src/db.rs

@@ -228,7 +228,7 @@ mod test {
     }
 
     #[tokio::test]
-    async fn test_database_finalized_batch_hash_at_height() {
+    async fn test_database_batches_by_finalized_block_range() {
         // Set up the test database.
         let db = setup_test_db().await;
 
@@ -240,7 +240,7 @@ mod test {
         // Generate randoms BatchInfoCommitData, insert in database and finalize.
         let mut block_number = 100;
         let mut batch_index = 100;
-        let mut highest_finalized_batch_hash = B256::ZERO;
+        let mut finalized_batches_hashes = vec![];
 
         for _ in 0..20 {
             let data = BatchCommitData {
@@ -251,13 +251,9 @@ mod test {
             let hash = data.hash;
             db.insert_batch(data).await.unwrap();
 
-            // save batch hash finalized at block number 109.
-            if block_number == 109 {
-                highest_finalized_batch_hash = hash;
-            }
-
             // Finalize batch up to block number 110.
             if block_number <= 110 {
+                finalized_batches_hashes.push(hash);
                 db.finalize_batch(hash, block_number).await.unwrap();
             }
 
@@ -266,9 +262,14 @@ mod test {
         }
 
         // Fetch the finalized batch for provided height and verify number.
-        let highest_batch_hash_from_db =
-            db.get_finalized_batch_hash_at_height(109).await.unwrap().unwrap();
-        assert_eq!(highest_finalized_batch_hash, highest_batch_hash_from_db);
+        let batch_infos = db
+            .get_batches_by_finalized_block_range(100, 110)
+            .await
+            .unwrap()
+            .into_iter()
+            .map(|b| b.hash)
+            .collect::<Vec<_>>();
+        assert_eq!(finalized_batches_hashes, batch_infos);
     }
 
     #[tokio::test]

crates/database/db/src/operations.rs

@@ -115,24 +115,31 @@ pub trait DatabaseOperations: DatabaseConnectionProvider {
             .map(|x| x.and_then(|x| x.parse::<u64>().ok()))?)
     }
 
-    /// Get the newest finalized batch hash up to or at the provided height.
-    async fn get_finalized_batch_hash_at_height(
+    /// Get the finalized batches between the provided range \[low; high\].
+    async fn get_batches_by_finalized_block_range(
         &self,
-        height: u64,
-    ) -> Result<Option<B256>, DatabaseError> {
+        low: u64,
+        high: u64,
+    ) -> Result<Vec<BatchInfo>, DatabaseError> {
         Ok(models::batch_commit::Entity::find()
             .filter(
                 Condition::all()
                     .add(models::batch_commit::Column::FinalizedBlockNumber.is_not_null())
-                    .add(models::batch_commit::Column::FinalizedBlockNumber.lte(height)),
+                    .add(models::batch_commit::Column::FinalizedBlockNumber.gte(low))
+                    .add(models::batch_commit::Column::FinalizedBlockNumber.lte(high)),
             )
-            .order_by_desc(models::batch_commit::Column::Index)
+            .order_by_asc(models::batch_commit::Column::Index)
             .select_only()
+            .column(models::batch_commit::Column::Index)
             .column(models::batch_commit::Column::Hash)
-            .into_tuple::<Vec<u8>>()
-            .one(self.get_connection())
+            .into_tuple::<(i64, Vec<u8>)>()
+            .all(self.get_connection())
             .await
-            .map(|x| x.map(|x| B256::from_slice(&x)))?)
+            .map(|x| {
+                x.into_iter()
+                    .map(|(index, hash)| BatchInfo::new(index as u64, B256::from_slice(&hash)))
+                    .collect()
+            })?)
     }
 
     /// Delete all [`BatchCommitData`]s with a block number greater than the provided block number.

crates/derivation-pipeline/benches/pipeline.rs

@@ -84,7 +84,7 @@ fn benchmark_pipeline_derivation(c: &mut Criterion) {
 
                     // commit 1000 batches.
                     for _ in 0..1000 {
-                        pipeline.handle_batch_commit(batch_info, 0);
+                        pipeline.push_batch(batch_info, 0);
                     }
 
                     tx.send(pipeline).unwrap();