chore: refactor and simplify stream_graph and unordered_stream_graph (#5893)

Author: hanabi1224

src/ipld/util.rs

@@ -119,8 +119,14 @@ pin_project! {
 }
 
 impl<DB, T> ChainStream<DB, T> {
-    pub fn with_seen(self, seen: CidHashSet) -> Self {
-        ChainStream { seen, ..self }
+    pub fn with_seen(mut self, seen: CidHashSet) -> Self {
+        self.seen = seen;
+        self
+    }
+
+    pub fn fail_on_dead_links(mut self, fail_on_dead_links: bool) -> Self {
+        self.fail_on_dead_links = fail_on_dead_links;
+        self
     }
 
     #[allow(dead_code)]
@@ -162,14 +168,7 @@ pub fn stream_graph<DB: Blockstore, T: Borrow<Tipset>, ITER: Iterator<Item = T>
     tipset_iter: ITER,
     stateroot_limit: ChainEpoch,
 ) -> ChainStream<DB, ITER> {
-    ChainStream {
-        tipset_iter,
-        db,
-        dfs: VecDeque::new(),
-        seen: CidHashSet::default(),
-        stateroot_limit,
-        fail_on_dead_links: false,
-    }
+    stream_chain(db, tipset_iter, stateroot_limit).fail_on_dead_links(false)
 }
 
 impl<DB: Blockstore, T: Borrow<Tipset>, ITER: Iterator<Item = T> + Unpin> Stream
@@ -287,7 +286,7 @@ pin_project! {
         block_receiver: flume::Receiver<anyhow::Result<CarBlock>>,
         extract_sender: flume::Sender<Cid>,
         stateroot_limit: ChainEpoch,
-        queue: Vec<(Cid,Option<Vec<u8>>)>,
+        queue: Vec<(Cid, Option<Vec<u8>>)>,
         fail_on_dead_links: bool,
     }
 
@@ -308,29 +307,18 @@ impl<DB, T> UnorderedChainStream<DB, T> {
     }
 }
 
-/// Stream all blocks that are reachable before the `stateroot_limit` epoch in an unordered fashion.
-/// After this limit, only block headers are streamed. Any dead links are reported as errors.
-///
-/// # Arguments
-///
-/// * `db` - A database that implements [`Blockstore`] interface.
-/// * `tipset_iter` - An iterator of [`Tipset`], descending order `$child -> $parent`.
-/// * `stateroot_limit` - An epoch that signifies how far back we need to inspect tipsets, in-depth.
-///   This has to be pre-calculated using this formula: `$cur_epoch - $depth`, where `$depth` is the
-///   number of `[`Tipset`]` that needs inspection.
-#[allow(dead_code)]
-pub fn unordered_stream_chain<
+fn unordered_stream_chain_inner<
     DB: Blockstore + Sync + Send + 'static,
     T: Borrow<Tipset>,
     ITER: Iterator<Item = T> + Unpin + Send + 'static,
 >(
     db: Arc<DB>,
     tipset_iter: ITER,
     stateroot_limit: ChainEpoch,
+    fail_on_dead_links: bool,
 ) -> UnorderedChainStream<DB, ITER> {
     let (sender, receiver) = flume::bounded(BLOCK_CHANNEL_LIMIT);
     let (extract_sender, extract_receiver) = flume::unbounded();
-    let fail_on_dead_links = true;
     let seen = Arc::new(Mutex::new(CidHashSet::default()));
     let handle = UnorderedChainStream::<DB, ITER>::start_workers(
         db.clone(),
@@ -353,6 +341,29 @@ pub fn unordered_stream_chain<
     }
 }
 
+/// Stream all blocks that are reachable before the `stateroot_limit` epoch in an unordered fashion.
+/// After this limit, only block headers are streamed. Any dead links are reported as errors.
+///
+/// # Arguments
+///
+/// * `db` - A database that implements [`Blockstore`] interface.
+/// * `tipset_iter` - An iterator of [`Tipset`], descending order `$child -> $parent`.
+/// * `stateroot_limit` - An epoch that signifies how far back we need to inspect tipsets, in-depth.
+///   This has to be pre-calculated using this formula: `$cur_epoch - $depth`, where `$depth` is the
+///   number of `[`Tipset`]` that needs inspection.
+#[allow(dead_code)]
+pub fn unordered_stream_chain<
+    DB: Blockstore + Sync + Send + 'static,
+    T: Borrow<Tipset>,
+    ITER: Iterator<Item = T> + Unpin + Send + 'static,
+>(
+    db: Arc<DB>,
+    tipset_iter: ITER,
+    stateroot_limit: ChainEpoch,
+) -> UnorderedChainStream<DB, ITER> {
+    unordered_stream_chain_inner(db, tipset_iter, stateroot_limit, true)
+}
+
 // Stream available graph in unordered search. All reachable nodes are touched and dead-links
 // are ignored.
 pub fn unordered_stream_graph<
@@ -364,29 +375,7 @@ pub fn unordered_stream_graph<
     tipset_iter: ITER,
     stateroot_limit: ChainEpoch,
 ) -> UnorderedChainStream<DB, ITER> {
-    let (sender, receiver) = flume::bounded(2048);
-    let (extract_sender, extract_receiver) = flume::unbounded();
-    let fail_on_dead_links = false;
-    let seen = Arc::new(Mutex::new(CidHashSet::default()));
-    let handle = UnorderedChainStream::<DB, ITER>::start_workers(
-        db.clone(),
-        sender.clone(),
-        extract_receiver,
-        seen.clone(),
-        fail_on_dead_links,
-    );
-
-    UnorderedChainStream {
-        seen,
-        db,
-        worker_handle: handle,
-        block_receiver: receiver,
-        queue: Vec::new(),
-        tipset_iter,
-        extract_sender,
-        stateroot_limit,
-        fail_on_dead_links,
-    }
+    unordered_stream_chain_inner(db, tipset_iter, stateroot_limit, false)
 }
 
 impl<DB: Blockstore + Send + Sync + 'static, T: Borrow<Tipset>, ITER: Iterator<Item = T> + Unpin>