Skip to content

Commit

Permalink
h3i: API cleanup
Browse files Browse the repository at this point in the history 
h3i's assertion API philosophy is to offload as much of the processing
as possible to the user, for maximum flexibility. As such, we shouldn't
expose a function that shows all HEADERS - rather, we should just allow
the user to provide their own lambdas.

If, in the future, we want to provide convenience methods for _all_
frame types, we can do so. However, we should think more on where we
want to draw that line.

This also changes StreamMap's backing store from a HashMap to a BTreeMap
to make inter-stream ordering deterministic when iterating over all
received frames.
  • Loading branch information
@evanrittenhouse
evanrittenhouse committed Dec 20, 2024
1 parent 4e1a3ed commit f6f9be3
Showing 1 changed file with 125 additions and 48 deletions.
173 changes: 125 additions & 48 deletions h3i/src/client/connection_summary.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,10 +35,9 @@ use serde::ser::SerializeStruct;
use serde::ser::Serializer;
use serde::Serialize;
use std::cmp;
use std::collections::HashMap;
use std::collections::BTreeMap;
use std::iter::FromIterator;

use crate::frame::EnrichedHeaders;
use crate::frame::ExpectedFrame;
use crate::frame::H3iFrame;

Expand Down Expand Up @@ -87,7 +86,7 @@ impl Serialize for ConnectionSummary {
/// stream ID over which they were received.
#[derive(Clone, Debug, Default, Serialize)]
pub struct StreamMap {
map: HashMap<u64, Vec<H3iFrame>>,
map: BTreeMap<u64, Vec<H3iFrame>>,
expected_frames: Option<ExpectedFrames>,
}

Expand All @@ -96,7 +95,7 @@ where
T: IntoIterator<Item = (u64, Vec<H3iFrame>)>,
{
fn from(value: T) -> Self {
let map = HashMap::from_iter(value);
let map = BTreeMap::from_iter(value);
Self {
map,
expected_frames: None,
Expand All @@ -105,8 +104,7 @@ where
}

impl StreamMap {
/// Flatten all received frames into a single vector. The ordering is
/// non-deterministic.
/// Get all frames on a given `stream_id`.
///
/// # Example
///
Expand All @@ -117,41 +115,150 @@ impl StreamMap {
/// use quiche::h3::Header;
/// use std::iter::FromIterator;
///
/// let mut stream_map = StreamMap::default();
///
/// let h = Header::new(b"hello", b"world");
/// let headers = H3iFrame::Headers(EnrichedHeaders::from(vec![h]));
///
/// let stream_map: StreamMap = [(0, vec![headers.clone()])].into();
/// assert_eq!(stream_map.all_frames(), vec![headers]);
/// assert_eq!(stream_map.stream(0), vec![headers]);
/// ```
pub fn all_frames(&self) -> Vec<H3iFrame> {
pub fn stream(&self, stream_id: u64) -> Vec<H3iFrame> {
self.map.get(&stream_id).cloned().unwrap_or_default()
}

/// Collect all frames that match the predicate on a given stream.
///
/// # Example
///
/// ```
/// use h3i::client::connection_summary::StreamMap;
/// use h3i::frame::EnrichedHeaders;
/// use h3i::frame::H3iFrame;
/// use quiche::h3::frame::Frame;
/// use quiche::h3::Header;
/// use std::iter::FromIterator;
///
/// let h = Header::new(b"hello", b"world");
/// let headers = H3iFrame::Headers(EnrichedHeaders::from(vec![h]));
/// let data = H3iFrame::QuicheH3(Frame::Data {
/// payload: b"nyj".to_vec(),
/// });
/// let stream_map: StreamMap =
/// [(0, vec![headers.clone(), data]), (4, vec![headers.clone()])].into();
///
/// assert_eq!(
/// stream_map.filter_stream(0, |frame| {
/// match frame {
/// H3iFrame::Headers(_) => true,
/// _ => false,
/// }
/// }),
/// vec![headers]
/// );
/// ```
pub fn filter_stream<F>(&self, stream_id: u64, f: F) -> Vec<H3iFrame>
where
F: Fn(&H3iFrame) -> bool,
{
self.stream(stream_id)
.into_iter()
.filter(|frame| f(frame))
.collect()
}

/// Collect all frames that match the predicate. The resulting vector
/// is ordered by stream in increasing order, with all frames received
/// on that stream in order.
///
/// # Example
///
/// ```
/// use h3i::client::connection_summary::StreamMap;
/// use h3i::frame::EnrichedHeaders;
/// use h3i::frame::H3iFrame;
/// use quiche::h3::frame::Frame;
/// use quiche::h3::Header;
/// use std::iter::FromIterator;
///
/// let first_headers =
/// H3iFrame::Headers(EnrichedHeaders::from(vec![Header::new(
/// b"hello", b"world",
/// )]));
/// let second_headers =
/// H3iFrame::Headers(EnrichedHeaders::from(vec![Header::new(
/// b"go", b"jets",
/// )]));
/// let data = H3iFrame::QuicheH3(Frame::Data {
/// payload: b"nyj".to_vec(),
/// });
///
/// let stream_map: StreamMap = [
/// (0, vec![first_headers.clone(), data.clone()]),
/// (4, vec![second_headers.clone()]),
/// ]
/// .into();
/// assert_eq!(
/// stream_map.filter(|frame| {
/// match frame {
/// H3iFrame::Headers(_) => true,
/// _ => false,
/// }
/// }),
/// vec![first_headers, second_headers]
/// );
/// ```
pub fn filter<F>(&self, f: F) -> Vec<H3iFrame>
where
F: Fn(&H3iFrame) -> bool,
{
self.map
.values()
.flatten()
.filter(|frame| f(frame))
.map(Clone::clone)
.collect::<Vec<H3iFrame>>()
.collect()
}

/// Get all frames on a given `stream_id`.
/// Flatten all received frames into a single vector. The resulting vector
/// is ordered by stream in increasing order, with all frames received
/// on that stream in order.
///
/// # Example
///
/// ```
/// use h3i::client::connection_summary::StreamMap;
/// use h3i::frame::EnrichedHeaders;
/// use h3i::frame::H3iFrame;
/// use quiche::h3::frame::Frame;
/// use quiche::h3::Header;
/// use std::iter::FromIterator;
///
/// let mut stream_map = StreamMap::default();
///
/// let h = Header::new(b"hello", b"world");
/// let headers = H3iFrame::Headers(EnrichedHeaders::from(vec![h]));
/// let first_headers =
/// H3iFrame::Headers(EnrichedHeaders::from(vec![Header::new(
/// b"hello", b"world",
/// )]));
/// let second_headers =
/// H3iFrame::Headers(EnrichedHeaders::from(vec![Header::new(
/// b"go", b"jets",
/// )]));
/// let data = H3iFrame::QuicheH3(Frame::Data {
/// payload: b"nyj".to_vec(),
/// });
///
/// let stream_map: StreamMap = [(0, vec![headers.clone()])].into();
/// assert_eq!(stream_map.stream(0), vec![headers]);
/// let stream_map: StreamMap = [
/// (0, vec![first_headers.clone(), data.clone()]),
/// (4, vec![second_headers.clone()]),
/// ]
/// .into();
/// assert_eq!(stream_map.all_frames(), vec![
/// first_headers,
/// data,
/// second_headers
/// ]);
/// ```
pub fn stream(&self, stream_id: u64) -> Vec<H3iFrame> {
self.map.get(&stream_id).cloned().unwrap_or_default()
pub fn all_frames(&self) -> Vec<H3iFrame> {
self.filter(|_| true)
}

/// Check if a provided [`H3iFrame`] was received, regardless of what stream
Expand Down Expand Up @@ -223,36 +330,6 @@ impl StreamMap {
self.map.is_empty()
}

/// See all HEADERS received on a given stream.
///
/// # Example
///
/// ```
/// use h3i::client::connection_summary::StreamMap;
/// use h3i::frame::EnrichedHeaders;
/// use h3i::frame::H3iFrame;
/// use quiche::h3::Header;
/// use std::iter::FromIterator;
///
/// let mut stream_map = StreamMap::default();
///
/// let h = Header::new(b"hello", b"world");
/// let enriched = EnrichedHeaders::from(vec![h]);
/// let headers = H3iFrame::Headers(enriched.clone());
/// let data = H3iFrame::QuicheH3(quiche::h3::frame::Frame::Data {
/// payload: b"hello world".to_vec(),
/// });
///
/// let stream_map: StreamMap = [(0, vec![headers.clone(), data.clone()])].into();
/// assert_eq!(stream_map.headers_on_stream(0), vec![enriched]);
/// ```
pub fn headers_on_stream(&self, stream_id: u64) -> Vec<EnrichedHeaders> {
self.stream(stream_id)
.into_iter()
.filter_map(|h3i_frame| h3i_frame.to_enriched_headers())
.collect()
}

pub(crate) fn new(expected_frames: Option<ExpectedFrames>) -> Self {
Self {
expected_frames,
Expand Down

0 comments on commit f6f9be3

Please sign in to comment.