From 138e1643e81669cae9dbe215197abd0e07f0c1e7 Mon Sep 17 00:00:00 2001
From: Joe Wilm <joe@jwilm.com>
Date: Thu, 6 Oct 2016 14:50:30 -0700
Subject: [PATCH] feat(client): DNS worker count is configurable

When loading up a client suddenly with thousands of connections, the
default DNS worker count of four cannot keep up and many requests
timeout as a result. Most people don't need a large pool, so making this
configurable is a natural choice.
---
 src/client/connect.rs | 13 ++++++++++++-
 src/client/mod.rs     | 13 +++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/client/connect.rs b/src/client/connect.rs
index 19aee61584..2c314d6e2a 100644
--- a/src/client/connect.rs
+++ b/src/client/connect.rs
@@ -24,6 +24,9 @@ pub trait Connect {
     /// Returns a connected socket and associated host.
     fn connected(&mut self) -> Option<(Self::Key, io::Result<Self::Output>)>;
     #[doc(hidden)]
+    /// Configure number of dns workers to use.
+    fn dns_workers(&mut self, usize);
+    #[doc(hidden)]
     fn register(&mut self, Registration);
 }
 
@@ -71,6 +74,10 @@ impl Connect for HttpConnector {
     type Output = HttpStream;
     type Key = (&'static str, String, u16);
 
+    fn dns_workers(&mut self, count: usize) {
+        self.threads = count;
+    }
+
     fn key(&self, url: &Url) -> Option<Self::Key> {
         if url.scheme() == "http" {
             Some((
@@ -119,7 +126,7 @@ impl Connect for HttpConnector {
     }
 
     fn register(&mut self, reg: Registration) {
-        self.dns = Some(Dns::new(reg.notify, 4));
+        self.dns = Some(Dns::new(reg.notify, self.threads));
     }
 }
 
@@ -144,6 +151,10 @@ impl<S: SslClient> Connect for HttpsConnector<S> {
     type Output = HttpsStream<S::Stream>;
     type Key = (&'static str, String, u16);
 
+    fn dns_workers(&mut self, count: usize) {
+        self.http.dns_workers(count)
+    }
+
     fn key(&self, url: &Url) -> Option<Self::Key> {
         let scheme = match url.scheme() {
             "http" => "http",
diff --git a/src/client/mod.rs b/src/client/mod.rs
index 515eef6457..1a90a246f0 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -87,6 +87,7 @@ impl<H: Send> Client<H> {
         let mut loop_ = try!(rotor::Loop::new(&rotor_config));
         let mut notifier = None;
         let mut connector = config.connector;
+        connector.dns_workers(config.dns_workers);
         {
             let not = &mut notifier;
             loop_.add_machine_with(move |scope| {
@@ -150,6 +151,7 @@ pub struct Config<C> {
     //TODO: make use of max_idle config
     max_idle: usize,
     max_sockets: usize,
+    dns_workers: usize,
 }
 
 impl<C> Config<C> where C: Connect + Send + 'static {
@@ -163,6 +165,7 @@ impl<C> Config<C> where C: Connect + Send + 'static {
             keep_alive_timeout: Some(Duration::from_secs(60 * 2)),
             max_idle: self.max_idle,
             max_sockets: self.max_sockets,
+            dns_workers: self.dns_workers,
         }
     }
 
@@ -204,6 +207,15 @@ impl<C> Config<C> where C: Connect + Send + 'static {
         self
     }
 
+    /// Set number of Dns workers to use for this client
+    ///
+    /// Default is 4
+    #[inline]
+    pub fn dns_workers(mut self, workers: usize) -> Config<C> {
+        self.dns_workers = workers;
+        self
+    }
+
     /// Construct the Client with this configuration.
     #[inline]
     pub fn build<H: Handler<C::Output>>(self) -> ::Result<Client<H>> {
@@ -220,6 +232,7 @@ impl Default for Config<DefaultConnector> {
             keep_alive_timeout: Some(Duration::from_secs(60 * 2)),
             max_idle: 5,
             max_sockets: 1024,
+            dns_workers: 4,
         }
     }
 }