fix(dash-spv): route header requests to the shortest queue

dash-spv/src/network/broker.rs

@@ -236,7 +236,7 @@ impl Broker {
     }
 
     /// Sends what fits in both the peers' and the global byte caps, each message to
-    /// the peer least loaded relative to its own cap.
+    /// the peer [`choose_peer`] picks.
     pub async fn route(
         &self,
         peers: &[PeerHandle],
@@ -271,12 +271,8 @@ impl Broker {
         let mut on_wire: Vec<(u32, NetworkMessage, PeerHandle)> = Vec::new();
         let mut msgs = msgs.into_iter();
         for (height, msg) in msgs.by_ref() {
-            let Some(peer) =
-                peers.iter().filter(|p| p.in_flight_bytes() < p.cap()).min_by(|a, b| {
-                    (a.in_flight_bytes() as u128 * b.cap() as u128)
-                        .cmp(&(b.in_flight_bytes() as u128 * a.cap() as u128))
-                })
-            else {
+            let open = peers.iter().filter(|p| p.in_flight_bytes() < p.cap());
+            let Some(peer) = choose_peer(open, classify(&msg)) else {
                 unsent.push((height, msg));
                 break;
             };
@@ -425,6 +421,36 @@ fn classify(msg: &NetworkMessage) -> MsgClass {
     }
 }
 
+/// Picks the peer among `open` for a message of `class`. A header request goes to the
+/// peer that will answer it soonest. Anything else goes to the peer least loaded relative
+/// to its own cap, and so does a header request while a busy peer has no serving rate yet.
+fn choose_peer<'a>(
+    open: impl Iterator<Item = &'a PeerHandle> + Clone,
+    class: MsgClass,
+) -> Option<&'a PeerHandle> {
+    if class == MsgClass::Headers {
+        let waits: Option<Vec<(f64, &PeerHandle)>> =
+            open.clone().map(|p| Some((queued_secs(p)?, p))).collect();
+        if let Some(waits) = waits {
+            return waits.into_iter().min_by(|a, b| a.0.total_cmp(&b.0)).map(|(_, p)| p);
+        }
+    }
+    open.min_by(|a, b| {
+        (a.in_flight_bytes() as u128 * b.cap() as u128)
+            .cmp(&(b.in_flight_bytes() as u128 * a.cap() as u128))
+    })
+}
+
+/// Seconds the peer needs to serve what it already holds, `None` while it holds work
+/// but has no measured serving rate. Peers answer in order, so a new request waits that long.
+fn queued_secs(peer: &PeerHandle) -> Option<f64> {
+    match (peer.in_flight_bytes(), peer.serving_rate()) {
+        (0, _) => Some(0.0),
+        (_, 0) => None,
+        (bytes, rate) => Some(bytes as f64 / rate as f64),
+    }
+}
+
 /// `(height, seq)`: lowest height first, then FIFO; front inserts take negative seqs.
 type QueueKey = (u32, i64);
 
@@ -545,6 +571,7 @@ impl MsgQueue {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::network::peer::ConnectedPeer;
     use crate::test_utils::test_socket_address;
     use dashcore::network::message_blockdata::GetHeadersMessage;
     use dashcore::network::message_filter::{GetCFHeaders, GetCFilters};
@@ -709,6 +736,37 @@ mod tests {
         assert_eq!(broker.requests.lock().await.len(), tracked);
     }
 
+    /// `slow` holds 1 MB at 100 KB/s, a 10 s queue at a tenth of its cap. `fast` holds
+    /// 256 KB at `fast_rate`, half its cap. Returns which of them `msg` is routed to.
+    #[test_case(get_headers(1), 1 << 20 => "fast"; "a header request goes to the shortest queue")]
+    #[test_case(get_headers(1), 0 => "slow"; "a header request falls back to the cap ratio while a busy peer is unmeasured")]
+    #[test_case(get_cfilters(0), 1 << 20 => "slow"; "bulk requests keep the cap ratio")]
+    #[tokio::test]
+    async fn route_picks_the_peer_by_class(msg: NetworkMessage, fast_rate: u64) -> &'static str {
+        let costs = ResponseCosts::new();
+        let slow = ConnectedPeer::for_test(conn(1, 0)).await;
+        let fast = ConnectedPeer::for_test(conn(2, 0)).await;
+        for start in [0, 1000, 2000, 3000] {
+            slow.handle().send(&get_cfilters(start)).await.expect("loopback send");
+        }
+        fast.handle().send(&get_cfilters(4000)).await.expect("loopback send");
+        slow.set_cap(slow.in_flight_bytes() * 10);
+        fast.set_cap(fast.in_flight_bytes() * 2);
+        slow.set_serving_rate(100 << 10);
+        fast.set_serving_rate(fast_rate);
+        let before = (slow.in_flight_bytes(), fast.in_flight_bytes());
+
+        let broker = Broker::new();
+        broker.enqueue(0, msg).await;
+        assert_eq!(broker.route(&[slow.handle(), fast.handle()], u64::MAX, &costs).await, 1);
+
+        match (slow.in_flight_bytes() > before.0, fast.in_flight_bytes() > before.1) {
+            (true, false) => "slow",
+            (false, true) => "fast",
+            other => panic!("routed to {other:?}"),
+        }
+    }
+
     #[test_case(get_headers(1) => MsgClass::Headers)]
     #[test_case(get_cfheaders(1) => MsgClass::CfHeaders)]
     #[test_case(get_cfilters(0) => MsgClass::CFilters)]

dash-spv/src/network/peer.rs

@@ -21,7 +21,7 @@ use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 #[cfg(test)]
-use tokio::net::TcpListener;
+use tokio::{io, net::TcpListener};
 use tokio::{
     io::{AsyncRead, AsyncWriteExt, ReadBuf},
     net::{
@@ -257,6 +257,10 @@ impl PeerHandle {
         self.id
     }
 
+    pub(super) fn serving_rate(&self) -> u64 {
+        self.serving_rate.load(Ordering::Relaxed)
+    }
+
     pub fn in_flight_bytes(&self) -> u64 {
         self.ledger.in_flight_bytes()
     }
@@ -357,19 +361,22 @@ impl ConnectedPeer {
     }
 
     pub fn serving_rate(&self) -> u64 {
-        self.handle.serving_rate.load(Ordering::Relaxed)
+        self.handle.serving_rate()
     }
 
     pub fn set_serving_rate(&self, bytes_per_sec: u64) {
         self.handle.serving_rate.store(bytes_per_sec, Ordering::Relaxed);
     }
 
     /// A handshaken connection `id` over a loopback socket, with no reader running.
+    /// The far end discards what it reads, so sends succeed.
     #[cfg(test)]
     pub(crate) async fn for_test(id: PeerId) -> Self {
         let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind loopback");
         let addr = listener.local_addr().expect("bound address");
         let stream = TcpStream::connect(addr).await.expect("connect loopback");
+        let (mut far, _) = listener.accept().await.expect("accept loopback");
+        tokio::spawn(async move { io::copy(&mut far, &mut io::sink()).await });
         let link = Link::v1(stream, Network::Regtest.magic());
         ConnectedPeer {
             handle: PeerHandle {