fix(derive): merge upstream changes

Author: refcell

crates/derive/src/stages/batch_queue.rs

@@ -52,14 +52,8 @@ where
     T: TelemetryProvider + Debug,
 {
     /// Create a new [ChannelBank] stage.
-    pub fn new(cfg: RollupConfig, prev: FrameQueue<DAP, CP, T>, telemetry: T) -> Self {
-        Self {
-            cfg: Arc::new(cfg),
-            telemetry,
-            channels: HashMap::new(),
-            channel_queue: VecDeque::new(),
-            prev,
-        }
+    pub fn new(cfg: Arc<RollupConfig>, prev: FrameQueue<DAP, CP, T>, telemetry: T) -> Self {
+        Self { cfg, telemetry, channels: HashMap::new(), channel_queue: VecDeque::new(), prev }
     }
 
     /// Returns the L1 origin [BlockInfo].
@@ -227,36 +221,35 @@ mod tests {
     use super::*;
     use crate::{
         stages::{
-            frame_queue::tests::new_test_frames, l1_retrieval::L1Retrieval,
-            l1_traversal::tests::new_test_traversal,
+            frame_queue::tests::new_test_frames, l1_retrieval::L1Retrieval, l1_traversal::tests::*,
         },
         traits::test_utils::{TestDAP, TestTelemetry},
     };
     use alloc::vec;
 
     #[test]
     fn test_ingest_empty_origin() {
-        let mut traversal = new_test_traversal(false, false);
+        let mut traversal = new_test_traversal(vec![], vec![]);
         traversal.block = None;
         let dap = TestDAP::default();
         let retrieval = L1Retrieval::new(traversal, dap, TestTelemetry::new());
         let frame_queue = FrameQueue::new(retrieval, TestTelemetry::new());
         let mut channel_bank =
-            ChannelBank::new(RollupConfig::default(), frame_queue, TestTelemetry::new());
+            ChannelBank::new(Arc::new(RollupConfig::default()), frame_queue, TestTelemetry::new());
         let frame = Frame::default();
         let err = channel_bank.ingest_frame(frame).unwrap_err();
         assert_eq!(err, StageError::Custom(anyhow!("No origin")));
     }
 
     #[test]
     fn test_ingest_and_prune_channel_bank() {
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let results = vec![Ok(Bytes::from(vec![0x00]))];
         let dap = TestDAP { results };
         let retrieval = L1Retrieval::new(traversal, dap, TestTelemetry::new());
         let frame_queue = FrameQueue::new(retrieval, TestTelemetry::new());
         let mut channel_bank =
-            ChannelBank::new(RollupConfig::default(), frame_queue, TestTelemetry::new());
+            ChannelBank::new(Arc::new(RollupConfig::default()), frame_queue, TestTelemetry::new());
         let mut frames = new_test_frames(100000);
         // Ingest frames until the channel bank is full and it stops increasing in size
         let mut current_size = 0;
@@ -279,16 +272,16 @@ mod tests {
 
     #[tokio::test]
     async fn test_read_empty_channel_bank() {
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let results = vec![Ok(Bytes::from(vec![0x00]))];
         let dap = TestDAP { results };
         let retrieval = L1Retrieval::new(traversal, dap, TestTelemetry::new());
         let frame_queue = FrameQueue::new(retrieval, TestTelemetry::new());
         let mut channel_bank =
-            ChannelBank::new(RollupConfig::default(), frame_queue, TestTelemetry::new());
+            ChannelBank::new(Arc::new(RollupConfig::default()), frame_queue, TestTelemetry::new());
         let err = channel_bank.read().unwrap_err();
         assert_eq!(err, StageError::Eof);
         let err = channel_bank.next_data().await.unwrap_err();
-        assert_eq!(err, StageError::Custom(anyhow!("Not Enough Data")));
+        assert_eq!(err, StageError::NotEnoughData);
     }
 }

crates/derive/src/stages/channel_bank.rs

@@ -5,6 +5,7 @@ use crate::{
     traits::{ChainProvider, DataAvailabilityProvider, LogLevel, TelemetryProvider},
     types::{Batch, BlockInfo, StageError, StageResult},
 };
+
 use alloc::vec::Vec;
 use anyhow::anyhow;
 use core::fmt::Debug;

crates/derive/src/stages/channel_reader.rs

@@ -13,7 +13,8 @@ use alloc::{boxed::Box, collections::VecDeque};
 use anyhow::anyhow;
 use async_trait::async_trait;
 
-/// The frame queue stage of the derivation pipeline.
+/// The [FrameQueue] stage of the derivation pipeline.
+/// This stage takes the output of the [L1Retrieval] stage and parses it into frames.
 #[derive(Debug)]
 pub struct FrameQueue<DAP, CP, T>
 where
@@ -35,28 +36,33 @@ where
     CP: ChainProvider + Debug,
     T: TelemetryProvider + Debug,
 {
-    /// Create a new frame queue stage.
+    /// Create a new [FrameQueue] stage with the given previous [L1Retrieval] stage.
     pub fn new(prev: L1Retrieval<DAP, CP, T>, telemetry: T) -> Self {
         Self { prev, telemetry, queue: VecDeque::new() }
     }
 
-    /// Returns the L1 origin [BlockInfo].
+    /// Returns the L1 [BlockInfo] origin.
     pub fn origin(&self) -> Option<&BlockInfo> {
         self.prev.origin()
     }
 
-    /// Fetches the next frame from the frame queue.
+    /// Fetches the next frame from the [FrameQueue].
     pub async fn next_frame(&mut self) -> StageResult<Frame> {
         if self.queue.is_empty() {
             match self.prev.next_data().await {
                 Ok(data) => {
-                    // TODO: what do we do with frame parsing errors?
                     if let Ok(frames) = Frame::parse_frames(data.as_ref()) {
                         self.queue.extend(frames);
+                    } else {
+                        // TODO: log parsing frame error
+                        // Failed to parse frames, but there may be more frames in the queue for
+                        // the pipeline to advance, so don't return an error here.
                     }
                 }
                 Err(e) => {
-                    return Err(anyhow!("Error fetching next data: {e}").into());
+                    // TODO: log retrieval error
+                    // The error must be bubbled up without a wrapper in case it's an EOF error.
+                    return Err(e);
                 }
             }
         }
@@ -68,7 +74,7 @@ where
                 ),
                 LogLevel::Debug,
             );
-            return Err(anyhow!("Not enough data").into());
+            return Err(StageError::NotEnoughData);
         }
 
         self.queue.pop_front().ok_or_else(|| anyhow!("Frame queue is impossibly empty.").into())
@@ -92,7 +98,7 @@ where
 pub(crate) mod tests {
     use super::*;
     use crate::{
-        stages::l1_traversal::tests::new_test_traversal,
+        stages::l1_traversal::tests::new_populated_test_traversal,
         traits::test_utils::{TestDAP, TestTelemetry},
         DERIVATION_VERSION_0,
     };
@@ -123,71 +129,71 @@ pub(crate) mod tests {
     #[tokio::test]
     async fn test_frame_queue_empty_bytes() {
         let telemetry = TestTelemetry::new();
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let results = vec![Ok(Bytes::from(vec![0x00]))];
         let dap = TestDAP { results };
         let retrieval = L1Retrieval::new(traversal, dap, TestTelemetry::new());
         let mut frame_queue = FrameQueue::new(retrieval, telemetry);
         let err = frame_queue.next_frame().await.unwrap_err();
-        assert_eq!(err, anyhow!("Not enough data").into());
+        assert_eq!(err, StageError::NotEnoughData);
     }
 
     #[tokio::test]
     async fn test_frame_queue_no_frames_decoded() {
         let telemetry = TestTelemetry::new();
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let results = vec![Err(StageError::Eof), Ok(Bytes::default())];
         let dap = TestDAP { results };
         let retrieval = L1Retrieval::new(traversal, dap, TestTelemetry::new());
         let mut frame_queue = FrameQueue::new(retrieval, telemetry);
         let err = frame_queue.next_frame().await.unwrap_err();
-        assert_eq!(err, anyhow!("Not enough data").into());
+        assert_eq!(err, StageError::NotEnoughData);
     }
 
     #[tokio::test]
     async fn test_frame_queue_wrong_derivation_version() {
         let telemetry = TestTelemetry::new();
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let results = vec![Ok(Bytes::from(vec![0x01]))];
         let dap = TestDAP { results };
         let retrieval = L1Retrieval::new(traversal, dap, TestTelemetry::new());
         let mut frame_queue = FrameQueue::new(retrieval, telemetry);
         let err = frame_queue.next_frame().await.unwrap_err();
-        assert_eq!(err, anyhow!("Unsupported derivation version").into());
+        assert_eq!(err, StageError::NotEnoughData);
     }
 
     #[tokio::test]
     async fn test_frame_queue_frame_too_short() {
         let telemetry = TestTelemetry::new();
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let results = vec![Ok(Bytes::from(vec![0x00, 0x01]))];
         let dap = TestDAP { results };
         let retrieval = L1Retrieval::new(traversal, dap, TestTelemetry::new());
         let mut frame_queue = FrameQueue::new(retrieval, telemetry);
         let err = frame_queue.next_frame().await.unwrap_err();
-        assert_eq!(err, anyhow!("Frame too short to decode").into());
+        assert_eq!(err, StageError::NotEnoughData);
     }
 
     #[tokio::test]
     async fn test_frame_queue_single_frame() {
         let data = new_encoded_test_frames(1);
         let telemetry = TestTelemetry::new();
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let dap = TestDAP { results: vec![Ok(data)] };
         let retrieval = L1Retrieval::new(traversal, dap, TestTelemetry::new());
         let mut frame_queue = FrameQueue::new(retrieval, telemetry);
         let frame_decoded = frame_queue.next_frame().await.unwrap();
         let frame = new_test_frames(1);
         assert_eq!(frame[0], frame_decoded);
         let err = frame_queue.next_frame().await.unwrap_err();
-        assert_eq!(err, anyhow!("Not enough data").into());
+        assert_eq!(err, StageError::Eof);
     }
 
     #[tokio::test]
     async fn test_frame_queue_multiple_frames() {
         let telemetry = TestTelemetry::new();
         let data = new_encoded_test_frames(3);
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let dap = TestDAP { results: vec![Ok(data)] };
         let retrieval = L1Retrieval::new(traversal, dap, TestTelemetry::new());
         let mut frame_queue = FrameQueue::new(retrieval, telemetry);
@@ -196,6 +202,6 @@ pub(crate) mod tests {
             assert_eq!(frame_decoded.number, i);
         }
         let err = frame_queue.next_frame().await.unwrap_err();
-        assert_eq!(err, anyhow!("Not enough data").into());
+        assert_eq!(err, StageError::Eof);
     }
 }

crates/derive/src/stages/frame_queue.rs

@@ -13,7 +13,11 @@ use alloy_primitives::Bytes;
 use anyhow::anyhow;
 use async_trait::async_trait;
 
-/// The L1 retrieval stage of the derivation pipeline.
+/// The [L1Retrieval] stage of the derivation pipeline.
+/// For each L1 [BlockInfo] pulled from the [L1Traversal] stage,
+/// [L1Retrieval] fetches the associated data from a specified
+/// [DataAvailabilityProvider]. This data is returned as a generic
+/// [DataIter] that can be iterated over.
 #[derive(Debug)]
 pub struct L1Retrieval<DAP, CP, T>
 where
@@ -37,20 +41,20 @@ where
     CP: ChainProvider,
     T: TelemetryProvider,
 {
-    /// Creates a new L1 retrieval stage with the given data availability provider and previous
-    /// stage.
+    /// Creates a new [L1Retrieval] stage with the previous [L1Traversal]
+    /// stage and given [DataAvailabilityProvider].
     pub fn new(prev: L1Traversal<CP, T>, provider: DAP, telemetry: T) -> Self {
         Self { prev, telemetry, provider, data: None }
     }
 
-    /// Returns the current L1 block in the traversal stage, if it exists.
+    /// Returns the current L1 [BlockInfo] origin from the previous
+    /// [L1Traversal] stage.
     pub fn origin(&self) -> Option<&BlockInfo> {
         self.prev.origin()
     }
 
-    /// Retrieves the next data item from the L1 retrieval stage.
-    /// If there is data, it pushes it into the next stage.
-    /// If there is no data, it returns an error.
+    /// Retrieves the next data item from the [L1Retrieval] stage.
+    /// Returns an error if there is no data.
     pub async fn next_data(&mut self) -> StageResult<Bytes> {
         if self.data.is_none() {
             self.telemetry.write(
@@ -94,15 +98,15 @@ where
 mod tests {
     use super::*;
     use crate::{
-        stages::l1_traversal::tests::new_test_traversal,
+        stages::l1_traversal::tests::*,
         traits::test_utils::{TestDAP, TestIter, TestTelemetry},
     };
     use alloc::vec;
     use alloy_primitives::Address;
 
     #[tokio::test]
     async fn test_l1_retrieval_origin() {
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let dap = TestDAP { results: vec![] };
         let telemetry = TestTelemetry::new();
         let retrieval = L1Retrieval::new(traversal, dap, telemetry);
@@ -112,7 +116,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_l1_retrieval_next_data() {
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let results = vec![Err(StageError::Eof), Ok(Bytes::default())];
         let dap = TestDAP { results };
         let telemetry = TestTelemetry::new();
@@ -140,8 +144,8 @@ mod tests {
         // Create a new traversal with no blocks or receipts.
         // This would bubble up an error if the prev stage
         // (traversal) is called in the retrieval stage.
-        let traversal = new_test_traversal(false, false);
         let telemetry = TestTelemetry::new();
+        let traversal = new_test_traversal(vec![], vec![]);
         let dap = TestDAP { results: vec![] };
         let mut retrieval =
             L1Retrieval { prev: traversal, telemetry, provider: dap, data: Some(data) };
@@ -159,7 +163,7 @@ mod tests {
             results: vec![Err(StageError::Eof)],
         };
         let telemetry = TestTelemetry::new();
-        let traversal = new_test_traversal(true, true);
+        let traversal = new_populated_test_traversal();
         let dap = TestDAP { results: vec![] };
         let mut retrieval =
             L1Retrieval { prev: traversal, telemetry, provider: dap, data: Some(data) };