From 563bcee4dc7f18264b99a708313d725c4fec5e4b Mon Sep 17 00:00:00 2001 From: Sam Rijs Date: Tue, 16 May 2017 00:28:04 +1000 Subject: [PATCH] feat(client): make dns resolver pluggable This caters for the use-case of switching out the DNS implementation, for instance to influence caching behaviour, or to exchange it for a futures-based async resolver. --- src/client/connect.rs | 54 +++++++++++++++++++++++++---------------- src/client/dns.rs | 56 ++++++++++++++++++++++++++++++++++++------- src/client/mod.rs | 1 + 3 files changed, 82 insertions(+), 29 deletions(-) diff --git a/src/client/connect.rs b/src/client/connect.rs index 745a281d1d..c866130330 100644 --- a/src/client/connect.rs +++ b/src/client/connect.rs @@ -2,6 +2,7 @@ use std::error::Error as StdError; use std::fmt; use std::io; use std::mem; +use std::rc::Rc; //use std::net::SocketAddr; use futures::{Future, Poll, Async}; @@ -11,7 +12,7 @@ use tokio::net::{TcpStream, TcpStreamNew}; use tokio_service::Service; use Uri; -use super::dns; +use super::dns::{self, DnsService}; /// A connector creates an Io to a remote address.. /// @@ -42,21 +43,34 @@ where T: Service + 'static, /// A connector for the `http` scheme. #[derive(Clone)] -pub struct HttpConnector { - dns: dns::Dns, +pub struct HttpConnector { + dns: Rc, enforce_http: bool, handle: Handle, } impl HttpConnector { - /// Construct a new HttpConnector. /// /// Takes number of DNS worker threads. #[inline] pub fn new(threads: usize, handle: &Handle) -> HttpConnector { HttpConnector { - dns: dns::Dns::new(threads), + dns: Rc::new(dns::DnsPool::new(threads)), + enforce_http: true, + handle: handle.clone(), + } + } +} + +impl HttpConnector { + /// Construct a new HttpConnector. + /// + /// Takes a `DnsService`. + #[inline] + pub fn with_dns(dns: D, handle: &Handle) -> HttpConnector { + HttpConnector { + dns: Rc::new(dns), enforce_http: true, handle: handle.clone(), } @@ -71,7 +85,7 @@ impl HttpConnector { } } -impl fmt::Debug for HttpConnector { +impl fmt::Debug for HttpConnector { #[inline] fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.debug_struct("HttpConnector") @@ -79,11 +93,11 @@ impl fmt::Debug for HttpConnector { } } -impl Service for HttpConnector { +impl Service for HttpConnector { type Request = Uri; type Response = TcpStream; type Error = io::Error; - type Future = HttpConnecting; + type Future = HttpConnecting; fn call(&self, uri: Uri) -> Self::Future { debug!("Http::connect({:?})", uri); @@ -117,7 +131,7 @@ impl Service for HttpConnector { } #[inline] -fn invalid_url(err: InvalidUrl, handle: &Handle) -> HttpConnecting { +fn invalid_url(err: InvalidUrl, handle: &Handle) -> HttpConnecting { HttpConnecting { state: State::Error(Some(io::Error::new(io::ErrorKind::InvalidInput, err))), handle: handle.clone(), @@ -148,19 +162,19 @@ impl StdError for InvalidUrl { } /// A Future representing work to connect to a URL. -pub struct HttpConnecting { - state: State, +pub struct HttpConnecting { + state: State, handle: Handle, } -enum State { - Lazy(dns::Dns, String, u16), - Resolving(dns::Query), - Connecting(ConnectingTcp), +enum State { + Lazy(Rc, String, u16), + Resolving(::Future), + Connecting(ConnectingTcp), Error(Option), } -impl Future for HttpConnecting { +impl Future for HttpConnecting { type Item = TcpStream; type Error = io::Error; @@ -177,7 +191,7 @@ impl Future for HttpConnecting { Async::NotReady => return Ok(Async::NotReady), Async::Ready(addrs) => { state = State::Connecting(ConnectingTcp { - addrs: addrs, + addrs: addrs.into_iter(), current: None, }) } @@ -197,12 +211,12 @@ impl fmt::Debug for HttpConnecting { } } -struct ConnectingTcp { - addrs: dns::IpAddrs, +struct ConnectingTcp { + addrs: ::IntoIter, current: Option, } -impl ConnectingTcp { +impl ConnectingTcp { // not a Future, since passing a &Handle to poll fn poll(&mut self, handle: &Handle) -> Poll { let mut err = None; diff --git a/src/client/dns.rs b/src/client/dns.rs index 5badd5c5f7..9b97965048 100644 --- a/src/client/dns.rs +++ b/src/client/dns.rs @@ -1,3 +1,4 @@ +use std::fmt; use std::io; use std::net::{SocketAddr, ToSocketAddrs}; use std::vec; @@ -5,26 +6,56 @@ use std::vec; use ::futures::{Future, Poll}; use ::futures_cpupool::{CpuPool, CpuFuture}; +/// Trait for resolving a host and port into a set of ip addresses. +pub trait DnsService { + /// Response representing a set of ip addresses. + type Response: IntoIterator; + /// Future being returned from `resolve`. + type Future: Future; + + /// Resolves a host and port into a set of ip addresses. + fn resolve(&self, host: String, port: u16) -> Self::Future; +} + #[derive(Clone)] -pub struct Dns { +/// An implementation of a dns service using a thread pool. +pub struct DnsPool { pool: CpuPool, } -impl Dns { - pub fn new(threads: usize) -> Dns { - Dns { +impl DnsPool { + /// Create a new dns pool given a number of worker threads. + pub fn new(threads: usize) -> DnsPool { + DnsPool { pool: CpuPool::new(threads) } } +} + +impl fmt::Debug for DnsPool { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.pad("DnsPool") + } +} - pub fn resolve(&self, host: String, port: u16) -> Query { - Query(self.pool.spawn_fn(move || work(host, port))) +impl DnsService for DnsPool { + type Response = IpAddrs; + type Future = DnsQuery; + fn resolve(&self, host: String, port: u16) -> Self::Future { + DnsQuery(self.pool.spawn_fn(move || work(host, port))) } } -pub struct Query(CpuFuture); +/// A `Future` that will resolve to `IpAddrs`. +pub struct DnsQuery(CpuFuture); -impl Future for Query { +impl fmt::Debug for DnsQuery { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.pad("DnsQuery") + } +} + +impl Future for DnsQuery { type Item = IpAddrs; type Error = io::Error; @@ -33,10 +64,17 @@ impl Future for Query { } } +/// A set of ip addresses. pub struct IpAddrs { iter: vec::IntoIter, } +impl fmt::Debug for IpAddrs { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.pad("IpAddrs") + } +} + impl Iterator for IpAddrs { type Item = SocketAddr; #[inline] @@ -45,7 +83,7 @@ impl Iterator for IpAddrs { } } -pub type Answer = io::Result; +type Answer = io::Result; fn work(hostname: String, port: u16) -> Answer { debug!("resolve {:?}:{:?}", hostname, port); diff --git a/src/client/mod.rs b/src/client/mod.rs index 1f217a7389..de2d7507f9 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -28,6 +28,7 @@ use uri::{self, Uri}; pub use http::response::Response; pub use http::request::Request; pub use self::connect::{HttpConnector, Connect}; +pub use self::dns::{DnsPool, DnsQuery, DnsService, IpAddrs}; mod connect; mod dns;