use Resolver service for Connector

src/connector.rs

@@ -1,54 +1,47 @@
 use std::collections::VecDeque;
 use std::io;
-use std::marker::PhantomData;
 use std::net::SocketAddr;
 
 use futures::{
     future::{ok, FutureResult},
     Async, Future, Poll,
 };
-use tokio;
 use tokio_tcp::{ConnectFuture, TcpStream};
 use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
-use trust_dns_resolver::lookup_ip::LookupIpFuture;
 use trust_dns_resolver::system_conf::read_system_conf;
-use trust_dns_resolver::{AsyncResolver, Background};
 
+use super::resolver::{HostAware, Resolver, ResolverError, ResolverFuture};
 use super::{NewService, Service};
 
-pub trait HostAware {
-    fn host(&self) -> &str;
-}
+// #[derive(Fail, Debug)]
 
-impl HostAware for String {
-    fn host(&self) -> &str {
-        self.as_ref()
-    }
-}
-
-#[derive(Fail, Debug)]
 pub enum ConnectorError {
     /// Failed to resolve the hostname
-    #[fail(display = "Failed resolving hostname: {}", _0)]
-    Resolver(String),
+    // #[fail(display = "Failed resolving hostname: {}", _0)]
+    Resolver(ResolverError),
 
-    /// Address is invalid
-    #[fail(display = "Invalid input: {}", _0)]
-    InvalidInput(&'static str),
+    /// Not dns records
+    // #[fail(display = "Invalid input: {}", _0)]
+    NoRecords,
 
     /// Connection io error
-    #[fail(display = "{}", _0)]
+    // #[fail(display = "{}", _0)]
     IoError(io::Error),
 }
 
+impl From<ResolverError> for ConnectorError {
+    fn from(err: ResolverError) -> Self {
+        ConnectorError::Resolver(err)
+    }
+}
+
 pub struct ConnectionInfo {
     pub host: String,
     pub addr: SocketAddr,
 }
 
 pub struct Connector<T = String> {
-    resolver: AsyncResolver,
-    req: PhantomData<T>,
+    resolver: Resolver<T>,
 }
 
 impl<T: HostAware> Default for Connector<T> {
@@ -65,11 +58,8 @@ impl<T: HostAware> Default for Connector<T> {
 
 impl<T: HostAware> Connector<T> {
     pub fn new(cfg: ResolverConfig, opts: ResolverOpts) -> Self {
-        let (resolver, bg) = AsyncResolver::new(cfg, opts);
-        tokio::spawn(bg);
         Connector {
-            resolver,
-            req: PhantomData,
+            resolver: Resolver::new(cfg, opts),
         }
     }
 
@@ -95,8 +85,7 @@ impl<T: HostAware> Connector<T> {
 
     pub fn change_request<T2: HostAware>(&self) -> Connector<T2> {
         Connector {
-            resolver: self.resolver.clone(),
-            req: PhantomData,
+            resolver: self.resolver.change_request(),
         }
     }
 }
@@ -105,7 +94,6 @@ impl<T> Clone for Connector<T> {
     fn clone(&self) -> Self {
         Connector {
             resolver: self.resolver.clone(),
-            req: PhantomData,
         }
     }
 }
@@ -121,14 +109,15 @@ impl<T: HostAware> Service for Connector<T> {
     }
 
     fn call(&mut self, req: Self::Request) -> Self::Future {
-        let fut = ResolveFut::new(req, 0, &self.resolver);
-
-        ConnectorFuture { fut, fut2: None }
+        ConnectorFuture {
+            fut: self.resolver.call(req),
+            fut2: None,
+        }
     }
 }
 
 pub struct ConnectorFuture<T: HostAware> {
-    fut: ResolveFut<T>,
+    fut: ResolverFuture<T>,
     fut2: Option<TcpConnector<T>>,
 }
 
@@ -140,11 +129,15 @@ impl<T: HostAware> Future for ConnectorFuture<T> {
         if let Some(ref mut fut) = self.fut2 {
             return fut.poll();
         }
-        match self.fut.poll()? {
+        match self.fut.poll().map_err(ConnectorError::from)? {
             Async::Ready((req, host, addrs)) => {
+                if addrs.is_empty() {
+                    Err(ConnectorError::NoRecords)
+                } else {
                     self.fut2 = Some(TcpConnector::new(req, host, addrs));
                     self.poll()
                 }
+            }
             Async::NotReady => Ok(Async::NotReady),
         }
     }
@@ -196,106 +189,6 @@ impl<T: HostAware> Future for DefaultConnectorFuture<T> {
     }
 }
 
-/// Resolver future
-struct ResolveFut<T> {
-    req: Option<T>,
-    host: Option<String>,
-    port: u16,
-    lookup: Option<Background<LookupIpFuture>>,
-    addrs: Option<VecDeque<SocketAddr>>,
-    error: Option<ConnectorError>,
-    error2: Option<String>,
-}
-
-impl<T: HostAware> ResolveFut<T> {
-    pub fn new(addr: T, port: u16, resolver: &AsyncResolver) -> Self {
-        // we need to do dns resolution
-        match ResolveFut::<T>::parse(addr.host(), port) {
-            Ok((host, port)) => {
-                let lookup = Some(resolver.lookup_ip(host.as_str()));
-                ResolveFut {
-                    port,
-                    lookup,
-                    req: Some(addr),
-                    host: Some(host),
-                    addrs: None,
-                    error: None,
-                    error2: None,
-                }
-            }
-            Err(err) => ResolveFut {
-                port,
-                req: None,
-                host: None,
-                lookup: None,
-                addrs: None,
-                error: Some(err),
-                error2: None,
-            },
-        }
-    }
-
-    fn parse(addr: &str, port: u16) -> Result<(String, u16), ConnectorError> {
-        macro_rules! try_opt {
-            ($e:expr, $msg:expr) => {
-                match $e {
-                    Some(r) => r,
-                    None => return Err(ConnectorError::InvalidInput($msg)),
-                }
-            };
-        }
-
-        // split the string by ':' and convert the second part to u16
-        let mut parts_iter = addr.splitn(2, ':');
-        let host = try_opt!(parts_iter.next(), "invalid socket address");
-        let port_str = parts_iter.next().unwrap_or("");
-        let port: u16 = port_str.parse().unwrap_or(port);
-
-        Ok((host.to_owned(), port))
-    }
-}
-
-impl<T> Future for ResolveFut<T> {
-    type Item = (T, String, VecDeque<SocketAddr>);
-    type Error = ConnectorError;
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        if let Some(err) = self.error.take() {
-            Err(err)
-        } else if let Some(err) = self.error2.take() {
-            Err(ConnectorError::Resolver(err))
-        } else if let Some(addrs) = self.addrs.take() {
-            Ok(Async::Ready((
-                self.req.take().unwrap(),
-                self.host.take().unwrap(),
-                addrs,
-            )))
-        } else {
-            match self.lookup.as_mut().unwrap().poll() {
-                Ok(Async::NotReady) => Ok(Async::NotReady),
-                Ok(Async::Ready(ips)) => {
-                    let addrs: VecDeque<_> = ips
-                        .iter()
-                        .map(|ip| SocketAddr::new(ip, self.port))
-                        .collect();
-                    if addrs.is_empty() {
-                        Err(ConnectorError::Resolver(
-                            "Expect at least one A dns record".to_owned(),
-                        ))
-                    } else {
-                        Ok(Async::Ready((
-                            self.req.take().unwrap(),
-                            self.host.take().unwrap(),
-                            addrs,
-                        )))
-                    }
-                }
-                Err(err) => Err(ConnectorError::Resolver(format!("{}", err))),
-            }
-        }
-    }
-}
-
 /// Tcp stream connector
 pub struct TcpConnector<T> {
     req: Option<T>,

src/lib.rs

@@ -19,7 +19,7 @@
 #[macro_use]
 extern crate log;
 extern crate bytes;
-#[macro_use]
+// #[macro_use]
 extern crate failure;
 #[macro_use]
 extern crate futures;

src/resolver.rs

@@ -54,6 +54,13 @@ impl<T: HostAware> Resolver<T> {
             req: PhantomData,
         }
     }
+
+    pub fn change_request<T2: HostAware>(&self) -> Resolver<T2> {
+        Resolver {
+            resolver: self.resolver.clone(),
+            req: PhantomData,
+        }
+    }
 }
 
 impl<T> Clone for Resolver<T> {
@@ -67,7 +74,7 @@ impl<T> Clone for Resolver<T> {
 
 impl<T: HostAware> Service for Resolver<T> {
     type Request = T;
-    type Response = VecDeque<SocketAddr>;
+    type Response = (T, String, VecDeque<SocketAddr>);
     type Error = ResolverError;
     type Future = ResolverFuture<T>;
 
@@ -83,11 +90,12 @@ impl<T: HostAware> Service for Resolver<T> {
 #[doc(hidden)]
 /// Resolver future
 pub struct ResolverFuture<T> {
+    req: Option<T>,
     port: u16,
     lookup: Option<Background<LookupIpFuture>>,
     addrs: Option<VecDeque<SocketAddr>>,
     error: Option<ResolverError>,
-    req: PhantomData<T>,
+    host: Option<String>,
 }
 
 impl<T: HostAware> ResolverFuture<T> {
@@ -99,17 +107,19 @@ impl<T: HostAware> ResolverFuture<T> {
                 ResolverFuture {
                     port,
                     lookup,
+                    req: Some(addr),
+                    host: Some(host.to_owned()),
                     addrs: None,
                     error: None,
-                    req: PhantomData,
                 }
             }
             Err(err) => ResolverFuture {
                 port,
+                req: None,
+                host: None,
                 lookup: None,
                 addrs: None,
                 error: Some(err),
-                req: PhantomData,
             },
         }
     }
@@ -126,14 +136,18 @@ impl<T: HostAware> ResolverFuture<T> {
 }
 
 impl<T: HostAware> Future for ResolverFuture<T> {
-    type Item = VecDeque<SocketAddr>;
+    type Item = (T, String, VecDeque<SocketAddr>);
     type Error = ResolverError;
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
         if let Some(err) = self.error.take() {
             Err(err)
         } else if let Some(addrs) = self.addrs.take() {
-            Ok(Async::Ready(addrs))
+            Ok(Async::Ready((
+                self.req.take().unwrap(),
+                self.host.take().unwrap(),
+                addrs,
+            )))
         } else {
             match self.lookup.as_mut().unwrap().poll() {
                 Ok(Async::NotReady) => Ok(Async::NotReady),
@@ -142,7 +156,11 @@ impl<T: HostAware> Future for ResolverFuture<T> {
                         .iter()
                         .map(|ip| SocketAddr::new(ip, self.port))
                         .collect();
-                    Ok(Async::Ready(addrs))
+                    Ok(Async::Ready((
+                        self.req.take().unwrap(),
+                        self.host.take().unwrap(),
+                        addrs,
+                    )))
                 }
                 Err(err) => Err(ResolverError::Resolve(err)),
             }