fix(consensus-manager): Add exponential backoff for artifact timeout

Author: tthebst

rs/p2p/consensus_manager/src/receiver.rs

@@ -11,6 +11,7 @@ use axum::{
     routing::any,
     Extension, Router,
 };
+use backoff::{backoff::Backoff, ExponentialBackoffBuilder};
 use bytes::Bytes;
 use crossbeam_channel::Sender as CrossbeamSender;
 use ic_interfaces::p2p::consensus::{PriorityFnAndFilterProducer, ValidatedPoolReader};
@@ -36,7 +37,8 @@ use tokio::{
     time::{self, sleep_until, timeout_at, Instant, MissedTickBehavior},
 };
 
-const ARTIFACT_RPC_TIMEOUT: Duration = Duration::from_secs(5);
+const MIN_ARTIFACT_RPC_TIMEOUT: Duration = Duration::from_secs(5);
+const MAX_ARTIFACT_RPC_TIMEOUT: Duration = Duration::from_secs(120);
 const PRIORITY_FUNCTION_UPDATE_INTERVAL: Duration = Duration::from_secs(3);
 
 type ValidatedPoolReaderRef<T> = Arc<RwLock<dyn ValidatedPoolReader<T> + Send + Sync>>;
@@ -504,6 +506,12 @@ where
         )
         .await?;
 
+        let mut artifact_download_timeout = ExponentialBackoffBuilder::new()
+            .with_initial_interval(MIN_ARTIFACT_RPC_TIMEOUT)
+            .with_max_interval(MAX_ARTIFACT_RPC_TIMEOUT)
+            .with_max_elapsed_time(None)
+            .build();
+
         match artifact {
             // Artifact was pushed by peer. In this case we don't need check that the artifact ID corresponds
             // to the artifact because we earlier derived the ID from the artifact.
@@ -527,7 +535,14 @@ where
                         .body(bytes)
                         .unwrap();
 
-                    let next_request_at = Instant::now() + ARTIFACT_RPC_TIMEOUT;
+                    if peer_rx.has_changed().unwrap_or(false) {
+                        artifact_download_timeout.reset();
+                    }
+
+                    let next_request_at = Instant::now()
+                        + artifact_download_timeout
+                            .next_backoff()
+                            .unwrap_or(MAX_ARTIFACT_RPC_TIMEOUT);
                     match timeout_at(next_request_at, transport.rpc(&peer, request)).await {
                         Ok(Ok(response)) if response.status() == StatusCode::OK => {
                             let body = response.into_body();

rs/p2p/memory_transport/src/lib.rs

@@ -127,7 +127,7 @@ impl TransportRouter {
     pub fn add_peer(
         &mut self,
         node_id: NodeId,
-        mut router: Router,
+        router: Router,
         latency: Duration,
         capacity: usize,
     ) -> PeerTransport {
@@ -151,15 +151,19 @@ impl TransportRouter {
                 parts.extensions.insert(ConnId::from(u64::MAX));
                 let req = Request::from_parts(parts, Body::from(body));
 
-                // Call request handler
-                let resp = router.ready().await.unwrap().call(req).await.unwrap();
+                let router_resp_tx = router_resp_tx.clone();
+                let mut router = router.clone();
+                tokio::spawn(async move {
+                    // Call request handler
+                    let resp = router.ready().await.unwrap().call(req).await.unwrap();
 
-                // Transform request back to `Request<Bytes>` and attach this node in the extension map.
-                let (mut parts, body) = resp.into_parts();
-                let body = to_bytes(body).await.unwrap();
-                parts.extensions.insert(this_node_id);
-                let resp = Response::from_parts(parts, body);
-                let _ = router_resp_tx.send((resp, origin_id, oneshot_tx));
+                    // Transform request back to `Request<Bytes>` and attach this node in the extension map.
+                    let (mut parts, body) = resp.into_parts();
+                    let body = to_bytes(body).await.unwrap();
+                    parts.extensions.insert(this_node_id);
+                    let resp = Response::from_parts(parts, body);
+                    let _ = router_resp_tx.send((resp, origin_id, oneshot_tx));
+                });
             }
         });
 

rs/p2p/test_utils/src/consensus.rs

@@ -180,20 +180,16 @@ impl TestConsensus<U64Artifact> {
 
     pub fn push_advert(&self, id: u64) {
         let mut inner = self.inner.lock().unwrap();
-        inner
-            .peer_pool
-            .entry(self.node_id)
-            .and_modify(|p| p.insert(id));
+        let my_pool = inner.peer_pool.get_mut(&self.node_id).unwrap();
+        my_pool.insert(id);
 
         inner.adverts.push_back(id);
     }
 
     pub fn push_purge(&self, id: u64) {
         let mut inner = self.inner.lock().unwrap();
-        inner
-            .peer_pool
-            .entry(self.node_id)
-            .and_modify(|p| p.remove(id));
+        let my_pool = inner.peer_pool.get_mut(&self.node_id).unwrap();
+        my_pool.remove(id);
 
         inner.purge.push_back(id);
     }