Improve L1 message reorg testing (#253)

* Refactor `can_handle_reorgs_while_sequencing` test to also include a follower node and add a bunch of convenience methods

* feat: attach L1 block number to ScrollPayloadAttributes

* feat: correctly handle L1 reorgs in driver

* adjust can_handle_l1_message_reorg test to assert correct reorg conditions

* fix linter errors

* make sequencer not issue blocks by default in tests

* improve test

* fixes after merge

* improve test

* fixes after merge

* improve test helper functions

* feat: answer comments

---------

Co-authored-by: Gregory Edison <[email protected]>

Author: jonastheis

Cargo.lock

@@ -586,7 +586,7 @@ pub trait DatabaseOperations: DatabaseConnectionProvider {
 ///
 /// It can either be an index, which is the queue index of the first message to return, or a hash,
 /// which is the hash of the first message to return.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum L1MessageStart {
     /// Start from the provided queue index.
     Index(u64),

crates/database/db/src/operations.rs

@@ -24,8 +24,8 @@ use core::{
 };
 use futures::{FutureExt, Stream};
 use rollup_node_primitives::{
-    BatchCommitData, BatchInfo, ScrollPayloadAttributesWithBatchInfo, WithFinalizedBatchInfo,
-    WithFinalizedBlockNumber,
+    BatchCommitData, BatchInfo, ScrollPayloadAttributesWithBatchInfo, WithBlockNumber,
+    WithFinalizedBatchInfo, WithFinalizedBlockNumber,
 };
 use rollup_node_providers::{BlockDataProvider, L1Provider};
 use scroll_alloy_rpc_types_engine::{BlockDataHint, ScrollPayloadAttributes};
@@ -186,7 +186,7 @@ impl<P> Stream for DerivationPipeline<P>
 where
     P: L1Provider + Clone + Unpin + Send + Sync + 'static,
 {
-    type Item = ScrollPayloadAttributesWithBatchInfo;
+    type Item = WithBlockNumber<ScrollPayloadAttributesWithBatchInfo>;
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.get_mut();
@@ -198,7 +198,7 @@ where
 
         // return attributes from the queue if any.
         if let Some(attribute) = this.attributes_queue.pop_front() {
-            return Poll::Ready(Some(attribute.inner))
+            return Poll::Ready(Some(attribute))
         }
 
         // if future is None and the batch queue is empty, store the waker and return.
@@ -487,8 +487,10 @@ mod tests {
 
         // check the correctness of the last attribute.
         let mut attribute = ScrollPayloadAttributes::default();
-        while let Some(ScrollPayloadAttributesWithBatchInfo { payload_attributes: a, .. }) =
-            pipeline.next().await
+        while let Some(WithBlockNumber {
+            inner: ScrollPayloadAttributesWithBatchInfo { payload_attributes: a, .. },
+            ..
+        }) = pipeline.next().await
         {
             if a.payload_attributes.timestamp == 1696935657 {
                 attribute = a;
@@ -545,8 +547,10 @@ mod tests {
 
         // check the correctness of the last attribute.
         let mut attribute = ScrollPayloadAttributes::default();
-        while let Some(ScrollPayloadAttributesWithBatchInfo { payload_attributes: a, .. }) =
-            pipeline.next().await
+        while let Some(WithBlockNumber {
+            inner: ScrollPayloadAttributesWithBatchInfo { payload_attributes: a, .. },
+            ..
+        }) = pipeline.next().await
         {
             if a.payload_attributes.timestamp == 1696935657 {
                 attribute = a;

crates/derivation-pipeline/src/lib.rs

@@ -8,7 +8,7 @@ use crate::{
 use alloy_provider::Provider;
 use futures::{ready, task::AtomicWaker, FutureExt, Stream};
 use rollup_node_primitives::{
-    BlockInfo, ChainImport, MeteredFuture, ScrollPayloadAttributesWithBatchInfo,
+    BlockInfo, ChainImport, MeteredFuture, ScrollPayloadAttributesWithBatchInfo, WithBlockNumber,
 };
 use scroll_alloy_hardforks::ScrollHardforks;
 use scroll_alloy_network::Scroll;
@@ -38,7 +38,7 @@ pub struct EngineDriver<EC, CS, P> {
     /// Block building duration.
     block_building_duration: Duration,
     /// The pending payload attributes derived from batches on L1.
-    l1_payload_attributes: VecDeque<ScrollPayloadAttributesWithBatchInfo>,
+    l1_payload_attributes: VecDeque<WithBlockNumber<ScrollPayloadAttributesWithBatchInfo>>,
     /// The pending block imports received over the network.
     chain_imports: VecDeque<ChainImport>,
     /// The latest optimistic sync target.
@@ -121,6 +121,67 @@ where
         }
     }
 
+    /// Handle L1 reorg, with the L1 block number reorged to, and whether this reorged the head or
+    /// batches.
+    pub fn handle_l1_reorg(
+        &mut self,
+        l1_block_number: u64,
+        reorged_unsafe_head: Option<BlockInfo>,
+        reorged_safe_head: Option<BlockInfo>,
+    ) {
+        // On an unsafe head reorg.
+        if let Some(l2_head_block_info) = reorged_unsafe_head {
+            // clear the payload building future.
+            self.payload_building_future = None;
+
+            // retain only blocks from chain imports for which the block number <= L2 reorged
+            // number.
+            for chain_import in &mut self.chain_imports {
+                chain_import.chain.retain(|block| block.number <= l2_head_block_info.number);
+            }
+
+            // reset the unsafe head.
+            self.set_head_block_info(l2_head_block_info);
+
+            // drop the engine future if it's a `NewPayload` or `BlockImport` with block number >
+            // L2 reorged number.
+            if let Some(MeteredFuture { fut, .. }) = self.engine_future.as_ref() {
+                match fut {
+                    EngineFuture::ChainImport(WithBlockNumber { number, .. })
+                        if number > &l2_head_block_info.number =>
+                    {
+                        self.engine_future = None
+                    }
+                    // `NewPayload` future is ONLY instantiated when the payload building future is
+                    // done, and we want to issue the payload to the EN. Thus, we also clear it on a
+                    // L2 reorg.
+                    EngineFuture::NewPayload(_) => self.engine_future = None,
+                    _ => {}
+                }
+            }
+        }
+
+        // On a safe head reorg: reset the safe head.
+        if let Some(safe_block_info) = reorged_safe_head {
+            self.set_safe_block_info(safe_block_info);
+        }
+
+        // drop the engine future if it's a `L1Consolidation` future associated with a L1 block
+        // number > l1_block_number.
+        if matches!(
+            self.engine_future.as_ref(),
+            Some(MeteredFuture {
+                fut: EngineFuture::L1Consolidation(WithBlockNumber { number, .. }),
+                ..
+            }) if number > &l1_block_number
+        ) {
+            self.engine_future = None;
+        }
+
+        // retain the L1 payload attributes with block number <= L1 block.
+        self.l1_payload_attributes.retain(|attribute| attribute.number <= l1_block_number);
+    }
+
     /// Handles a block import request by adding it to the queue and waking up the driver.
     pub fn handle_chain_import(&mut self, chain_import: ChainImport) {
         tracing::trace!(target: "scroll::engine", head = %chain_import.chain.last().unwrap().hash_slow(), "new chain import request received");
@@ -145,7 +206,10 @@ where
 
     /// Handles a [`ScrollPayloadAttributes`] sourced from L1 by initiating a task sending the
     /// attribute to the EN via the [`EngineDriver`].
-    pub fn handle_l1_consolidation(&mut self, attributes: ScrollPayloadAttributesWithBatchInfo) {
+    pub fn handle_l1_consolidation(
+        &mut self,
+        attributes: WithBlockNumber<ScrollPayloadAttributesWithBatchInfo>,
+    ) {
         self.l1_payload_attributes.push_back(attributes);
         self.waker.wake();
     }
@@ -193,7 +257,7 @@ where
                         self.metrics.block_import_duration.record(duration.as_secs_f64());
 
                         // Return the block import outcome
-                        return block_import_outcome.map(EngineDriverEvent::BlockImportOutcome)
+                        return block_import_outcome.map(EngineDriverEvent::BlockImportOutcome);
                     }
                     Err(err) => {
                         tracing::error!(target: "scroll::engine", ?err, "failed to import block");
@@ -223,7 +287,7 @@ where
                         // record the metric.
                         self.metrics.l1_consolidation_duration.record(duration.as_secs_f64());
 
-                        return Some(EngineDriverEvent::L1BlockConsolidated(consolidation_outcome))
+                        return Some(EngineDriverEvent::L1BlockConsolidated(consolidation_outcome));
                     }
                     Err(err) => {
                         tracing::error!(target: "scroll::engine", ?err, "failed to consolidate block derived from L1");
@@ -250,7 +314,7 @@ where
                         self.metrics.build_new_payload_duration.record(duration.as_secs_f64());
                         self.metrics.gas_per_block.record(block.gas_used as f64);
 
-                        return Some(EngineDriverEvent::NewPayload(block))
+                        return Some(EngineDriverEvent::NewPayload(block));
                     }
                     Err(err) => {
                         tracing::error!(target: "scroll::engine", ?err, "failed to build new payload");

crates/engine/src/driver.rs

@@ -1,5 +1,5 @@
 use alloy_rpc_types_engine::PayloadError;
-use rollup_node_primitives::ScrollPayloadAttributesWithBatchInfo;
+use rollup_node_primitives::{ScrollPayloadAttributesWithBatchInfo, WithBlockNumber};
 use scroll_alloy_provider::ScrollEngineApiError;
 use scroll_alloy_rpc_types_engine::ScrollPayloadAttributes;
 
@@ -21,7 +21,7 @@ pub enum EngineDriverError {
     /// The payload id field is missing in the forkchoice update response for an L1 consolidation
     /// job.
     #[error("Forkchoice update response missing payload id for L1 consolidation job")]
-    L1ConsolidationMissingPayloadId(ScrollPayloadAttributesWithBatchInfo),
+    L1ConsolidationMissingPayloadId(WithBlockNumber<ScrollPayloadAttributesWithBatchInfo>),
     /// The payload id field is missing in the forkchoice update response for a payload building
     /// job.
     #[error("Forkchoice update response missing payload id for payload building job")]

crates/engine/src/error.rs

@@ -11,7 +11,7 @@ use reth_scroll_engine_primitives::try_into_block;
 use reth_scroll_primitives::ScrollBlock;
 use rollup_node_primitives::{
     BatchInfo, BlockInfo, ChainImport, L2BlockInfoWithL1Messages, MeteredFuture,
-    ScrollPayloadAttributesWithBatchInfo,
+    ScrollPayloadAttributesWithBatchInfo, WithBlockNumber,
 };
 use scroll_alloy_hardforks::ScrollHardforks;
 use scroll_alloy_network::Scroll;
@@ -47,7 +47,7 @@ type L1ConsolidationFuture =
     Pin<Box<dyn Future<Output = Result<ConsolidationOutcome, EngineDriverError>> + Send>>;
 
 /// An enum that represents the different outcomes of an L1 consolidation job.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum ConsolidationOutcome {
     /// Represents a successful consolidation outcome with the consolidated block info and batch
     /// info.
@@ -97,8 +97,8 @@ pub(crate) type OptimisticSyncFuture =
 /// An enum that represents the different types of futures that can be executed on the engine API.
 /// It can be a block import job, an L1 consolidation job, or a new payload processing.
 pub(crate) enum EngineFuture {
-    ChainImport(ChainImportFuture),
-    L1Consolidation(L1ConsolidationFuture),
+    ChainImport(WithBlockNumber<ChainImportFuture>),
+    L1Consolidation(WithBlockNumber<L1ConsolidationFuture>),
     NewPayload(NewPayloadFuture),
     OptimisticSync(OptimisticSyncFuture),
 }
@@ -112,7 +112,11 @@ impl EngineFuture {
     where
         EC: ScrollEngineApi + Unpin + Send + Sync + 'static,
     {
-        Self::ChainImport(Box::pin(handle_chain_import(client, chain_import, fcs)))
+        let highest_block_number = chain_import.chain.last().unwrap().number;
+        Self::ChainImport(WithBlockNumber::new(
+            highest_block_number,
+            Box::pin(handle_chain_import(client, chain_import, fcs)),
+        ))
     }
 
     pub(crate) fn optimistic_sync<EC>(client: Arc<EC>, fcs: AlloyForkchoiceState) -> Self
@@ -127,18 +131,21 @@ impl EngineFuture {
         client: Arc<EC>,
         execution_payload_provider: P,
         fcs: ForkchoiceState,
-        payload_attributes: ScrollPayloadAttributesWithBatchInfo,
+        payload_attributes: WithBlockNumber<ScrollPayloadAttributesWithBatchInfo>,
     ) -> Self
     where
         EC: ScrollEngineApi + Unpin + Send + Sync + 'static,
         P: Provider<Scroll> + Unpin + Send + Sync + 'static,
     {
-        Self::L1Consolidation(Box::pin(handle_payload_attributes(
-            client,
-            execution_payload_provider,
-            fcs,
-            payload_attributes,
-        )))
+        Self::L1Consolidation(WithBlockNumber::new(
+            payload_attributes.number,
+            Box::pin(handle_payload_attributes(
+                client,
+                execution_payload_provider,
+                fcs,
+                payload_attributes,
+            )),
+        ))
     }
 
     /// Creates a new [`EngineFuture::NewPayload`] future from the provided parameters.
@@ -162,8 +169,8 @@ impl Future for EngineFuture {
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<EngineDriverFutureResult> {
         let this = self.get_mut();
         match this {
-            Self::ChainImport(fut) => fut.as_mut().poll(cx).map(Into::into),
-            Self::L1Consolidation(fut) => fut.as_mut().poll(cx).map(Into::into),
+            Self::ChainImport(fut) => fut.inner.as_mut().poll(cx).map(Into::into),
+            Self::L1Consolidation(fut) => fut.inner.as_mut().poll(cx).map(Into::into),
             Self::NewPayload(fut) => fut.as_mut().poll(cx).map(Into::into),
             Self::OptimisticSync(fut) => fut.as_mut().poll(cx).map(Into::into),
         }
@@ -258,7 +265,7 @@ async fn handle_payload_attributes<EC, P>(
     client: Arc<EC>,
     provider: P,
     fcs: ForkchoiceState,
-    payload_attributes_with_batch_info: ScrollPayloadAttributesWithBatchInfo,
+    payload_attributes_with_batch_info: WithBlockNumber<ScrollPayloadAttributesWithBatchInfo>,
 ) -> Result<ConsolidationOutcome, EngineDriverError>
 where
     EC: ScrollEngineApi + Unpin + Send + Sync + 'static,
@@ -267,7 +274,7 @@ where
     tracing::trace!(target: "scroll::engine::future", ?fcs, ?payload_attributes_with_batch_info, "handling payload attributes");
 
     let ScrollPayloadAttributesWithBatchInfo { mut payload_attributes, batch_info } =
-        payload_attributes_with_batch_info.clone();
+        payload_attributes_with_batch_info.inner.clone();
 
     let maybe_execution_payload = provider
         .get_block((fcs.safe_block_info().number + 1).into())