From 6fd280bf6202383d63c34084afbfcdaf8e9cd200 Mon Sep 17 00:00:00 2001
From: Floris Bruynooghe <flub@n0.computer>
Date: Tue, 11 Feb 2025 17:30:15 +0100
Subject: [PATCH] Kind of do tail-loss probes correctly

This is a bit sad because it splits the use in poll_transmit.  It will
need more adjustments once poll_transmit starts sending on multiple
paths though.
---
 quinn-proto/src/connection/mod.rs    | 20 ++++++++++++++------
 quinn-proto/src/connection/spaces.rs |  4 +++-
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
index 1cb89a8d7..a6f536099 100644
--- a/quinn-proto/src/connection/mod.rs
+++ b/quinn-proto/src/connection/mod.rs
@@ -1,6 +1,6 @@
 use std::{
     cmp,
-    collections::VecDeque,
+    collections::{BTreeSet, VecDeque},
     convert::TryFrom,
     fmt, io, mem,
     net::{IpAddr, SocketAddr},
@@ -525,11 +525,12 @@ impl Connection {
         }
 
         // If we need to send a probe, make sure we have something to send.
-        for space in SpaceId::iter() {
-            let request_immediate_ack =
-                space == SpaceId::Data && self.peer_supports_ack_frequency();
-            self.spaces[space].maybe_queue_probe(request_immediate_ack, &self.streams);
-        }
+        self.spaces[SpaceId::Initial].maybe_queue_probe(PathId(0), false, &self.streams);
+        self.spaces[SpaceId::Handshake].maybe_queue_probe(PathId(0), false, &self.streams);
+
+        // For the data paths we need to call maybe_queue_probe once for each path.  This
+        // keeps track if it was already done for a path.
+        let mut data_tail_probes: BTreeSet<PathId> = BTreeSet::new();
 
         // Check whether we need to send a close message
         let close = match self.state {
@@ -581,6 +582,13 @@ impl Connection {
         while space_idx < spaces.len() {
             let space_id = spaces[space_idx];
 
+            // If we need to send a tail-loss probe, make sure there is something to send.
+            if space_id == SpaceId::Data && !data_tail_probes.contains(&path_id) {
+                let immediate_ack = self.peer_supports_ack_frequency();
+                self.spaces[space_id].maybe_queue_probe(path_id, immediate_ack, &self.streams);
+                data_tail_probes.insert(path_id);
+            }
+
             // Number of bytes available for frames if this is a 1-RTT packet. We're guaranteed to
             // be able to send an individual frame at least this large in the next 1-RTT
             // packet. This could be generalized to support every space, but it's only needed to
diff --git a/quinn-proto/src/connection/spaces.rs b/quinn-proto/src/connection/spaces.rs
index f61bfee92..53ed20a25 100644
--- a/quinn-proto/src/connection/spaces.rs
+++ b/quinn-proto/src/connection/spaces.rs
@@ -114,10 +114,11 @@ impl PacketSpace {
     //    might be lost.
     pub(super) fn maybe_queue_probe(
         &mut self,
+        path_id: PathId,
         request_immediate_ack: bool,
         streams: &StreamsState,
     ) {
-        if self.number_spaces.values().all(|s| s.loss_probes == 0) {
+        if self.for_path(path_id).loss_probes == 0 {
             return;
         }
 
@@ -133,6 +134,7 @@ impl PacketSpace {
             return;
         }
 
+        // We use retransmits from any path.
         for packet in self
             .number_spaces
             .values_mut()