diff --git a/Cargo.lock b/Cargo.lock index 9cb37b518..f3e3a8f10 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1443,6 +1443,7 @@ dependencies = [ "rustls-pki-types", "rustls-platform-verifier", "slab", + "sorted-index-buffer", "thiserror 2.0.17", "tinyvec", "tracing", @@ -2396,6 +2397,12 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "sorted-index-buffer" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d96a9d278ef78c991bf2faabd578b41199037846ff96fbb093f84bb50adaa00" + [[package]] name = "stable_deref_trait" version = "1.2.0" diff --git a/Cargo.toml b/Cargo.toml index e7c8d109d..4ffbe75d0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,6 +45,7 @@ serde_json = "1" slab = "0.4.9" smol = "2" socket2 = ">=0.5, <0.7" +sorted-index-buffer = { version = "0.2.0" } thiserror = "2.0.3" tinyvec = { version = "1.1", features = ["alloc"] } tokio = { version = "1.28.1", features = ["sync"] } diff --git a/quinn-proto/Cargo.toml b/quinn-proto/Cargo.toml index 0fb36ab5c..1e1cec4ee 100644 --- a/quinn-proto/Cargo.toml +++ b/quinn-proto/Cargo.toml @@ -46,6 +46,7 @@ arbitrary = { workspace = true, optional = true } aws-lc-rs = { workspace = true, optional = true } bytes = { workspace = true } fastbloom = { workspace = true, optional = true } +identity-hash = { workspace = true } lru-slab = { workspace = true } qlog = { workspace = true, optional = true } rustc-hash = { workspace = true } @@ -54,10 +55,10 @@ ring = { workspace = true, optional = true } rustls = { workspace = true, optional = true } rustls-platform-verifier = { workspace = true, optional = true } slab = { workspace = true } +sorted-index-buffer = { workspace = true } thiserror = { workspace = true } tinyvec = { workspace = true, features = ["alloc"] } tracing = { workspace = true } -identity-hash = { workspace = true } # Feature flags & dependencies for wasm # wasm-bindgen is assumed for a wasm*-*-unknown target diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs index dfcb8e43d..111799074 100644 --- a/quinn-proto/src/connection/mod.rs +++ b/quinn-proto/src/connection/mod.rs @@ -2343,7 +2343,7 @@ impl Connection { let space = &mut self.spaces[space].for_path(path); if space.largest_acked_packet.is_none_or(|pn| ack.largest > pn) { space.largest_acked_packet = Some(ack.largest); - if let Some(info) = space.sent_packets.get(&ack.largest) { + if let Some(info) = space.sent_packets.get(ack.largest) { // This should always succeed, but a misbehaving peer might ACK a packet we // haven't sent. At worst, that will result in us spuriously reducing the // congestion window. @@ -2365,7 +2365,11 @@ impl Connection { let mut newly_acked = ArrayRangeSet::new(); for range in ack.iter() { self.spaces[space].for_path(path).check_ack(range.clone())?; - for (&pn, _) in self.spaces[space].for_path(path).sent_packets.range(range) { + for (pn, _) in self.spaces[space] + .for_path(path) + .sent_packets + .iter_range(range) + { newly_acked.insert_one(pn); } } @@ -2478,13 +2482,12 @@ impl Connection { for range in ack.iter() { let spurious_losses: Vec = lost_packets - .range(range.clone()) + .iter_range(range.clone()) .map(|(pn, _info)| pn) - .copied() .collect(); for pn in spurious_losses { - lost_packets.remove(&pn); + lost_packets.remove(pn); } } @@ -2700,7 +2703,7 @@ impl Connection { let mut prev_packet = None; let space = self.spaces[pn_space].for_path(path_id); - for (&packet, info) in space.sent_packets.range(0..largest_acked_packet) { + for (packet, info) in space.sent_packets.iter_range(0..largest_acked_packet) { if prev_packet != Some(packet.wrapping_sub(1)) { // An intervening packet was acknowledged persistent_congestion_start = None; @@ -2772,10 +2775,10 @@ impl Connection { .for_path(path_id) .sent_packets .iter() - .filter(|(pn, _info)| Some(**pn) != in_flight_mtu_probe) + .filter(|(pn, _info)| Some(*pn) != in_flight_mtu_probe) .map(|(pn, info)| { size_of_lost_packets += info.size as u64; - *pn + pn }) .collect(); @@ -2835,8 +2838,12 @@ impl Connection { // OnPacketsLost if let Some(largest_lost) = lost_packets.last().cloned() { let old_bytes_in_flight = self.path_data_mut(path_id).in_flight.bytes; - let largest_lost_sent = - self.spaces[pn_space].for_path(path_id).sent_packets[&largest_lost].time_sent; + let largest_lost_sent = self.spaces[pn_space] + .for_path(path_id) + .sent_packets + .get(largest_lost) + .unwrap() + .time_sent; let path_stats = self.path_stats.entry(path_id).or_default(); path_stats.lost_packets += lost_packets.len() as u64; path_stats.lost_bytes += size_of_lost_packets; @@ -3463,7 +3470,7 @@ impl Connection { pns.loss_probes = 0; let sent_packets = mem::take(&mut pns.sent_packets); let path = self.paths.get_mut(&PathId::ZERO).unwrap(); - for packet in sent_packets.into_values() { + for (_, packet) in sent_packets.into_iter() { path.data.remove_in_flight(&packet); } @@ -3866,7 +3873,7 @@ impl Connection { .for_path(PathId::ZERO) .sent_packets, ); - for info in zero_rtt.into_values() { + for (_, info) in zero_rtt.into_iter() { self.paths .get_mut(&PathId::ZERO) .unwrap() @@ -3938,7 +3945,7 @@ impl Connection { let sent_packets = mem::take( &mut self.spaces[SpaceId::Data].for_path(path_id).sent_packets, ); - for packet in sent_packets.into_values() { + for (_, packet) in sent_packets.into_iter() { self.paths .get_mut(&path_id) .unwrap() diff --git a/quinn-proto/src/connection/spaces.rs b/quinn-proto/src/connection/spaces.rs index 89ea2509e..a41ae1e40 100644 --- a/quinn-proto/src/connection/spaces.rs +++ b/quinn-proto/src/connection/spaces.rs @@ -6,6 +6,7 @@ use std::{ ops::{Bound, Index, IndexMut}, }; +use sorted_index_buffer::SortedIndexBuffer; use rand::Rng; use rustc_hash::{FxHashMap, FxHashSet}; use tracing::{error, trace}; @@ -221,10 +222,10 @@ pub(super) struct PacketNumberSpace { pub(super) unacked_non_ack_eliciting_tail: u64, /// Transmitted but not acked // We use a BTreeMap here so we can efficiently query by range on ACK and for loss detection - pub(super) sent_packets: BTreeMap, + pub(super) sent_packets: SortedIndexBuffer, /// Packets that were deemed lost // Older packets are regularly removed in `Connection::drain_lost_packets`. - pub(super) lost_packets: BTreeMap, + pub(super) lost_packets: SortedIndexBuffer, /// Number of explicit congestion notification codepoints seen on incoming packets pub(super) ecn_counters: frame::EcnCounts, /// Recent ECN counters sent by the peer in ACK frames @@ -276,8 +277,8 @@ impl PacketNumberSpace { largest_acked_packet_sent: now, largest_ack_eliciting_sent: 0, unacked_non_ack_eliciting_tail: 0, - sent_packets: BTreeMap::new(), - lost_packets: BTreeMap::new(), + sent_packets: SortedIndexBuffer::new(), + lost_packets: SortedIndexBuffer::new(), ecn_counters: frame::EcnCounts::ZERO, ecn_feedback: frame::EcnCounts::ZERO, sent_with_keys: 0, @@ -305,8 +306,8 @@ impl PacketNumberSpace { largest_acked_packet_sent: now, largest_ack_eliciting_sent: 0, unacked_non_ack_eliciting_tail: 0, - sent_packets: BTreeMap::new(), - lost_packets: BTreeMap::new(), + sent_packets: SortedIndexBuffer::new(), + lost_packets: SortedIndexBuffer::new(), ecn_counters: frame::EcnCounts::ZERO, ecn_feedback: frame::EcnCounts::ZERO, sent_with_keys: 0, @@ -335,8 +336,8 @@ impl PacketNumberSpace { largest_acked_packet_sent: Instant::now(), largest_ack_eliciting_sent: 0, unacked_non_ack_eliciting_tail: 0, - sent_packets: BTreeMap::new(), - lost_packets: BTreeMap::new(), + sent_packets: SortedIndexBuffer::new(), + lost_packets: SortedIndexBuffer::new(), ecn_counters: frame::EcnCounts::ZERO, ecn_feedback: frame::EcnCounts::ZERO, sent_with_keys: 0, @@ -434,7 +435,7 @@ impl PacketNumberSpace { /// Stop tracking sent packet `number`, and return what we knew about it pub(super) fn take(&mut self, number: u64) -> Option { - let packet = self.sent_packets.remove(&number)?; + let packet = self.sent_packets.remove(number)?; if !packet.ack_eliciting && number > self.largest_ack_eliciting_sent { self.unacked_non_ack_eliciting_tail = self.unacked_non_ack_eliciting_tail.checked_sub(1).unwrap(); @@ -457,22 +458,21 @@ impl PacketNumberSpace { self.unacked_non_ack_eliciting_tail = 0; self.largest_ack_eliciting_sent = number; } else if self.unacked_non_ack_eliciting_tail > MAX_UNACKED_NON_ACK_ELICTING_TAIL { - let oldest_after_ack_eliciting = *self + let oldest_after_ack_eliciting = self .sent_packets - .range(( + .keys_range(( Bound::Excluded(self.largest_ack_eliciting_sent), Bound::Unbounded, )) .next() - .unwrap() - .0; + .unwrap(); // Per https://www.rfc-editor.org/rfc/rfc9000.html#name-frames-and-frame-types, // non-ACK-eliciting packets must only contain PADDING, ACK, and CONNECTION_CLOSE // frames, which require no special handling on ACK or loss beyond removal from // in-flight counters if padded. let packet = self .sent_packets - .remove(&oldest_after_ack_eliciting) + .remove(oldest_after_ack_eliciting) .unwrap(); debug_assert!(!packet.ack_eliciting); forgotten = Some(packet);