[dag] Integrate Dag Fetcher with Dag Driver

consensus/src/dag/dag_driver.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::{
+    dag_fetcher::FetchRequester,
     order_rule::OrderRule,
     storage::DAGStorage,
     types::{CertifiedAck, DAGMessage},
@@ -14,9 +15,10 @@ use crate::{
     },
     state_replication::PayloadClient,
 };
-use anyhow::bail;
+use anyhow::{bail, Ok};
 use aptos_consensus_types::common::{Author, Payload};
 use aptos_infallible::RwLock;
+use aptos_logger::error;
 use aptos_reliable_broadcast::ReliableBroadcast;
 use aptos_time_service::{TimeService, TimeServiceTrait};
 use aptos_types::{block_info::Round, epoch_state::EpochState};
@@ -45,6 +47,7 @@ pub(crate) struct DagDriver {
     rb_abort_handle: Option<AbortHandle>,
     storage: Arc<dyn DAGStorage>,
     order_rule: OrderRule,
+    fetch_requester: Arc<FetchRequester>,
 }
 
 impl DagDriver {
@@ -58,6 +61,7 @@ impl DagDriver {
         time_service: TimeService,
         storage: Arc<dyn DAGStorage>,
         order_rule: OrderRule,
+        fetch_requester: Arc<FetchRequester>,
     ) -> Self {
         // TODO: rebroadcast nodes after recovery
         Self {
@@ -71,6 +75,7 @@ impl DagDriver {
             rb_abort_handle: None,
             storage,
             order_rule,
+            fetch_requester,
         }
     }
 
@@ -88,7 +93,9 @@ impl DagDriver {
         let round = node.metadata().round();
 
         if !dag_writer.all_exists(node.parents_metadata()) {
-            // TODO(ibalajiarun): implement fetching logic.
+            if let Err(err) = self.fetch_requester.request_for_certified_node(node) {
+                error!("request to fetch failed: {}", err);
+            }
             bail!(DagDriverError::MissingParents);
         }
 

consensus/src/dag/dag_fetcher.rs

@@ -13,14 +13,58 @@ use aptos_infallible::RwLock;
 use aptos_logger::error;
 use aptos_time_service::TimeService;
 use aptos_types::epoch_state::EpochState;
-use futures::StreamExt;
+use futures::{stream::FuturesUnordered, StreamExt};
+use tokio::sync::{oneshot, mpsc::{Sender, Receiver}};
 use std::{collections::HashMap, sync::Arc, time::Duration};
 use thiserror::Error as ThisError;
-use tokio::sync::{
-    mpsc::{Receiver, Sender},
-    oneshot,
-};
 
+pub struct FetchRequester {
+    request_tx: Sender<LocalFetchRequest>,
+    node_rx_futures: FuturesUnordered<oneshot::Receiver<Node>>,
+    certified_node_rx_futures: FuturesUnordered<oneshot::Receiver<CertifiedNode>>,
+}
+
+impl FetchRequester {
+    pub fn new(request_tx: Sender<LocalFetchRequest>) -> Self {
+        Self {
+            request_tx,
+            node_rx_futures: FuturesUnordered::new(),
+            certified_node_rx_futures: FuturesUnordered::new(),
+        }
+    }
+
+    pub fn request_for_node(&self, node: Node) -> anyhow::Result<()> {
+        let (res_tx, res_rx) = oneshot::channel();
+        let fetch_req = LocalFetchRequest::Node(node, res_tx);
+        self.request_tx
+            .try_send(fetch_req)
+            .map_err(|e| anyhow::anyhow!("unable to send fetch request to channel: {}", e))?;
+        self.node_rx_futures.push(res_rx);
+        Ok(())
+    }
+
+    pub fn request_for_certified_node(&self, node: CertifiedNode) -> anyhow::Result<()> {
+        let (res_tx, res_rx) = oneshot::channel();
+        let fetch_req = LocalFetchRequest::CertifiedNode(node, res_tx);
+        self.request_tx
+            .try_send(fetch_req)
+            .map_err(|e| anyhow::anyhow!("unable to send fetch request to channel: {}", e))?;
+        self.certified_node_rx_futures.push(res_rx);
+        Ok(())
+    }
+
+    pub async fn next_ready_node(&mut self) -> Option<Result<Node, oneshot::error::RecvError>> {
+        self.node_rx_futures.next().await
+    }
+
+    pub async fn next_ready_certified_node(
+        &mut self,
+    ) -> Option<Result<CertifiedNode, oneshot::error::RecvError>> {
+        self.certified_node_rx_futures.next().await
+    }
+}
+
+#[derive(Debug)]
 pub enum LocalFetchRequest {
     Node(Node, oneshot::Sender<Node>),
     CertifiedNode(CertifiedNode, oneshot::Sender<CertifiedNode>),
@@ -55,7 +99,7 @@ impl LocalFetchRequest {
     }
 }
 
-struct DagFetcher {
+pub struct DagFetcher {
     epoch_state: Arc<EpochState>,
     network: Arc<dyn DAGNetworkSender>,
     dag: Arc<RwLock<Dag>>,
@@ -69,7 +113,7 @@ impl DagFetcher {
         network: Arc<dyn DAGNetworkSender>,
         dag: Arc<RwLock<Dag>>,
         time_service: TimeService,
-    ) -> (Self, Sender<LocalFetchRequest>) {
+    ) -> (Self, FetchRequester) {
         let (request_tx, request_rx) = tokio::sync::mpsc::channel(16);
         (
             Self {
@@ -79,7 +123,11 @@ impl DagFetcher {
                 request_rx,
                 time_service,
             },
-            request_tx,
+            FetchRequester {
+                request_tx,
+                node_rx_futures: FuturesUnordered::new(),
+                certified_node_rx_futures: FuturesUnordered::new(),
+            },
         )
     }
 

consensus/src/dag/dag_handler.rs

@@ -1,7 +1,12 @@
 // Copyright © Aptos Foundation
 
 use super::{
-    order_rule::OrderRule, dag_driver::DagDriver, dag_fetcher::FetchRequestHandler, storage::DAGStorage, dag_network::DAGNetworkSender, types::TDAGMessage,
+    dag_driver::DagDriver,
+    dag_fetcher::{DagFetcher, FetchRequestHandler},
+    dag_network::DAGNetworkSender,
+    order_rule::OrderRule,
+    storage::DAGStorage,
+    types::TDAGMessage,
 };
 use crate::{
     dag::{
@@ -41,7 +46,7 @@ impl NetworkHandler {
         epoch_state: Arc<EpochState>,
         storage: Arc<dyn DAGStorage>,
         payload_client: Arc<dyn PayloadClient>,
-        _dag_network_sender: Arc<dyn DAGNetworkSender>,
+        dag_network_sender: Arc<dyn DAGNetworkSender>,
         rb_network_sender: Arc<dyn RBNetworkSender<DAGMessage>>,
         time_service: TimeService,
         order_rule: OrderRule,
@@ -52,6 +57,13 @@ impl NetworkHandler {
             ExponentialBackoff::from_millis(10),
             time_service.clone(),
         ));
+        // TODO: wire dag fetcher
+        let (_dag_fetcher, fetch_requester) = DagFetcher::new(
+            epoch_state.clone(),
+            dag_network_sender,
+            dag.clone(),
+            time_service.clone(),
+        );
         Self {
             dag_rpc_rx,
             node_receiver: NodeBroadcastHandler::new(
@@ -70,6 +82,7 @@ impl NetworkHandler {
                 time_service,
                 storage,
                 order_rule,
+                Arc::new(fetch_requester),
             ),
             epoch_state: epoch_state.clone(),
             fetch_receiver: FetchRequestHandler::new(dag, epoch_state),

consensus/src/dag/tests/dag_driver_tests.rs

@@ -2,13 +2,15 @@
 
 use crate::{
     dag::{
+        anchor_election::RoundRobinAnchorElection,
         dag_driver::{DagDriver, DagDriverError},
+        dag_fetcher::FetchRequester,
         dag_network::{DAGNetworkSender, RpcWithFallback},
         dag_store::Dag,
         order_rule::OrderRule,
         tests::{dag_test::MockStorage, helpers::new_certified_node},
         types::{CertifiedAck, DAGMessage},
-        RpcHandler, anchor_election::RoundRobinAnchorElection,
+        RpcHandler,
     },
     test_utils::MockPayloadManager,
 };
@@ -90,6 +92,10 @@ fn test_certified_node_handler() {
         Box::new(RoundRobinAnchorElection::new(validators)),
         ordered_nodes_sender,
     );
+
+    let (request_tx, _) = tokio::sync::mpsc::channel(10);
+    let fetch_requester = Arc::new(FetchRequester::new(request_tx));
+
     let mut driver = DagDriver::new(
         signers[0].author(),
         epoch_state,
@@ -100,6 +106,7 @@ fn test_certified_node_handler() {
         time_service,
         storage,
         order_rule,
+        fetch_requester,
     );
 
     // expect an ack for a valid message