Merge branch 'main-iroh' into fix-send-buffer

quinn-proto/src/connection/assembler.rs

@@ -6,7 +6,7 @@ use std::{
 
 use bytes::{Buf, Bytes, BytesMut};
 
-use crate::range_set::RangeSet;
+use crate::range_set::ArrayRangeSet;
 
 /// Helper to assemble unordered stream frames into an ordered stream
 #[derive(Debug, Default)]
@@ -46,7 +46,7 @@ impl Assembler {
                 // Get rid of possible duplicates
                 self.defragment();
             }
-            let mut recvd = RangeSet::new();
+            let mut recvd = ArrayRangeSet::new();
             recvd.insert(0..self.bytes_read);
             for chunk in &self.data {
                 recvd.insert(chunk.offset..chunk.offset + chunk.bytes.len() as u64);
@@ -108,7 +108,7 @@ impl Assembler {
         let mut buffers = old.into_sorted_vec();
         self.buffered = 0;
         let mut fragmented_buffered = 0;
-        let mut offset = 0;
+        let mut offset = self.bytes_read;
         for chunk in buffers.iter_mut().rev() {
             chunk.try_mark_defragment(offset);
             let size = chunk.bytes.len();
@@ -120,7 +120,7 @@ impl Assembler {
         }
         self.allocated = self.buffered;
         let mut buffer = BytesMut::with_capacity(fragmented_buffered);
-        let mut offset = 0;
+        let mut offset = self.bytes_read;
         for chunk in buffers.into_iter().rev() {
             if chunk.defragmented {
                 // bytes might be empty after try_mark_defragment
@@ -157,7 +157,8 @@ impl Assembler {
         self.end = self.end.max(offset + bytes.len() as u64);
         if let State::Unordered { ref mut recvd } = self.state {
             // Discard duplicate data
-            for duplicate in recvd.replace(offset..offset + bytes.len() as u64) {
+            let range = offset..offset + bytes.len() as u64;
+            for duplicate in recvd.iter_range(range.clone()) {
                 if duplicate.start > offset {
                     let buffer = Buffer::new(
                         offset,
@@ -172,6 +173,7 @@ impl Assembler {
                 bytes.advance((duplicate.end - offset) as usize);
                 offset = duplicate.end;
             }
+            recvd.insert(range);
         } else if offset < self.bytes_read {
             if (offset + bytes.len() as u64) <= self.bytes_read {
                 return;
@@ -321,7 +323,7 @@ enum State {
     Unordered {
         /// The set of offsets that have been received from the peer, including portions not yet
         /// read by the application.
-        recvd: RangeSet,
+        recvd: ArrayRangeSet,
     },
 }
 
@@ -648,6 +650,20 @@ mod test {
         assert_eq!(x.read(3, false), None);
     }
 
+    #[test]
+    fn no_duplicate_after_mode_switch() {
+        // Regression test: bytes read in ordered mode should not be returned again in unordered mode
+        let mut x = Assembler::new();
+        x.insert(0, Bytes::from_static(b"a"), 1);
+        x.insert(0, Bytes::from_static(b"a"), 1); // duplicate
+        assert_eq!(
+            x.read(1, true),
+            Some(Chunk::new(0, Bytes::from_static(b"a")))
+        );
+        x.ensure_ordering(false).unwrap();
+        assert_eq!(x.read(1, false), None); // should be None, byte 0 already returned
+    }
+
     fn next_unordered(x: &mut Assembler) -> Chunk {
         x.read(usize::MAX, false).unwrap()
     }
@@ -656,3 +672,174 @@ mod test {
         x.read(size, true).map(|chunk| chunk.bytes)
     }
 }
+
+#[cfg(all(test, not(target_family = "wasm")))]
+mod proptests {
+    use proptest::prelude::*;
+    use test_strategy::{Arbitrary, proptest};
+
+    use super::*;
+
+    const MAX_OFFSET: u64 = 512;
+    const MAX_LEN: usize = 64;
+
+    #[derive(Debug, Clone, Arbitrary)]
+    enum Op {
+        #[weight(10)]
+        Insert {
+            #[strategy(0..MAX_OFFSET)]
+            offset: u64,
+            #[strategy(1..MAX_LEN)]
+            len: usize,
+        },
+        #[weight(10)]
+        Read {
+            #[strategy(1..MAX_LEN)]
+            max_len: usize,
+        },
+        #[weight(1)]
+        EnsureOrdering { ordered: bool },
+        #[weight(1)]
+        Defragment,
+    }
+
+    /// Tracks the state of the assembler for verification
+    struct RefState {
+        received: Vec<bool>,
+        returned: Vec<bool>,
+        ordered: bool,
+    }
+
+    fn set_range(bits: &mut [bool], start: u64, len: usize) {
+        for i in start..(start + len as u64).min(bits.len() as u64) {
+            bits[i as usize] = true;
+        }
+    }
+
+    impl RefState {
+        fn new() -> Self {
+            Self {
+                received: vec![false; MAX_OFFSET as usize],
+                returned: vec![false; MAX_OFFSET as usize],
+                ordered: true,
+            }
+        }
+
+        fn insert(&mut self, offset: u64, len: usize) {
+            set_range(&mut self.received, offset, len);
+        }
+
+        fn ensure_ordering(&mut self, ordered: bool) -> bool {
+            if ordered && !self.ordered {
+                return false;
+            }
+            self.ordered = ordered;
+            true
+        }
+
+        fn bytes_read(&self) -> u64 {
+            self.returned.iter().filter(|&&x| x).count() as u64
+        }
+    }
+
+    fn make_data() -> Vec<u8> {
+        use rand::{Rng, SeedableRng};
+        let mut rng = rand::rngs::StdRng::seed_from_u64(0xDEADBEEF);
+        let mut data = vec![0u8; MAX_OFFSET as usize];
+        rng.fill(data.as_mut_slice());
+        data
+    }
+
+    fn get_slice(data: &[u8], offset: u64, len: usize) -> Bytes {
+        let start = offset as usize;
+        let end = (start + len).min(data.len());
+        Bytes::copy_from_slice(&data[start..end])
+    }
+
+    fn verify_chunk(data: &[u8], chunk: &Chunk) -> bool {
+        let start = chunk.offset as usize;
+        chunk.bytes[..] == data[start..start + chunk.bytes.len()]
+    }
+
+    #[proptest]
+    fn assembler_matches_reference(
+        #[strategy(proptest::collection::vec(any::<Op>(), 1..100))] ops: Vec<Op>,
+    ) {
+        let data = make_data();
+        let mut asm = Assembler::new();
+        let mut reference = RefState::new();
+
+        for op in ops {
+            match op {
+                Op::Insert { offset, len } => {
+                    let bytes = get_slice(&data, offset, len);
+                    asm.insert(offset, bytes, len);
+                    reference.insert(offset, len);
+                }
+                Op::Read { max_len } => {
+                    let ordered = reference.ordered;
+                    let actual = asm.read(max_len, ordered);
+
+                    match actual {
+                        None => {
+                            // Should only be None if no unreturned received bytes available
+                            let has_available = if ordered {
+                                // In ordered mode, check if the first unreturned byte is received
+                                reference
+                                    .returned
+                                    .iter()
+                                    .position(|&x| !x)
+                                    .is_some_and(|pos| reference.received[pos])
+                            } else {
+                                // In unordered mode, check if any unreturned received byte exists
+                                reference
+                                    .received
+                                    .iter()
+                                    .zip(&reference.returned)
+                                    .any(|(&r, &ret)| r && !ret)
+                            };
+                            prop_assert!(
+                                !has_available,
+                                "read returned None but data was available"
+                            );
+                        }
+                        Some(chunk) => {
+                            prop_assert!(chunk.bytes.len() <= max_len, "chunk exceeds max_len");
+                            prop_assert!(verify_chunk(&data, &chunk), "data corruption");
+                            // Mark as returned, check for duplicates
+                            for i in 0..chunk.bytes.len() {
+                                let offset = chunk.offset as usize + i;
+                                prop_assert!(
+                                    reference.received[offset],
+                                    "returned unreceived byte at {offset}"
+                                );
+                                prop_assert!(
+                                    !reference.returned[offset],
+                                    "duplicate byte at {offset}"
+                                );
+                                reference.returned[offset] = true;
+                            }
+                        }
+                    }
+                }
+                Op::EnsureOrdering { ordered } => {
+                    let actual = asm.ensure_ordering(ordered).is_ok();
+                    let expected = reference.ensure_ordering(ordered);
+                    prop_assert_eq!(actual, expected, "ensure_ordering result mismatch");
+                }
+                Op::Defragment => {
+                    if asm.state.is_ordered() {
+                        asm.defragment();
+                    }
+                }
+            }
+        }
+
+        // Invariant: bytes_read matches
+        prop_assert_eq!(
+            asm.bytes_read(),
+            reference.bytes_read(),
+            "bytes_read mismatch"
+        );
+    }
+}

quinn-proto/src/connection/send_buffer.rs

@@ -2,7 +2,7 @@ use std::{collections::VecDeque, ops::Range};
 
 use bytes::{Buf, BufMut, Bytes};
 
-use crate::{VarInt, range_set::RangeSet};
+use crate::{VarInt, range_set::ArrayRangeSet};
 
 /// Buffer of outgoing retransmittable stream data
 #[derive(Default, Debug)]
@@ -25,14 +25,14 @@ pub(super) struct SendBuffer {
     /// All ranges must be within `data.range().start..(data.range().end - unsent)`, since data
     /// that has never been sent can't be acknowledged.
     // TODO: Recover storage from these by compacting (#700)
-    acks: RangeSet,
+    acks: ArrayRangeSet,
     /// Previously transmitted ranges deemed lost and marked for retransmission
     ///
     /// All ranges must be within `data.range().start..(data.range().end - unsent)`, since data
     /// that has never been sent can't be retransmitted.
     ///
     /// This should usually ot overlap with `acks`, but this is not strictly enforced.
-    retransmits: RangeSet,
+    retransmits: ArrayRangeSet,
 }
 
 /// This is where the data of the send buffer lives. It supports appending at the end,

quinn-proto/src/range_set/array_range_set.rs

@@ -68,13 +68,14 @@ impl ArrayRangeSet {
         false
     }
 
-    pub fn subtract(&mut self, other: &Self) {
-        // TODO: This can potentially be made more efficient, since the we know
-        // individual ranges are not overlapping, and the next range must start
-        // after the last one finished
-        for range in &other.0 {
-            self.remove(range.clone());
-        }
+    pub fn iter_range(&self, range: Range<u64>) -> impl Iterator<Item = Range<u64>> + '_ {
+        self.iter().filter_map(move |r| {
+            if r.end > range.start && r.start < range.end {
+                Some(r.start.max(range.start)..r.end.min(range.end))
+            } else {
+                None
+            }
+        })
     }
 
     pub fn insert_one(&mut self, x: u64) -> bool {

quinn-proto/src/range_set/btree_range_set.rs

@@ -1,381 +0,0 @@
-use std::{
-    cmp,
-    cmp::Ordering,
-    collections::{BTreeMap, btree_map},
-    ops::{
-        Bound::{Excluded, Included},
-        Range,
-    },
-};
-
-/// A set of u64 values optimized for long runs and random insert/delete/contains
-#[derive(Debug, Default, Clone)]
-pub struct RangeSet(BTreeMap<u64, u64>);
-
-impl RangeSet {
-    pub fn new() -> Self {
-        Default::default()
-    }
-
-    pub fn contains(&self, x: u64) -> bool {
-        self.pred(x).is_some_and(|(_, end)| end > x)
-    }
-
-    pub fn insert_one(&mut self, x: u64) -> bool {
-        if let Some((start, end)) = self.pred(x) {
-            match end.cmp(&x) {
-                // Wholly contained
-                Ordering::Greater => {
-                    return false;
-                }
-                Ordering::Equal => {
-                    // Extend existing
-                    self.0.remove(&start);
-                    let mut new_end = x + 1;
-                    if let Some((next_start, next_end)) = self.succ(x) {
-                        if next_start == new_end {
-                            self.0.remove(&next_start);
-                            new_end = next_end;
-                        }
-                    }
-                    self.0.insert(start, new_end);
-                    return true;
-                }
-                _ => {}
-            }
-        }
-        let mut new_end = x + 1;
-        if let Some((next_start, next_end)) = self.succ(x) {
-            if next_start == new_end {
-                self.0.remove(&next_start);
-                new_end = next_end;
-            }
-        }
-        self.0.insert(x, new_end);
-        true
-    }
-
-    pub fn insert(&mut self, mut x: Range<u64>) -> bool {
-        if x.is_empty() {
-            return false;
-        }
-        if let Some((start, end)) = self.pred(x.start) {
-            if end >= x.end {
-                // Wholly contained
-                return false;
-            } else if end >= x.start {
-                // Extend overlapping predecessor
-                self.0.remove(&start);
-                x.start = start;
-            }
-        }
-        while let Some((next_start, next_end)) = self.succ(x.start) {
-            if next_start > x.end {
-                break;
-            }
-            // Overlaps with successor
-            self.0.remove(&next_start);
-            x.end = cmp::max(next_end, x.end);
-        }
-        self.0.insert(x.start, x.end);
-        true
-    }
-
-    /// Find closest range to `x` that begins at or before it
-    fn pred(&self, x: u64) -> Option<(u64, u64)> {
-        self.0
-            .range((Included(0), Included(x)))
-            .next_back()
-            .map(|(&x, &y)| (x, y))
-    }
-
-    /// Find the closest range to `x` that begins after it
-    fn succ(&self, x: u64) -> Option<(u64, u64)> {
-        self.0
-            .range((Excluded(x), Included(u64::MAX)))
-            .next()
-            .map(|(&x, &y)| (x, y))
-    }
-
-    pub fn remove(&mut self, x: Range<u64>) -> bool {
-        if x.is_empty() {
-            return false;
-        }
-
-        let before = match self.pred(x.start) {
-            Some((start, end)) if end > x.start => {
-                self.0.remove(&start);
-                if start < x.start {
-                    self.0.insert(start, x.start);
-                }
-                if end > x.end {
-                    self.0.insert(x.end, end);
-                }
-                // Short-circuit if we cannot possibly overlap with another range
-                if end >= x.end {
-                    return true;
-                }
-                true
-            }
-            Some(_) | None => false,
-        };
-        let mut after = false;
-        while let Some((start, end)) = self.succ(x.start) {
-            if start >= x.end {
-                break;
-            }
-            after = true;
-            self.0.remove(&start);
-            if end > x.end {
-                self.0.insert(x.end, end);
-                break;
-            }
-        }
-        before || after
-    }
-
-    /// Add a range to the set, returning the intersection of current ranges with the new one
-    pub fn replace(&mut self, mut range: Range<u64>) -> Replace<'_> {
-        let pred = if let Some((prev_start, prev_end)) = self
-            .pred(range.start)
-            .filter(|&(_, end)| end >= range.start)
-        {
-            self.0.remove(&prev_start);
-            let replaced_start = range.start;
-            range.start = range.start.min(prev_start);
-            let replaced_end = range.end.min(prev_end);
-            range.end = range.end.max(prev_end);
-            if replaced_start != replaced_end {
-                Some(replaced_start..replaced_end)
-            } else {
-                None
-            }
-        } else {
-            None
-        };
-        Replace {
-            set: self,
-            range,
-            pred,
-        }
-    }
-
-    pub fn add(&mut self, other: &Self) {
-        for (&start, &end) in &other.0 {
-            self.insert(start..end);
-        }
-    }
-
-    pub fn subtract(&mut self, other: &Self) {
-        for (&start, &end) in &other.0 {
-            self.remove(start..end);
-        }
-    }
-
-    pub fn is_empty(&self) -> bool {
-        self.0.is_empty()
-    }
-
-    pub fn min(&self) -> Option<u64> {
-        self.0.first_key_value().map(|(&start, _)| start)
-    }
-
-    pub fn max(&self) -> Option<u64> {
-        self.0.last_key_value().map(|(_, &end)| end - 1)
-    }
-
-    pub fn len(&self) -> usize {
-        self.0.len()
-    }
-    pub fn iter(&self) -> Iter<'_> {
-        Iter(self.0.iter())
-    }
-    pub fn elts(&self) -> EltIter<'_> {
-        EltIter {
-            inner: self.0.iter(),
-            next: 0,
-            end: 0,
-        }
-    }
-
-    pub fn peek_min(&self) -> Option<Range<u64>> {
-        let (&start, &end) = self.0.iter().next()?;
-        Some(start..end)
-    }
-
-    pub fn pop_min(&mut self) -> Option<Range<u64>> {
-        let result = self.peek_min()?;
-        self.0.remove(&result.start);
-        Some(result)
-    }
-}
-
-pub struct Iter<'a>(btree_map::Iter<'a, u64, u64>);
-
-impl Iterator for Iter<'_> {
-    type Item = Range<u64>;
-    fn next(&mut self) -> Option<Range<u64>> {
-        let (&start, &end) = self.0.next()?;
-        Some(start..end)
-    }
-}
-
-impl DoubleEndedIterator for Iter<'_> {
-    fn next_back(&mut self) -> Option<Range<u64>> {
-        let (&start, &end) = self.0.next_back()?;
-        Some(start..end)
-    }
-}
-
-impl<'a> IntoIterator for &'a RangeSet {
-    type Item = Range<u64>;
-    type IntoIter = Iter<'a>;
-    fn into_iter(self) -> Iter<'a> {
-        self.iter()
-    }
-}
-
-pub struct EltIter<'a> {
-    inner: btree_map::Iter<'a, u64, u64>,
-    next: u64,
-    end: u64,
-}
-
-impl Iterator for EltIter<'_> {
-    type Item = u64;
-    fn next(&mut self) -> Option<u64> {
-        if self.next == self.end {
-            let (&start, &end) = self.inner.next()?;
-            self.next = start;
-            self.end = end;
-        }
-        let x = self.next;
-        self.next += 1;
-        Some(x)
-    }
-}
-
-impl DoubleEndedIterator for EltIter<'_> {
-    fn next_back(&mut self) -> Option<u64> {
-        if self.next == self.end {
-            let (&start, &end) = self.inner.next_back()?;
-            self.next = start;
-            self.end = end;
-        }
-        self.end -= 1;
-        Some(self.end)
-    }
-}
-
-/// Iterator returned by `RangeSet::replace`
-pub struct Replace<'a> {
-    set: &'a mut RangeSet,
-    /// Portion of the intersection arising from a range beginning at or before the newly inserted
-    /// range
-    pred: Option<Range<u64>>,
-    /// Union of the input range and all ranges that have been visited by the iterator so far
-    range: Range<u64>,
-}
-
-impl Iterator for Replace<'_> {
-    type Item = Range<u64>;
-    fn next(&mut self) -> Option<Range<u64>> {
-        if let Some(pred) = self.pred.take() {
-            // If a range starting before the inserted range overlapped with it, return the
-            // corresponding overlap first
-            return Some(pred);
-        }
-
-        let (next_start, next_end) = self.set.succ(self.range.start)?;
-        if next_start > self.range.end {
-            // If the next successor range starts after the current range ends, there can be no more
-            // overlaps. This is sound even when `self.range.end` is increased because `RangeSet` is
-            // guaranteed not to contain pairs of ranges that could be simplified.
-            return None;
-        }
-        // Remove the redundant range...
-        self.set.0.remove(&next_start);
-        // ...and handle the case where the redundant range ends later than the new range.
-        let replaced_end = self.range.end.min(next_end);
-        self.range.end = self.range.end.max(next_end);
-        if next_start == replaced_end {
-            // If the redundant range started exactly where the new range ended, there was no
-            // overlap with it or any later range.
-            None
-        } else {
-            Some(next_start..replaced_end)
-        }
-    }
-}
-
-impl Drop for Replace<'_> {
-    fn drop(&mut self) {
-        // Ensure we drain all remaining overlapping ranges
-        for _ in &mut *self {}
-        // Insert the final aggregate range
-        self.set.0.insert(self.range.start, self.range.end);
-    }
-}
-
-/// This module contains tests which only apply for this `RangeSet` implementation
-///
-/// Tests which apply for all implementations can be found in the `tests.rs` module
-#[cfg(test)]
-mod tests {
-    #![allow(clippy::single_range_in_vec_init)] // https://github.com/rust-lang/rust-clippy/issues/11086
-    use super::*;
-
-    #[test]
-    fn replace_contained() {
-        let mut set = RangeSet::new();
-        set.insert(2..4);
-        assert_eq!(set.replace(1..5).collect::<Vec<_>>(), &[2..4]);
-        assert_eq!(set.len(), 1);
-        assert_eq!(set.peek_min().unwrap(), 1..5);
-    }
-
-    #[test]
-    fn replace_contains() {
-        let mut set = RangeSet::new();
-        set.insert(1..5);
-        assert_eq!(set.replace(2..4).collect::<Vec<_>>(), &[2..4]);
-        assert_eq!(set.len(), 1);
-        assert_eq!(set.peek_min().unwrap(), 1..5);
-    }
-
-    #[test]
-    fn replace_pred() {
-        let mut set = RangeSet::new();
-        set.insert(2..4);
-        assert_eq!(set.replace(3..5).collect::<Vec<_>>(), &[3..4]);
-        assert_eq!(set.len(), 1);
-        assert_eq!(set.peek_min().unwrap(), 2..5);
-    }
-
-    #[test]
-    fn replace_succ() {
-        let mut set = RangeSet::new();
-        set.insert(2..4);
-        assert_eq!(set.replace(1..3).collect::<Vec<_>>(), &[2..3]);
-        assert_eq!(set.len(), 1);
-        assert_eq!(set.peek_min().unwrap(), 1..4);
-    }
-
-    #[test]
-    fn replace_exact_pred() {
-        let mut set = RangeSet::new();
-        set.insert(2..4);
-        assert_eq!(set.replace(4..6).collect::<Vec<_>>(), &[]);
-        assert_eq!(set.len(), 1);
-        assert_eq!(set.peek_min().unwrap(), 2..6);
-    }
-
-    #[test]
-    fn replace_exact_succ() {
-        let mut set = RangeSet::new();
-        set.insert(2..4);
-        assert_eq!(set.replace(0..2).collect::<Vec<_>>(), &[]);
-        assert_eq!(set.len(), 1);
-        assert_eq!(set.peek_min().unwrap(), 0..4);
-    }
-}

quinn-proto/src/range_set/mod.rs

@@ -1,7 +1,5 @@
 mod array_range_set;
-mod btree_range_set;
 #[cfg(test)]
 mod tests;
 
 pub(crate) use array_range_set::ArrayRangeSet;
-pub(crate) use btree_range_set::RangeSet;

quinn-proto/src/range_set/tests.rs

@@ -2,194 +2,185 @@ use std::ops::Range;
 
 use super::*;
 
-macro_rules! common_set_tests {
-    ($set_name:ident, $set_type:ident) => {
-        mod $set_name {
-            use super::*;
+mod array_range_set {
+    use super::*;
 
-            #[test]
-            fn merge_and_split() {
-                let mut set = $set_type::new();
-                assert!(set.insert(0..2));
-                assert!(set.insert(2..4));
-                assert!(!set.insert(1..3));
-                assert_eq!(set.len(), 1);
-                assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 3]);
-                assert!(!set.contains(4));
-                assert!(set.remove(2..3));
-                assert_eq!(set.len(), 2);
-                assert!(!set.contains(2));
-                assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 3]);
-            }
+    #[test]
+    fn merge_and_split() {
+        let mut set = ArrayRangeSet::new();
+        assert!(set.insert(0..2));
+        assert!(set.insert(2..4));
+        assert!(!set.insert(1..3));
+        assert_eq!(set.len(), 1);
+        assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 3]);
+        assert!(!set.contains(4));
+        assert!(set.remove(2..3));
+        assert_eq!(set.len(), 2);
+        assert!(!set.contains(2));
+        assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 3]);
+    }
 
-            #[test]
-            fn double_merge_exact() {
-                let mut set = $set_type::new();
-                assert!(set.insert(0..2));
-                assert!(set.insert(4..6));
-                assert_eq!(set.len(), 2);
-                assert!(set.insert(2..4));
-                assert_eq!(set.len(), 1);
-                assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 3, 4, 5]);
-            }
+    #[test]
+    fn double_merge_exact() {
+        let mut set = ArrayRangeSet::new();
+        assert!(set.insert(0..2));
+        assert!(set.insert(4..6));
+        assert_eq!(set.len(), 2);
+        assert!(set.insert(2..4));
+        assert_eq!(set.len(), 1);
+        assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 3, 4, 5]);
+    }
 
-            #[test]
-            fn single_merge_low() {
-                let mut set = $set_type::new();
-                assert!(set.insert(0..2));
-                assert!(set.insert(4..6));
-                assert_eq!(set.len(), 2);
-                assert!(set.insert(2..3));
-                assert_eq!(set.len(), 2);
-                assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 4, 5]);
-            }
+    #[test]
+    fn single_merge_low() {
+        let mut set = ArrayRangeSet::new();
+        assert!(set.insert(0..2));
+        assert!(set.insert(4..6));
+        assert_eq!(set.len(), 2);
+        assert!(set.insert(2..3));
+        assert_eq!(set.len(), 2);
+        assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 4, 5]);
+    }
 
-            #[test]
-            fn single_merge_high() {
-                let mut set = $set_type::new();
-                assert!(set.insert(0..2));
-                assert!(set.insert(4..6));
-                assert_eq!(set.len(), 2);
-                assert!(set.insert(3..4));
-                assert_eq!(set.len(), 2);
-                assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 3, 4, 5]);
-            }
+    #[test]
+    fn single_merge_high() {
+        let mut set = ArrayRangeSet::new();
+        assert!(set.insert(0..2));
+        assert!(set.insert(4..6));
+        assert_eq!(set.len(), 2);
+        assert!(set.insert(3..4));
+        assert_eq!(set.len(), 2);
+        assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 3, 4, 5]);
+    }
 
-            #[test]
-            fn double_merge_wide() {
-                let mut set = $set_type::new();
-                assert!(set.insert(0..2));
-                assert!(set.insert(4..6));
-                assert_eq!(set.len(), 2);
-                assert!(set.insert(1..5));
-                assert_eq!(set.len(), 1);
-                assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 3, 4, 5]);
-            }
+    #[test]
+    fn double_merge_wide() {
+        let mut set = ArrayRangeSet::new();
+        assert!(set.insert(0..2));
+        assert!(set.insert(4..6));
+        assert_eq!(set.len(), 2);
+        assert!(set.insert(1..5));
+        assert_eq!(set.len(), 1);
+        assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 1, 2, 3, 4, 5]);
+    }
 
-            #[test]
-            fn double_remove() {
-                let mut set = $set_type::new();
-                assert!(set.insert(0..2));
-                assert!(set.insert(4..6));
-                assert!(set.remove(1..5));
-                assert_eq!(set.len(), 2);
-                assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 5]);
-            }
+    #[test]
+    fn double_remove() {
+        let mut set = ArrayRangeSet::new();
+        assert!(set.insert(0..2));
+        assert!(set.insert(4..6));
+        assert!(set.remove(1..5));
+        assert_eq!(set.len(), 2);
+        assert_eq!(&set.elts().collect::<Vec<_>>()[..], [0, 5]);
+    }
 
-            #[test]
-            fn insert_multiple() {
-                let mut set = $set_type::new();
-                assert!(set.insert(0..1));
-                assert!(set.insert(2..3));
-                assert!(set.insert(4..5));
-                assert!(set.insert(0..5));
-                assert_eq!(set.len(), 1);
-            }
+    #[test]
+    fn insert_multiple() {
+        let mut set = ArrayRangeSet::new();
+        assert!(set.insert(0..1));
+        assert!(set.insert(2..3));
+        assert!(set.insert(4..5));
+        assert!(set.insert(0..5));
+        assert_eq!(set.len(), 1);
+    }
 
-            #[test]
-            fn remove_multiple() {
-                let mut set = $set_type::new();
-                assert!(set.insert(0..1));
-                assert!(set.insert(2..3));
-                assert!(set.insert(4..5));
-                assert!(set.remove(0..5));
-                assert!(set.is_empty());
-            }
+    #[test]
+    fn remove_multiple() {
+        let mut set = ArrayRangeSet::new();
+        assert!(set.insert(0..1));
+        assert!(set.insert(2..3));
+        assert!(set.insert(4..5));
+        assert!(set.remove(0..5));
+        assert!(set.is_empty());
+    }
 
-            #[test]
-            fn double_insert() {
-                let mut set = $set_type::new();
-                assert!(set.insert(0..2));
-                assert!(!set.insert(0..2));
-                assert!(set.insert(2..4));
-                assert!(!set.insert(2..4));
-                assert!(!set.insert(0..4));
-                assert!(!set.insert(1..2));
-                assert!(!set.insert(1..3));
-                assert!(!set.insert(1..4));
-                assert_eq!(set.len(), 1);
-            }
+    #[test]
+    fn double_insert() {
+        let mut set = ArrayRangeSet::new();
+        assert!(set.insert(0..2));
+        assert!(!set.insert(0..2));
+        assert!(set.insert(2..4));
+        assert!(!set.insert(2..4));
+        assert!(!set.insert(0..4));
+        assert!(!set.insert(1..2));
+        assert!(!set.insert(1..3));
+        assert!(!set.insert(1..4));
+        assert_eq!(set.len(), 1);
+    }
 
-            #[test]
-            fn skip_empty_ranges() {
-                let mut set = $set_type::new();
-                assert!(!set.insert(2..2));
-                assert_eq!(set.len(), 0);
-                assert!(!set.insert(4..4));
-                assert_eq!(set.len(), 0);
-                assert!(!set.insert(0..0));
-                assert_eq!(set.len(), 0);
-            }
+    #[test]
+    fn skip_empty_ranges() {
+        let mut set = ArrayRangeSet::new();
+        assert!(!set.insert(2..2));
+        assert_eq!(set.len(), 0);
+        assert!(!set.insert(4..4));
+        assert_eq!(set.len(), 0);
+        assert!(!set.insert(0..0));
+        assert_eq!(set.len(), 0);
+    }
 
-            #[test]
-            fn compare_insert_to_reference() {
-                const MAX_RANGE: u64 = 50;
+    #[test]
+    fn compare_insert_to_reference() {
+        const MAX_RANGE: u64 = 50;
 
-                for start in 0..=MAX_RANGE {
-                    for end in 0..=MAX_RANGE {
-                        println!("insert({}..{})", start, end);
-                        let (mut set, mut reference) = create_initial_sets(MAX_RANGE);
-                        assert_eq!(set.insert(start..end), reference.insert(start..end));
-                        assert_sets_equal(&set, &reference);
-                    }
-                }
-            }
-
-            #[test]
-            fn compare_remove_to_reference() {
-                const MAX_RANGE: u64 = 50;
-
-                for start in 0..=MAX_RANGE {
-                    for end in 0..=MAX_RANGE {
-                        println!("remove({}..{})", start, end);
-                        let (mut set, mut reference) = create_initial_sets(MAX_RANGE);
-                        assert_eq!(set.remove(start..end), reference.remove(start..end));
-                        assert_sets_equal(&set, &reference);
-                    }
-                }
-            }
-
-            #[test]
-            fn min_max() {
-                let mut set = $set_type::new();
-                set.insert(1..3);
-                set.insert(4..5);
-                set.insert(6..10);
-                assert_eq!(set.min(), Some(1));
-                assert_eq!(set.max(), Some(9));
-            }
-
-            fn create_initial_sets(max_range: u64) -> ($set_type, RefRangeSet) {
-                let mut set = $set_type::new();
-                let mut reference = RefRangeSet::new(max_range as usize);
+        for start in 0..=MAX_RANGE {
+            for end in 0..=MAX_RANGE {
+                let (mut set, mut reference) = create_initial_sets(MAX_RANGE);
+                assert_eq!(set.insert(start..end), reference.insert(start..end));
                 assert_sets_equal(&set, &reference);
-
-                assert_eq!(set.insert(2..6), reference.insert(2..6));
-                assert_eq!(set.insert(10..14), reference.insert(10..14));
-                assert_eq!(set.insert(14..14), reference.insert(14..14));
-                assert_eq!(set.insert(18..19), reference.insert(18..19));
-                assert_eq!(set.insert(20..21), reference.insert(20..21));
-                assert_eq!(set.insert(22..24), reference.insert(22..24));
-                assert_eq!(set.insert(26..30), reference.insert(26..30));
-                assert_eq!(set.insert(34..38), reference.insert(34..38));
-                assert_eq!(set.insert(42..44), reference.insert(42..44));
-
-                assert_sets_equal(&set, &reference);
-
-                (set, reference)
-            }
-
-            fn assert_sets_equal(set: &$set_type, reference: &RefRangeSet) {
-                assert_eq!(set.len(), reference.len());
-                assert_eq!(set.is_empty(), reference.is_empty());
-                assert_eq!(set.elts().collect::<Vec<_>>()[..], reference.elts()[..]);
             }
         }
-    };
-}
+    }
 
-common_set_tests!(range_set, RangeSet);
-common_set_tests!(array_range_set, ArrayRangeSet);
+    #[test]
+    fn compare_remove_to_reference() {
+        const MAX_RANGE: u64 = 50;
+
+        for start in 0..=MAX_RANGE {
+            for end in 0..=MAX_RANGE {
+                let (mut set, mut reference) = create_initial_sets(MAX_RANGE);
+                assert_eq!(set.remove(start..end), reference.remove(start..end));
+                assert_sets_equal(&set, &reference);
+            }
+        }
+    }
+
+    #[test]
+    fn min_max() {
+        let mut set = ArrayRangeSet::new();
+        set.insert(1..3);
+        set.insert(4..5);
+        set.insert(6..10);
+        assert_eq!(set.min(), Some(1));
+        assert_eq!(set.max(), Some(9));
+    }
+
+    fn create_initial_sets(max_range: u64) -> (ArrayRangeSet, RefRangeSet) {
+        let mut set = ArrayRangeSet::new();
+        let mut reference = RefRangeSet::new(max_range as usize);
+        assert_sets_equal(&set, &reference);
+
+        assert_eq!(set.insert(2..6), reference.insert(2..6));
+        assert_eq!(set.insert(10..14), reference.insert(10..14));
+        assert_eq!(set.insert(14..14), reference.insert(14..14));
+        assert_eq!(set.insert(18..19), reference.insert(18..19));
+        assert_eq!(set.insert(20..21), reference.insert(20..21));
+        assert_eq!(set.insert(22..24), reference.insert(22..24));
+        assert_eq!(set.insert(26..30), reference.insert(26..30));
+        assert_eq!(set.insert(34..38), reference.insert(34..38));
+        assert_eq!(set.insert(42..44), reference.insert(42..44));
+
+        assert_sets_equal(&set, &reference);
+
+        (set, reference)
+    }
+
+    fn assert_sets_equal(set: &ArrayRangeSet, reference: &RefRangeSet) {
+        assert_eq!(set.len(), reference.len());
+        assert_eq!(set.is_empty(), reference.is_empty());
+        assert_eq!(set.elts().collect::<Vec<_>>()[..], reference.elts()[..]);
+    }
+}
 
 /// A very simple reference implementation of a RangeSet
 struct RefRangeSet {

quinn-proto/src/token_memory_cache.rs

@@ -157,6 +157,7 @@ mod tests {
     use super::*;
     use rand::prelude::*;
     use rand_pcg::Pcg32;
+    use tracing::info;
 
     fn new_rng() -> impl Rng {
         Pcg32::from_seed(0xdeadbeefdeadbeefdeadbeefdeadbeefu128.to_le_bytes())
@@ -176,7 +177,7 @@ mod tests {
                 if rng.random_bool(0.666) {
                     // store
                     let token = Bytes::from(vec![i]);
-                    println!("STORE {server_name} {token:?}");
+                    info!("STORE {server_name} {token:?}");
                     if let Some((j, _)) = cache_1
                         .iter()
                         .enumerate()
@@ -199,7 +200,7 @@ mod tests {
                     cache_2.insert(&server_name.to_string(), token);
                 } else {
                     // take
-                    println!("TAKE {server_name}");
+                    info!("TAKE {server_name}");
                     let expecting = cache_1
                         .iter()
                         .enumerate()
@@ -213,7 +214,7 @@ mod tests {
                             }
                             token
                         });
-                    println!("EXPECTING {expecting:?}");
+                    info!("EXPECTING {expecting:?}");
                     assert_eq!(cache_2.take(&server_name.to_string()), expecting);
                 }
             }