From b946d06c5e54e140a516f103496647652a62ca6d Mon Sep 17 00:00:00 2001 From: Rob Ede Date: Mon, 8 Feb 2021 07:47:47 +0000 Subject: [PATCH] add draining iterator --- actix-http/src/client/h1proto.rs | 4 +- actix-http/src/header/map.rs | 172 +++++++++++++++++++++++-------- actix-http/src/response.rs | 4 +- awc/src/frozen.rs | 4 +- awc/src/request.rs | 8 +- awc/src/sender.rs | 4 +- 6 files changed, 147 insertions(+), 49 deletions(-) diff --git a/actix-http/src/client/h1proto.rs b/actix-http/src/client/h1proto.rs index 24f4207e8..92c3c0e1b 100644 --- a/actix-http/src/client/h1proto.rs +++ b/actix-http/src/client/h1proto.rs @@ -48,11 +48,11 @@ where match wrt.get_mut().split().freeze().try_into_value() { Ok(value) => match head { RequestHeadType::Owned(ref mut head) => { - head.headers.insert(HOST, value) + head.headers.insert(HOST, value); } RequestHeadType::Rc(_, ref mut extra_headers) => { let headers = extra_headers.get_or_insert(HeaderMap::new()); - headers.insert(HOST, value) + headers.insert(HOST, value); } }, Err(e) => log::error!("Can not set HOST header {}", e), diff --git a/actix-http/src/header/map.rs b/actix-http/src/header/map.rs index 69ac5885a..67bee6295 100644 --- a/actix-http/src/header/map.rs +++ b/actix-http/src/header/map.rs @@ -131,24 +131,6 @@ impl HeaderMap { self.inner.clear(); } - /// Returns the number of single-value headers the map can hold without needing to reallocate. - /// - /// Since this is a multi-value map, the actual capacity is much larger when considering - /// each header name can be associated with an arbitrary number of values. The effect is that - /// the size of `len` may be greater than `capacity` since it counts all the values. - /// Conversely, [`len_keys`](Self::len_keys) will never be larger than capacity. - pub fn capacity(&self) -> usize { - self.inner.capacity() - } - - /// Reserves capacity for at least `additional` more headers to be inserted in the map. - /// - /// The header map may reserve more space to avoid frequent reallocations. Additional capacity - /// only considers single-value headers. - pub fn reserve(&mut self, additional: usize) { - self.inner.reserve(additional) - } - fn get_value(&self, key: impl AsHeaderName) -> Option<&Value> { match key.try_as_name().ok()? { Cow::Borrowed(name) => self.inner.get(name), @@ -197,22 +179,6 @@ impl HeaderMap { } } - /// An iterator over all name-value pairs. - /// - /// Names will be yielded for each associated value. So, if a key has 3 associated values, it - /// will be yielded 3 times. The iteration order should be considered arbitrary. - pub fn iter(&self) -> Iter<'_> { - Iter::new(self.inner.iter()) - } - - /// An iterator over all contained header names. - /// - /// Each name will only be yielded once even if it has multiple associated values. The iteration - /// order should be considered arbitrary. - pub fn keys(&self) -> Keys<'_> { - Keys(self.inner.keys()) - } - /// Inserts a name-value pair into the map. /// /// If the map already contained this key, the new value is associated with the key and all @@ -248,6 +214,51 @@ impl HeaderMap { Removed::new(value) } + + /// Returns the number of single-value headers the map can hold without needing to reallocate. + /// + /// Since this is a multi-value map, the actual capacity is much larger when considering + /// each header name can be associated with an arbitrary number of values. The effect is that + /// the size of `len` may be greater than `capacity` since it counts all the values. + /// Conversely, [`len_keys`](Self::len_keys) will never be larger than capacity. + pub fn capacity(&self) -> usize { + self.inner.capacity() + } + + /// Reserves capacity for at least `additional` more headers to be inserted in the map. + /// + /// The header map may reserve more space to avoid frequent reallocations. Additional capacity + /// only considers single-value headers. + pub fn reserve(&mut self, additional: usize) { + self.inner.reserve(additional) + } + + /// An iterator over all name-value pairs. + /// + /// Names will be yielded for each associated value. So, if a key has 3 associated values, it + /// will be yielded 3 times. The iteration order should be considered arbitrary. + pub fn iter(&self) -> Iter<'_> { + Iter::new(self.inner.iter()) + } + + /// An iterator over all contained header names. + /// + /// Each name will only be yielded once even if it has multiple associated values. The iteration + /// order should be considered arbitrary. + pub fn keys(&self) -> Keys<'_> { + Keys(self.inner.keys()) + } + + /// Clears the map, returning all name-value sets as an iterator. + /// + /// Header names will only be yielded for the first value in each set. All items that are + /// yielded without a name and after an item with a name are associated with that same name. + /// The first item will always contain a name. + /// + /// Keeps the allocated memory for reuse. + pub fn drain(&mut self) -> Drain<'_> { + Drain::new(self.inner.drain()) + } } /// Note that this implementation will clone a [HeaderName] for each value. @@ -335,7 +346,7 @@ impl<'a> Iterator for GetAll<'a> { type Item = &'a HeaderValue; #[inline] - fn next(&mut self) -> Option<&'a HeaderValue> { + fn next(&mut self) -> Option { let val = self.value?; match val { @@ -382,7 +393,7 @@ impl Iterator for Removed { type Item = HeaderValue; #[inline] - fn next(&mut self) -> Option { + fn next(&mut self) -> Option { self.inner.next() } } @@ -395,14 +406,14 @@ impl<'a> Iterator for Keys<'a> { type Item = &'a HeaderName; #[inline] - fn next(&mut self) -> Option<&'a HeaderName> { + fn next(&mut self) -> Option { self.0.next() } } #[derive(Debug)] pub struct Iter<'a> { - iter: hash_map::Iter<'a, HeaderName, Value>, + inner: hash_map::Iter<'a, HeaderName, Value>, multi_inner: Option<(&'a HeaderName, &'a SmallVec<[HeaderValue; 4]>)>, multi_idx: usize, } @@ -410,7 +421,7 @@ pub struct Iter<'a> { impl<'a> Iter<'a> { fn new(iter: hash_map::Iter<'a, HeaderName, Value>) -> Self { Self { - iter, + inner: iter, multi_idx: 0, multi_inner: None, } @@ -421,7 +432,7 @@ impl<'a> Iterator for Iter<'a> { type Item = (&'a HeaderName, &'a HeaderValue); #[inline] - fn next(&mut self) -> Option<(&'a HeaderName, &'a HeaderValue)> { + fn next(&mut self) -> Option { // handle in-progress multi value lists first if let Some((ref name, ref mut vals)) = self.multi_inner { match vals.get(self.multi_idx) { @@ -437,7 +448,7 @@ impl<'a> Iterator for Iter<'a> { } } - let (name, value) = self.iter.next()?; + let (name, value) = self.inner.next()?; match value { Value::One(ref value) => Some((name, value)), @@ -450,6 +461,56 @@ impl<'a> Iterator for Iter<'a> { } } +/// Iterator over drained name-value pairs. +/// +/// Iterator items are `(Option, HeaderValue)` to avoid cloning. +#[derive(Debug)] +pub struct Drain<'a> { + inner: hash_map::Drain<'a, HeaderName, Value>, + multi_inner: Option<(Option, SmallVec<[HeaderValue; 4]>)>, + multi_inner_idx: usize, +} + +impl<'a> Drain<'a> { + fn new(iter: hash_map::Drain<'a, HeaderName, Value>) -> Self { + Self { + inner: iter, + multi_inner: None, + multi_inner_idx: 0, + } + } +} + +impl<'a> Iterator for Drain<'a> { + type Item = (Option, HeaderValue); + + #[inline] + fn next(&mut self) -> Option { + // handle in-progress multi value iterators first + if let Some((ref mut name, ref mut vals)) = self.multi_inner { + if vals.len() > 0 { + // OPTIMISE: array removals + return Some((name.take(), vals.remove(0))); + } else { + // no more items in value iterator; reset state + self.multi_inner = None; + self.multi_inner_idx = 0; + } + } + + let (name, mut value) = self.inner.next()?; + + match value { + Value::One(value) => Some((Some(name), value)), + Value::Multi(ref mut vals) => { + // set up new multi value inner iter and recurse into it + self.multi_inner = Some((Some(name), mem::take(vals))); + self.next() + } + } + } +} + /// Iterator over owned name-value pairs. /// /// Implementation necessarily clones header names for each value. @@ -472,7 +533,7 @@ impl Iterator for IntoIter { type Item = (HeaderName, HeaderValue); #[inline] - fn next(&mut self) -> Option<(HeaderName, HeaderValue)> { + fn next(&mut self) -> Option { // handle in-progress multi value iterators first if let Some((ref name, ref mut vals)) = self.multi_inner { match vals.next() { @@ -556,6 +617,33 @@ mod tests { assert!(pairs.contains(&(&header::COOKIE, &HeaderValue::from_static("two=2")))); } + #[test] + fn drain_iter() { + let mut map = HeaderMap::new(); + + map.append(header::COOKIE, HeaderValue::from_static("one=1")); + map.append(header::COOKIE, HeaderValue::from_static("two=2")); + + let mut vals = vec![]; + let mut iter = map.drain(); + + let (name, val) = iter.next().unwrap(); + assert_eq!(name, Some(header::COOKIE)); + vals.push(val); + + let (name, val) = iter.next().unwrap(); + assert!(name.is_none()); + vals.push(val); + + assert!(vals.contains(&HeaderValue::from_static("one=1"))); + assert!(vals.contains(&HeaderValue::from_static("two=2"))); + + assert!(iter.next().is_none()); + drop(iter); + + assert!(map.is_empty()); + } + #[test] fn entries_into_iter() { let mut map = HeaderMap::new(); diff --git a/actix-http/src/response.rs b/actix-http/src/response.rs index b9850f750..56a5471fd 100644 --- a/actix-http/src/response.rs +++ b/actix-http/src/response.rs @@ -363,7 +363,9 @@ impl ResponseBuilder { { if let Some(parts) = parts(&mut self.head, &self.err) { match header.try_into_header_pair() { - Ok((key, value)) => parts.headers.insert(key, value), + Ok((key, value)) => { + parts.headers.insert(key, value); + } Err(e) => self.err = Some(e.into()), }; } diff --git a/awc/src/frozen.rs b/awc/src/frozen.rs index 878f404c6..46b4063a0 100644 --- a/awc/src/frozen.rs +++ b/awc/src/frozen.rs @@ -145,7 +145,9 @@ impl FrozenSendBuilder { { match HeaderName::try_from(key) { Ok(key) => match value.try_into_value() { - Ok(value) => self.extra_headers.insert(key, value), + Ok(value) => { + self.extra_headers.insert(key, value); + } Err(e) => self.err = Some(e.into()), }, Err(e) => self.err = Some(e.into()), diff --git a/awc/src/request.rs b/awc/src/request.rs index b9a333b18..c87df9b3b 100644 --- a/awc/src/request.rs +++ b/awc/src/request.rs @@ -159,7 +159,9 @@ impl ClientRequest { H: IntoHeaderPair, { match header.try_into_header_pair() { - Ok((key, value)) => self.head.headers.insert(key, value), + Ok((key, value)) => { + self.head.headers.insert(key, value); + } Err(e) => self.err = Some(e.into()), }; @@ -232,7 +234,9 @@ impl ClientRequest { >::Error: Into, { match HeaderValue::try_from(value) { - Ok(value) => self.head.headers.insert(header::CONTENT_TYPE, value), + Ok(value) => { + self.head.headers.insert(header::CONTENT_TYPE, value); + } Err(e) => self.err = Some(e.into()), } self diff --git a/awc/src/sender.rs b/awc/src/sender.rs index 1cf863d96..ca220be71 100644 --- a/awc/src/sender.rs +++ b/awc/src/sender.rs @@ -304,7 +304,9 @@ impl RequestSender { RequestSender::Owned(head) => { if !head.headers.contains_key(&key) { match value.try_into_value() { - Ok(value) => head.headers.insert(key, value), + Ok(value) => { + head.headers.insert(key, value); + } Err(e) => return Err(e.into()), } }