[resolver] Async Consumer Delivery (#3691)

Author: patrick-ogrady

consensus/src/marshal/resolver/handler.rs

@@ -58,7 +58,6 @@ impl<D: Digest> Handler<D> {
 impl<D: Digest> Consumer for Handler<D> {
     type Key = Request<D>;
     type Value = Bytes;
-    type Failure = ();
 
     async fn deliver(&mut self, key: Self::Key, value: Self::Value) -> bool {
         let (response, receiver) = oneshot::channel();
@@ -77,10 +76,6 @@ impl<D: Digest> Consumer for Handler<D> {
         }
         receiver.await.unwrap_or(false)
     }
-
-    async fn failed(&mut self, _: Self::Key, _: Self::Failure) {
-        // We don't need to do anything on failure, the resolver will retry.
-    }
 }
 
 impl<D: Digest> Producer for Handler<D> {

consensus/src/simplex/actors/resolver/ingress.rs

@@ -68,7 +68,6 @@ impl Handler {
 impl Consumer for Handler {
     type Key = U64;
     type Value = Bytes;
-    type Failure = ();
 
     async fn deliver(&mut self, key: Self::Key, value: Self::Value) -> bool {
         self.sender
@@ -82,10 +81,6 @@ impl Consumer for Handler {
             )
             .await
     }
-
-    async fn failed(&mut self, _: Self::Key, _: Self::Failure) {
-        // We don't need to do anything on failure, the resolver will retry.
-    }
 }
 
 impl Producer for Handler {

resolver/src/lib.rs

@@ -20,29 +20,22 @@ commonware_macros::stability_scope!(BETA {
         /// Type of data to retrieve.
         type Value;
 
-        /// Type used to indicate why data is not available.
-        type Failure;
-
         /// Deliver data to the consumer.
         ///
         /// Returns `true` if the data is valid.
         ///
+        /// The returned future may be dropped before completion if the
+        /// application cancels the fetch via [`Resolver::cancel`],
+        /// [`Resolver::clear`], or [`Resolver::retain`]. When this happens,
+        /// the resolver discards the validation result.
+        ///
         /// Implementations of [`Resolver`] must only invoke `deliver` for keys that were
         /// previously requested via [`Resolver::fetch`] (or its variants).
         fn deliver(
             &mut self,
             key: Self::Key,
             value: Self::Value,
         ) -> impl Future<Output = bool> + Send;
-
-        /// Let the consumer know that the data is not being fetched anymore.
-        ///
-        /// The failure is used to indicate why.
-        fn failed(
-            &mut self,
-            key: Self::Key,
-            failure: Self::Failure,
-        ) -> impl Future<Output = ()> + Send;
     }
 
     /// Responsible for fetching data and notifying a `Consumer`.
@@ -90,12 +83,22 @@ commonware_macros::stability_scope!(BETA {
         ) -> impl Future<Output = ()> + Send;
 
         /// Cancel a fetch request.
+        ///
+        /// If response validation is in progress, cancellation may drop the
+        /// [`Consumer::deliver`] future before it reports whether the data was
+        /// valid.
         fn cancel(&mut self, key: Self::Key) -> impl Future<Output = ()> + Send;
 
         /// Cancel all fetch requests.
+        ///
+        /// See [`cancel`](Self::cancel) for how cancellation affects
+        /// in-progress response validation.
         fn clear(&mut self) -> impl Future<Output = ()> + Send;
 
         /// Retain only the fetch requests that satisfy the predicate.
+        ///
+        /// Fetches not retained are canceled. See [`cancel`](Self::cancel) for
+        /// how cancellation affects in-progress response validation.
         fn retain(
             &mut self,
             predicate: impl Fn(&Self::Key) -> bool + Send + 'static,

resolver/src/p2p/config.rs

@@ -11,7 +11,7 @@ pub struct Config<
     D: Provider<PublicKey = P>,
     B: Blocker<PublicKey = P>,
     Key: Span,
-    Con: Consumer<Key = Key, Value = Bytes, Failure = ()>,
+    Con: Consumer<Key = Key, Value = Bytes>,
     Pro: Producer<Key = Key>,
 > {
     /// Manages the current set of peers.

resolver/src/p2p/engine.rs

@@ -1,6 +1,7 @@
 use super::{
     config::Config,
     fetcher::{Config as FetcherConfig, Fetcher},
+    inflight::Inflight,
     ingress::{FetchRequest, Mailbox, Message},
     metrics, wire, Producer,
 };
@@ -24,7 +25,7 @@ use commonware_utils::{
 };
 use futures::future::{self, Either};
 use rand::Rng;
-use std::{collections::HashMap, marker::PhantomData};
+use std::marker::PhantomData;
 use tracing::{debug, error, trace, warn};
 
 /// Represents a pending serve operation.
@@ -42,17 +43,14 @@ pub struct Engine<
     D: Provider<PublicKey = P>,
     B: Blocker<PublicKey = P>,
     Key: Span,
-    Con: Consumer<Key = Key, Value = Bytes, Failure = ()>,
+    Con: Consumer<Key = Key, Value = Bytes>,
     Pro: Producer<Key = Key>,
     NetS: Sender<PublicKey = P>,
     NetR: Receiver<PublicKey = P>,
 > {
     /// Context used to spawn tasks, manage time, etc.
     context: ContextCell<E>,
 
-    /// Consumes data that is fetched from the network
-    consumer: Con,
-
     /// Produces data for incoming requests
     producer: Pro,
 
@@ -71,8 +69,8 @@ pub struct Engine<
     /// Manages outgoing fetch requests
     fetcher: Fetcher<E, P, Key, NetS>,
 
-    /// Track the start time of fetch operations
-    fetch_timers: HashMap<Key, histogram::Timer<E>>,
+    /// Tracks all in-flight fetch state
+    inflight: Inflight<E, Con, P, Key>,
 
     /// Holds futures that resolve once the `Producer` has produced the data.
     /// Once the future is resolved, the data (or an error) is sent to the peer.
@@ -96,7 +94,7 @@ impl<
         D: Provider<PublicKey = P>,
         B: Blocker<PublicKey = P>,
         Key: Span,
-        Con: Consumer<Key = Key, Value = Bytes, Failure = ()>,
+        Con: Consumer<Key = Key, Value = Bytes>,
         Pro: Producer<Key = Key>,
         NetS: Sender<PublicKey = P>,
         NetR: Receiver<PublicKey = P>,
@@ -123,17 +121,16 @@ impl<
         (
             Self {
                 context: ContextCell::new(context),
-                consumer: cfg.consumer,
                 producer: cfg.producer,
                 peer_provider: cfg.peer_provider,
                 blocker: cfg.blocker,
                 last_peer_set_id: None,
                 mailbox: receiver,
                 fetcher,
+                inflight: Inflight::new(cfg.consumer),
                 serves: FuturesPool::default(),
                 priority_responses: cfg.priority_responses,
                 metrics,
-                fetch_timers: HashMap::new(),
                 _r: PhantomData,
             },
             Mailbox::new(sender),
@@ -189,6 +186,7 @@ impl<
             },
             on_stopped => {
                 debug!("shutdown");
+                self.inflight.drain();
                 self.serves.cancel_all();
             },
             // Handle peer set updates
@@ -224,7 +222,7 @@ impl<
                             trace!(?key, "mailbox: fetch");
 
                             // Check if the fetch is already in progress
-                            let is_new = !self.fetch_timers.contains_key(&key);
+                            let is_new = !self.inflight.contains(&key);
 
                             // Update targets
                             match targets {
@@ -241,7 +239,7 @@ impl<
 
                             // Only start new fetch if not already in progress
                             if is_new {
-                                self.fetch_timers
+                                self.inflight
                                     .insert(key.clone(), self.metrics.fetch_duration.timer());
                                 self.fetcher.add_ready(key);
                             } else {
@@ -252,59 +250,36 @@ impl<
                     Message::Cancel { key } => {
                         trace!(?key, "mailbox: cancel");
                         let mut guard = self.metrics.cancel.guard(Status::Dropped);
-                        if self.fetcher.cancel(&key) {
+                        self.fetcher.cancel(&key);
+                        if self.inflight.cancel(&key) {
                             guard.set(Status::Success);
-                            self.fetch_timers.remove(&key).unwrap().cancel(); // must exist, don't record metric
-                            self.consumer.failed(key.clone(), ()).await;
                         }
                     }
                     Message::Retain { predicate } => {
                         trace!("mailbox: retain");
 
-                        // Remove from fetcher
                         self.fetcher.retain(&predicate);
-
-                        // Clean up timers and notify consumer
-                        let before = self.fetch_timers.len();
-                        let removed = self
-                            .fetch_timers
-                            .extract_if(|k, _| !predicate(k))
-                            .collect::<Vec<_>>();
-                        for (key, timer) in removed {
-                            timer.cancel();
-                            self.consumer.failed(key, ()).await;
-                        }
-
-                        // Metrics
-                        let removed = (before - self.fetch_timers.len()) as u64;
-                        if removed == 0 {
-                            self.metrics.cancel.inc(Status::Dropped);
-                        } else {
-                            self.metrics.cancel.inc_by(Status::Success, removed);
-                        }
+                        let count = self.inflight.retain(predicate) as u64;
+                        self.record_cancellations(count);
                     }
                     Message::Clear => {
                         trace!("mailbox: clear");
 
-                        // Clear fetcher
                         self.fetcher.clear();
-
-                        // Drain timers and notify consumer
-                        let removed = self.fetch_timers.len() as u64;
-                        for (key, timer) in self.fetch_timers.drain() {
-                            timer.cancel();
-                            self.consumer.failed(key, ()).await;
-                        }
-
-                        // Metrics
-                        if removed == 0 {
-                            self.metrics.cancel.inc(Status::Dropped);
-                        } else {
-                            self.metrics.cancel.inc_by(Status::Success, removed);
-                        }
+                        let count = self.inflight.drain() as u64;
+                        self.record_cancellations(count);
                     }
                 }
-                assert_eq!(self.fetcher.len(), self.fetch_timers.len());
+            },
+            // Handle completed consumer deliveries
+            delivery = self.inflight.next_delivery() => {
+                // If the delivery was aborted, its inflight entry was dropped (via
+                // Cancel, Retain, Clear, or shutdown) before the consumer finished validating.
+                let (peer, key, valid) = match delivery {
+                    Ok(delivery) => delivery,
+                    Err(_) => continue,
+                };
+                self.handle_delivery(peer, key, valid).await;
             },
             // Handle completed server requests
             serve = self.serves.next_completed() => {
@@ -352,15 +327,23 @@ impl<
                 };
                 match msg.payload {
                     wire::Payload::Request(key) => self.handle_network_request(peer, msg.id, key),
-                    wire::Payload::Response(response) => {
-                        self.handle_network_response(peer, msg.id, response).await
-                    }
+                    wire::Payload::Response(response) =>
+                        self.handle_network_response(peer, msg.id, response),
                     wire::Payload::Error => self.handle_network_error_response(peer, msg.id),
                 };
             },
         }
     }
 
+    /// Record cancellation metrics for a retain-style operation.
+    fn record_cancellations(&mut self, count: u64) {
+        if count == 0 {
+            self.metrics.cancel.inc(Status::Dropped);
+        } else {
+            self.metrics.cancel.inc_by(Status::Success, count);
+        }
+    }
+
     /// Handles the case where the application responds to a request from an external peer.
     async fn handle_serve(
         &mut self,
@@ -409,7 +392,7 @@ impl<
     }
 
     /// Handle a network response from a peer.
-    async fn handle_network_response(&mut self, peer: P, id: u64, response: Bytes) {
+    fn handle_network_response(&mut self, peer: P, id: u64, response: Bytes) {
         trace!(?peer, ?id, "peer response: data");
 
         // Get the key associated with the response, if any
@@ -418,13 +401,15 @@ impl<
             return;
         };
 
-        // The peer had the data, so we can deliver it to the consumer
-        if self.consumer.deliver(key.clone(), response).await {
-            // Record metrics
-            self.metrics.fetch.inc(Status::Success);
-            self.fetch_timers.remove(&key).unwrap(); // must exist in the map, records metric on drop
+        // The peer had the data, so deliver it to the consumer without blocking the engine.
+        self.inflight.deliver(key, peer, response);
+    }
 
-            // Clear all targets for this key
+    /// Handle completed delivery to the consumer.
+    async fn handle_delivery(&mut self, peer: P, key: Key, valid: bool) {
+        if valid {
+            self.metrics.fetch.inc(Status::Success);
+            self.inflight.complete(&key); // records duration on drop
             self.fetcher.clear_targets(&key);
             return;
         }

resolver/src/p2p/fetcher.rs

@@ -534,6 +534,7 @@ where
     }
 
     /// Returns the number of fetches.
+    #[cfg(test)]
     pub fn len(&self) -> usize {
         self.pending.len() + self.requests.len()
     }