diff --git a/actix-connect/src/connect.rs b/actix-connect/src/connect.rs index 435a3497..49e75cf6 100644 --- a/actix-connect/src/connect.rs +++ b/actix-connect/src/connect.rs @@ -1,5 +1,6 @@ -use std::collections::VecDeque; +use std::collections::{vec_deque, VecDeque}; use std::fmt; +use std::iter::{FromIterator, FusedIterator}; use std::net::SocketAddr; use either::Either; @@ -77,6 +78,20 @@ impl Connect { self } + /// Use addresses. + pub fn set_addrs(mut self, addrs: I) -> Self + where + I: IntoIterator, + { + let mut addrs = VecDeque::from_iter(addrs); + self.addr = if addrs.len() < 2 { + addrs.pop_front().map(Either::Left) + } else { + Some(Either::Right(addrs)) + }; + self + } + /// Host name pub fn host(&self) -> &str { self.req.host() @@ -86,6 +101,28 @@ impl Connect { pub fn port(&self) -> u16 { self.req.port().unwrap_or(self.port) } + + /// Preresolved addresses of the request. + pub fn addrs(&self) -> ConnectAddrsIter<'_> { + let inner = match self.addr { + None => Either::Left(None), + Some(Either::Left(addr)) => Either::Left(Some(addr)), + Some(Either::Right(ref addrs)) => Either::Right(addrs.iter()), + }; + + ConnectAddrsIter { inner } + } + + /// Takes preresolved addresses of the request. + pub fn take_addrs(&mut self) -> ConnectTakeAddrsIter { + let inner = match self.addr.take() { + None => Either::Left(None), + Some(Either::Left(addr)) => Either::Left(Some(addr)), + Some(Either::Right(addrs)) => Either::Right(addrs.into_iter()), + }; + + ConnectTakeAddrsIter { inner } + } } impl From for Connect { @@ -100,6 +137,70 @@ impl fmt::Display for Connect { } } +/// Iterator over addresses in a [`Connect`](struct.Connect.html) request. +#[derive(Clone)] +pub struct ConnectAddrsIter<'a> { + inner: Either, vec_deque::Iter<'a, SocketAddr>>, +} + +impl Iterator for ConnectAddrsIter<'_> { + type Item = SocketAddr; + + fn next(&mut self) -> Option { + match self.inner { + Either::Left(ref mut opt) => opt.take(), + Either::Right(ref mut iter) => iter.next().copied(), + } + } + + fn size_hint(&self) -> (usize, Option) { + match self.inner { + Either::Left(Some(_)) => (1, Some(1)), + Either::Left(None) => (0, Some(0)), + Either::Right(ref iter) => iter.size_hint(), + } + } +} + +impl fmt::Debug for ConnectAddrsIter<'_> { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_list().entries(self.clone()).finish() + } +} + +impl ExactSizeIterator for ConnectAddrsIter<'_> {} + +impl FusedIterator for ConnectAddrsIter<'_> {} + +/// Owned iterator over addresses in a [`Connect`](struct.Connect.html) request. +#[derive(Debug)] +pub struct ConnectTakeAddrsIter { + inner: Either, vec_deque::IntoIter>, +} + +impl Iterator for ConnectTakeAddrsIter { + type Item = SocketAddr; + + fn next(&mut self) -> Option { + match self.inner { + Either::Left(ref mut opt) => opt.take(), + Either::Right(ref mut iter) => iter.next(), + } + } + + fn size_hint(&self) -> (usize, Option) { + match self.inner { + Either::Left(Some(_)) => (1, Some(1)), + Either::Left(None) => (0, Some(0)), + Either::Right(ref iter) => iter.size_hint(), + } + } +} + +impl ExactSizeIterator for ConnectTakeAddrsIter {} + +impl FusedIterator for ConnectTakeAddrsIter {} + fn parse(host: &str) -> (&str, Option) { let mut parts_iter = host.splitn(2, ':'); if let Some(host) = parts_iter.next() { diff --git a/actix-connect/src/resolver.rs b/actix-connect/src/resolver.rs index a5adc4ae..957bedc1 100644 --- a/actix-connect/src/resolver.rs +++ b/actix-connect/src/resolver.rs @@ -1,4 +1,3 @@ -use std::collections::VecDeque; use std::marker::PhantomData; use std::net::SocketAddr; @@ -162,23 +161,20 @@ impl Future for ResolverFuture { })? { Async::NotReady => Ok(Async::NotReady), Async::Ready(ips) => { - let mut req = self.req.take().unwrap(); - let mut addrs: VecDeque<_> = ips - .iter() - .map(|ip| SocketAddr::new(ip, req.port())) - .collect(); + let req = self.req.take().unwrap(); + + let port = req.port(); + let req = req.set_addrs(ips.iter().map(|ip| SocketAddr::new(ip, port))); + trace!( "DNS resolver: host {:?} resolved to {:?}", req.host(), - addrs + req.addrs() ); - if addrs.is_empty() { + + if req.addr.is_none() { Err(ConnectError::NoRecords) - } else if addrs.len() == 1 { - req.addr = Some(either::Either::Left(addrs.pop_front().unwrap())); - Ok(Async::Ready(req)) } else { - req.addr = Some(either::Either::Right(addrs)); Ok(Async::Ready(req)) } }