Repository: kafka-rust/kafka-rust (language: rust)

Issue:
client: re-connect if connection closed by broker
As of Kafka 0.9, brokers will actively shutdown client connections which are idle for more than a broker configured timeout (the broker configuration parameter is `connections.max.idle.ms`; defaults to 10 minutes). `kafka-rust` should re-establish its broken connections to brokers.

Using `example/console-producer.rs` to send a message, leaving it untouched for more than 10 minutes, and then sending another message should not result in an error (currently `UnexpectedEOF`.)



Discussion / hints:

Fix patch:
diff --git a/examples/console-producer.rs b/examples/console-producer.rs
index f2ac9ba8d..0c3f692cc 100644
--- a/examples/console-producer.rs
+++ b/examples/console-producer.rs
@@ -7,10 +7,10 @@ use std::fs::File;
 use std::io::{self, stdin, stderr, Write, BufRead, BufReader};
 use std::ops::{Deref, DerefMut};
 
-use kafka::client::{KafkaClient, Compression, RequiredAcks};
+use kafka::client::{KafkaClient, Compression, RequiredAcks, DEFAULT_CONNECTION_IDLE_TIMEOUT};
 use kafka::producer::{AsBytes, Producer, Record};
 
-// how long do we allow wainting for the acknowledgement
+// how long do we allow waiting for the acknowledgement
 const ACK_TIMEOUT: i32 = 100;
 
 /// This is a very simple command line application sending every
@@ -37,7 +37,6 @@ fn main() {
 
 fn produce(cfg: &Config) -> Result<(), Error> {
     let mut client = KafkaClient::new(cfg.brokers.clone());
-    client.set_compression(cfg.compression);
     try!(client.load_metadata_all());
 
     // ~ verify that the remote brokers do know about the target topic
@@ -61,6 +60,8 @@ fn produce_impl(src: &mut BufRead, client: KafkaClient, cfg: &Config) -> Result<
     let mut producer = try!(Producer::from_client(client)
                             .with_ack_timeout(ACK_TIMEOUT)
                             .with_required_acks(cfg.required_acks)
+                            .with_compression(cfg.compression)
+                            .with_connection_idle_timeout(cfg.conn_idle_timeout)
                             .create());
     if cfg.batch_size < 2 {
         produce_impl_nobatch(&mut producer, src, cfg)
@@ -195,6 +196,7 @@ struct Config {
     compression: Compression,
     required_acks: RequiredAcks,
     batch_size: usize,
+    conn_idle_timeout: u32,
 }
 
 impl Config {
@@ -211,6 +213,7 @@ impl Config {
         opts.optopt("", "required-acks", "Specify amount of required broker acknowledgments \
                                           [NONE, ONE, ALL]", "TYPE");
         opts.optopt("", "batch-size", "Send N message in one batch.", "N");
+        opts.optopt("", "idle-timeout", "Specify timeout for idle connections", "SECS");
 
         let m = match opts.parse(&args[1..]) {
             Ok(m) => m,
@@ -256,6 +259,13 @@ impl Config {
                     }
                 }
             },
+            conn_idle_timeout: match m.opt_str("idle-timeout") {
+                None => DEFAULT_CONNECTION_IDLE_TIMEOUT,
+                Some(s) => match s.parse::<u32>() {
+                    Err(_) => return Err(Error::Literal(format!("Not a number: {}", s))),
+                    Ok(n) => n * 1_000,
+                },
+            }
         })
     }
 }
diff --git a/src/client/mod.rs b/src/client/mod.rs
index 780909a7a..93b44af5c 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -9,23 +9,23 @@ use std::io::Cursor;
 use std::iter::Iterator;
 use std::mem;
 use std::thread;
-use std::time::Duration;
-
-#[cfg(feature = "security")]
-use openssl::ssl::SslContext;
+use std::time::{Duration, Instant};
 
 // pub re-export
 pub use compression::Compression;
 pub use utils::PartitionOffset;
 pub use utils::TopicPartitionOffset;
 
+#[cfg(feature = "security")]
+pub use self::network::SecurityConfig;
+
 use codecs::{ToByte, FromByte};
-use connection::KafkaConnection;
 use error::{Result, Error, KafkaCode};
 use protocol::{self, ResponseParser};
 
 pub mod metadata;
 mod state;
+mod network;
 
 // ~ re-export (only) certain types from the protocol::fetch module as
 // 'client::fetch'.
@@ -36,7 +36,14 @@ pub mod fetch {
 }
 
 const CLIENTID: &'static str = "kafka-rust";
-const DEFAULT_SO_TIMEOUT_SECS: i32 = 120; // socket read, write timeout seconds
+const DEFAULT_CONNECTION_RW_TIMEOUT_SECS: u64 = 120;
+
+fn default_conn_rw_timeout() -> Option<Duration> {
+    match DEFAULT_CONNECTION_RW_TIMEOUT_SECS {
+        0 => None,
+        n => Some(Duration::from_secs(n)),
+    }
+}
 
 /// The default value for `KafkaClient::set_compression(..)`
 pub const DEFAULT_COMPRESSION: Compression = Compression::NONE;
@@ -64,6 +71,13 @@ pub const DEFAULT_RETRY_BACKOFF_TIME: u32 = 100;
 // approximetaly up to two minutes
 pub const DEFAULT_RETRY_MAX_ATTEMPTS: u32 = 120_000 / DEFAULT_RETRY_BACKOFF_TIME;
 
+/// The default value for `KafkaClient::set_connection_idle_timeout(..)`
+pub const DEFAULT_CONNECTION_IDLE_TIMEOUT: u32 = 540_000;
+
+fn default_conn_idle_timeout() -> Duration {
+    Duration::from_millis(DEFAULT_CONNECTION_IDLE_TIMEOUT as u64)
+}
+
 /// Client struct keeping track of brokers and topic metadata.
 ///
 /// Implements methods described by the [Kafka Protocol](http://kafka.apache.org/protocol.html).
@@ -75,7 +89,7 @@ pub struct KafkaClient {
     config: ClientConfig,
 
     // ~ a pool of re-usable connections to kafka brokers
-    conn_pool: ConnectionPool,
+    conn_pool: network::Connections,
 
     // ~ the current state of this client
     state: state::ClientState,
@@ -109,65 +123,6 @@ struct ClientConfig {
     retry_max_attempts: u32,
 }
 
-#[derive(Debug)]
-struct ConnectionPool {
-    conns: HashMap<String, KafkaConnection>,
-    timeout: i32,
-    #[cfg(feature = "security")]
-    security_config: Option<SecurityConfig>
-}
-
-impl ConnectionPool {
-    #[cfg(not(feature = "security"))]
-    fn new(timeout: i32) -> ConnectionPool {
-        ConnectionPool {
-            conns: HashMap::new(),
-            timeout: timeout,
-        }
-    }
-
-    #[cfg(feature = "security")]
-    fn new(timeout: i32) -> ConnectionPool {
-        Self::new_with_security(timeout, None)
-    }
-
-    #[cfg(feature = "security")]
-    fn new_with_security(timeout: i32, security: Option<SecurityConfig>) -> ConnectionPool {
-        ConnectionPool {
-            conns: HashMap::new(),
-            timeout: timeout,
-            security_config: security,
-        }
-    }
-
-    fn get_conn<'a>(&'a mut self, host: &str) -> Result<&'a mut KafkaConnection> {
-        if let Some(conn) = self.conns.get_mut(host) {
-            // ~ decouple the lifetimes to make borrowck happy; this
-            // is actually safe since we're immediatelly returning
-            // this, so the follow up code is not affected here (this
-            // method is no longer recursive).
-            return Ok(unsafe { mem::transmute(conn) });
-        }
-        let conn = try!(self.new_conn(host));
-        self.conns.insert(host.to_owned(), conn);
-        Ok(self.conns.get_mut(host).unwrap())
-    }
-
-    fn get_conn_any<'a>(&'a mut self) -> Option<&'a mut KafkaConnection> {
-        self.conns.iter_mut().next().map(|(_, conn)| conn)
-    }
-
-    #[cfg(not(feature = "security"))]
-    fn new_conn(&self, host: &str) -> Result<KafkaConnection> {
-        KafkaConnection::new(host, self.timeout)
-    }
-
-    #[cfg(feature = "security")]
-    fn new_conn(&self, host: &str) -> Result<KafkaConnection> {
-        KafkaConnection::new(host, self.timeout, self.security_config.as_ref().map(|c| &c.0))
-    }
-}
-
 // --------------------------------------------------------------------
 
 /// Possible values when querying a topic's offset.
@@ -384,22 +339,6 @@ impl<'a> AsRef<FetchPartition<'a>> for FetchPartition<'a> {
 
 // --------------------------------------------------------------------
 
-/// Security relevant configuration options for KafkaClient.
-// This will be expanded in the future. See #51.
-#[cfg(feature = "security")]
-#[derive(Debug)]
-pub struct SecurityConfig(SslContext);
-
-#[cfg(feature = "security")]
-impl SecurityConfig {
-    /// In the future this will also support a kerbos via #51.
-    pub fn new(ssl: SslContext) -> SecurityConfig {
-        SecurityConfig(ssl)
-    }
-}
-
-// --------------------------------------------------------------------
-
 impl KafkaClient {
 
     /// Creates a new instance of KafkaClient. Before being able to
@@ -426,7 +365,8 @@ impl KafkaClient {
                 retry_backoff_time: DEFAULT_RETRY_BACKOFF_TIME,
                 retry_max_attempts: DEFAULT_RETRY_MAX_ATTEMPTS,
             },
-            conn_pool: ConnectionPool::new(DEFAULT_SO_TIMEOUT_SECS),
+            conn_pool: network::Connections::new(
+                default_conn_rw_timeout(), default_conn_idle_timeout()),
             state: state::ClientState::new(),
         }
     }
@@ -482,7 +422,9 @@ impl KafkaClient {
                 retry_backoff_time: DEFAULT_RETRY_BACKOFF_TIME,
                 retry_max_attempts: DEFAULT_RETRY_MAX_ATTEMPTS,
             },
-            conn_pool: ConnectionPool::new_with_security(DEFAULT_SO_TIMEOUT_SECS, Some(security)),
+            conn_pool: network::Connections::new_with_security(
+                default_conn_rw_timeout(), default_conn_idle_timeout(),
+                Some(security)),
             state: state::ClientState::new(),
         }
     }
@@ -683,6 +625,26 @@ impl KafkaClient {
         self.config.retry_max_attempts
     }
 
+    /// Specifies the timeout (in number of milliseconds) after which
+    /// idle connections will transparently be closed/re-established
+    /// by `KafkaClient`.
+    ///
+    /// To be effective this value must be smaller than the [remote
+    /// broker's `connections.max.idle.ms`
+    /// setting](https://kafka.apache.org/documentation.html#brokerconfigs).
+    #[inline]
+    pub fn set_connection_idle_timeout(&mut self, millis: u32) {
+        self.conn_pool.set_idle_timeout(Duration::from_millis(millis as u64))
+    }
+
+    /// Retrieves the current
+    /// `KafkaClient::set_connection_idle_timeout` setting.
+    #[inline]
+    pub fn connection_idle_timeout(&self) -> u32 {
+        let d = self.conn_pool.idle_timeout();
+        d.as_secs() as u32 * 1_000 + d.subsec_nanos() / 1_000_000
+    }
+
     /// Provides a view onto the currently loaded metadata of known .
     ///
     /// # Examples
@@ -767,9 +729,10 @@ impl KafkaClient {
     /// underlying brokers (`self.hosts`).
     fn fetch_metadata<T: AsRef<str>>(&mut self, topics: &[T]) -> Result<protocol::MetadataResponse> {
         let correlation = self.state.next_correlation_id();
+        let now = Instant::now();
         for host in &self.config.hosts {
             debug!("fetch_metadata: requesting metadata from {}", host);
-            match self.conn_pool.get_conn(host) {
+            match self.conn_pool.get_conn(host, now) {
                 Ok(conn) => {
                     let req = protocol::MetadataRequest::new(correlation, &self.config.client_id, topics);
                     match __send_request(conn, req) {
@@ -824,9 +787,10 @@ impl KafkaClient {
         }
 
         // Call each broker with the request formed earlier
+        let now = Instant::now();
         let mut res: HashMap<String, Vec<PartitionOffset>> = HashMap::with_capacity(n_topics);
         for (host, req) in reqs {
-            let resp = try!(__send_receive::<_, protocol::OffsetResponse>(&mut self.conn_pool, &host, req));
+            let resp = try!(__send_receive::<_, protocol::OffsetResponse>(&mut self.conn_pool, &host, now, req));
             for tp in resp.topic_partitions {
                 let e = res.entry(tp.topic).or_insert(vec!());
                 for p in tp.partitions {
@@ -1170,8 +1134,9 @@ impl KafkaClient {
 
 fn __get_group_coordinator<'a>(group: &str,
                                state: &'a mut state::ClientState,
-                               conn_pool: &mut ConnectionPool,
-                               config: &ClientConfig)
+                               conn_pool: &mut network::Connections,
+                               config: &ClientConfig,
+                               now: Instant)
                                -> Result<&'a str>
 {
     if let Some(host) = state.group_coordinator(group) {
@@ -1185,11 +1150,11 @@ fn __get_group_coordinator<'a>(group: &str,
     let req = protocol::GroupCoordinatorRequest::new(group, correlation_id, &config.client_id);
     let mut attempt = 1;
     loop {
-        // ~ idealy we make make this work even if load_metadata has not
-        // been called yet; if there are not connections available we can
+        // ~ idealy we'd make this work even if `load_metadata` has not
+        // been called yet; if there are no connections available we can
         // try connecting to the user specified bootstrap server similar
-        // to the way load_metadata works.
-        let conn = conn_pool.get_conn_any().expect("available connection");
+        // to the way `load_metadata` works.
+        let conn = conn_pool.get_conn_any(now).expect("available connection");
         debug!("get_group_coordinator: asking for coordinator of '{}' on: {:?}", group, conn);
         let r = try!(__send_receive_conn::<_, protocol::GroupCoordinatorResponse>(conn, &req));
         let retry_code;
@@ -1217,17 +1182,19 @@ fn __get_group_coordinator<'a>(group: &str,
 
 fn __commit_offsets(req: protocol::OffsetCommitRequest,
                     state: &mut state::ClientState,
-                    conn_pool: &mut ConnectionPool,
+                    conn_pool: &mut network::Connections,
                     config: &ClientConfig)
                     -> Result<()>
 {
     let mut attempt = 1;
     loop {
+        let now = Instant::now();
+
         let tps = {
-            let host = try!(__get_group_coordinator(req.group, state, conn_pool, config));
+            let host = try!(__get_group_coordinator(req.group, state, conn_pool, config, now));
             debug!("commit_offsets: committing on protocol v{} for '{}' to: {}",
                    req.header.api_version, req.group, host);
-            try!(__send_receive::<_, protocol::OffsetCommitResponse>(conn_pool, host, &req))
+            try!(__send_receive::<_, protocol::OffsetCommitResponse>(conn_pool, host, now, &req))
                 .topic_partitions
         };
 
@@ -1274,16 +1241,18 @@ fn __commit_offsets(req: protocol::OffsetCommitRequest,
 // ~ 'num_partitions' is merely a hint to pre-allocation inside this method
 fn __fetch_group_offsets(
     req: protocol::OffsetFetchRequest, num_partitions: usize,
-    state: &mut state::ClientState, conn_pool: &mut ConnectionPool,
+    state: &mut state::ClientState, conn_pool: &mut network::Connections,
     config: &ClientConfig) -> Result<Vec<TopicPartitionOffset>>
 {
     let mut attempt = 1;
     loop {
+        let now = Instant::now();
+
         let r = {
-            let host = try!(__get_group_coordinator(req.group, state, conn_pool, config));
+            let host = try!(__get_group_coordinator(req.group, state, conn_pool, config, now));
             debug!("fetch_group_offsets: fetching on protocol v{} for '{}' from: {}",
                    req.header.api_version, req.group, host);
-            try!(__send_receive::<_, protocol::OffsetFetchResponse>(conn_pool, host, &req))
+            try!(__send_receive::<_, protocol::OffsetFetchResponse>(conn_pool, host, now, &req))
         };
 
         let mut retry_code = None;
@@ -1342,40 +1311,39 @@ fn __fetch_group_offsets(
 }
 
 /// ~ carries out the given fetch requests and returns the response
-fn __fetch_messages(conn_pool: &mut ConnectionPool,
+fn __fetch_messages(conn_pool: &mut network::Connections,
                     config: &ClientConfig,
                     reqs: HashMap<&str, protocol::FetchRequest>)
                     -> Result<Vec<fetch::Response>>
 {
-    // Call each broker with the request formed earlier
+    let now = Instant::now();
     let mut res = Vec::with_capacity(reqs.len());
     for (host, req) in reqs {
         let p = protocol::fetch::ResponseParser {
             validate_crc: config.fetch_crc_validation,
             requests: Some(&req),
         };
-
-        res.push(try!(__z_send_receive::<&protocol::FetchRequest, _>(conn_pool, host, &req, &p)));
+        res.push(try!(__z_send_receive(conn_pool, host, now, &req, &p)));
     }
     Ok(res)
 }
 
 /// ~ carries out the given produce requests and returns the response
-fn __produce_messages(conn_pool: &mut ConnectionPool,
+fn __produce_messages(conn_pool: &mut network::Connections,
                       reqs: HashMap<&str, protocol::ProduceRequest>,
                       no_acks: bool)
                       -> Result<Vec<TopicPartitionOffset>>
 {
-    // Call each broker with the request formed earlier
+    let now = Instant::now();
     if no_acks {
         for (host, req) in reqs {
-            try!(__send_noack::<_, protocol::ProduceResponse>(conn_pool, host, req));
+            try!(__send_noack::<_, protocol::ProduceResponse>(conn_pool, host, now, req));
         }
         Ok(vec!())
     } else {
         let mut res: Vec<TopicPartitionOffset> = vec![];
         for (host, req) in reqs {
-            let resp = try!(__send_receive::<_, protocol::ProduceResponse>(conn_pool, &host, req));
+            let resp = try!(__send_receive::<_, protocol::ProduceResponse>(conn_pool, &host, now, req));
             for tpo in resp.get_response() {
                 res.push(tpo);
             }
@@ -1384,29 +1352,29 @@ fn __produce_messages(conn_pool: &mut ConnectionPool,
     }
 }
 
-fn __send_receive<T, V>(conn_pool: &mut ConnectionPool, host: &str, req: T)
+fn __send_receive<T, V>(conn_pool: &mut network::Connections, host: &str, now: Instant, req: T)
                         -> Result<V::R>
     where T: ToByte, V: FromByte
 {
-    __send_receive_conn::<T, V>(try!(conn_pool.get_conn(host)), req)
+    __send_receive_conn::<T, V>(try!(conn_pool.get_conn(host, now)), req)
 }
 
-fn __send_receive_conn<T, V>(conn: &mut KafkaConnection, req: T) -> Result<V::R>
+fn __send_receive_conn<T, V>(conn: &mut network::KafkaConnection, req: T) -> Result<V::R>
     where T: ToByte, V: FromByte
 {
     try!(__send_request(conn, req));
     __get_response::<V>(conn)
 }
 
-fn __send_noack<T, V>(conn_pool: &mut ConnectionPool, host: &str, req: T)
+fn __send_noack<T, V>(conn_pool: &mut network::Connections, host: &str, now: Instant, req: T)
                       -> Result<usize>
     where T: ToByte, V: FromByte
 {
-    let mut conn = try!(conn_pool.get_conn(&host));
+    let mut conn = try!(conn_pool.get_conn(host, now));
     __send_request(&mut conn, req)
 }
 
-fn __send_request<T: ToByte>(conn: &mut KafkaConnection, request: T)
+fn __send_request<T: ToByte>(conn: &mut network::KafkaConnection, request: T)
                              -> Result<usize>
 {
     // ~ buffer to receive data to be sent
@@ -1422,7 +1390,7 @@ fn __send_request<T: ToByte>(conn: &mut KafkaConnection, request: T)
     conn.send(&buffer)
 }
 
-fn __get_response<T: FromByte>(conn: &mut KafkaConnection)
+fn __get_response<T: FromByte>(conn: &mut network::KafkaConnection)
                              -> Result<T::R>
 {
     let size = try!(__get_response_size(conn));
@@ -1443,16 +1411,16 @@ fn __get_response<T: FromByte>(conn: &mut KafkaConnection)
     T::decode_new(&mut Cursor::new(resp))
 }
 
-fn __z_send_receive<R, P>(conn_pool: &mut ConnectionPool, host: &str, req: R, parser: &P)
+fn __z_send_receive<R, P>(conn_pool: &mut network::Connections, host: &str, now: Instant, req: R, parser: &P)
                           -> Result<P::T>
     where R: ToByte, P: ResponseParser
 {
-    let mut conn = try!(conn_pool.get_conn(host));
+    let mut conn = try!(conn_pool.get_conn(host, now));
     try!(__send_request(&mut conn, req));
     __z_get_response(&mut conn, parser)
 }
 
-fn __z_get_response<P>(conn: &mut KafkaConnection, parser: &P) -> Result<P::T>
+fn __z_get_response<P>(conn: &mut network::KafkaConnection, parser: &P) -> Result<P::T>
     where P: ResponseParser
 {
     let size = try!(__get_response_size(conn));
@@ -1473,7 +1441,7 @@ fn __z_get_response<P>(conn: &mut KafkaConnection, parser: &P) -> Result<P::T>
     parser.parse(resp)
 }
 
-fn __get_response_size(conn: &mut KafkaConnection) -> Result<i32> {
+fn __get_response_size(conn: &mut network::KafkaConnection) -> Result<i32> {
     let mut buf = [0u8; 4];
     try!(conn.read_exact(&mut buf));
     i32::decode_new(&mut Cursor::new(&buf))
diff --git a/src/client/network.rs b/src/client/network.rs
new file mode 100644
index 000000000..40917ea8d
--- /dev/null
+++ b/src/client/network.rs
@@ -0,0 +1,365 @@
+//! Network related functionality for `KafkaClient`.
+//!
+//! This module is crate private and not exposed to the public except
+//! through re-exports of individual items from within
+//! `kafka::client`.
+
+use std::collections::HashMap;
+use std::fmt;
+use std::io::{Read, Write};
+use std::mem;
+use std::net::{TcpStream, Shutdown};
+use std::time::{Instant, Duration};
+
+#[cfg(feature = "security")]
+use openssl::ssl::{SslContext, SslStream};
+
+use error::Result;
+
+// --------------------------------------------------------------------
+
+/// Security relevant configuration options for `KafkaClient`.
+// This will be expanded in the future. See #51.
+#[cfg(feature = "security")]
+#[derive(Debug)]
+pub struct SecurityConfig(SslContext);
+
+#[cfg(feature = "security")]
+impl SecurityConfig {
+    /// In the future this will also support a kerbos via #51.
+    pub fn new(ssl: SslContext) -> SecurityConfig {
+        SecurityConfig(ssl)
+    }
+}
+
+// --------------------------------------------------------------------
+
+struct Pooled<T> {
+    last_checkout: Instant,
+    item: T,
+}
+
+impl<T> Pooled<T> {
+    fn new(last_checkout: Instant, item: T) -> Self {
+        Pooled { last_checkout: last_checkout, item: item }
+    }
+}
+
+impl<T: fmt::Debug> fmt::Debug for Pooled<T> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Pooled {{ last_checkout: {:?}, item: {:?} }}",
+               self.last_checkout, self.item)
+    }
+}
+
+#[derive(Debug)]
+pub struct Config {
+    rw_timeout: Option<Duration>,
+    idle_timeout: Duration,
+    #[cfg(feature = "security")]
+    security_config: Option<SecurityConfig>,
+}
+
+impl Config {
+    #[cfg(not(feature = "security"))]
+    fn new_conn(&self, id: u32, host: &str) -> Result<KafkaConnection> {
+        KafkaConnection::new(id, host, self.rw_timeout)
+            .map(|c| {
+                debug!("Established: {:?}", c);
+                c
+            })
+    }
+
+    #[cfg(feature = "security")]
+    fn new_conn(&self, id: u32, host: &str) -> Result<KafkaConnection> {
+        KafkaConnection::new(id,
+                             host,
+                             self.rw_timeout,
+                             self.security_config.as_ref().map(|c| &c.0))
+            .map(|c| {
+                debug!("Established: {:?}", c);
+                c
+            })
+    }
+}
+
+#[derive(Debug)]
+struct State {
+    num_conns: u32,
+}
+
+impl State {
+    fn new() -> State {
+        State { num_conns: 0 }
+    }
+
+    fn next_conn_id(&mut self) -> u32 {
+        let c = self.num_conns;
+        self.num_conns = self.num_conns.wrapping_add(1);
+        c
+    }
+}
+
+#[derive(Debug)]
+pub struct Connections {
+    conns: HashMap<String, Pooled<KafkaConnection>>,
+    state: State,
+    config: Config,
+}
+
+impl Connections {
+    #[cfg(not(feature = "security"))]
+    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration) -> Connections {
+        Connections {
+            conns: HashMap::new(),
+            state: State::new(),
+            config: Config {
+                rw_timeout: rw_timeout,
+                idle_timeout: idle_timeout,
+            },
+        }
+    }
+
+    #[cfg(feature = "security")]
+    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration) -> Connections {
+        Self::new_with_security(rw_timeout, idle_timeout, None)
+    }
+
+    #[cfg(feature = "security")]
+    pub fn new_with_security(rw_timeout: Option<Duration>,
+                             idle_timeout: Duration,
+                             security: Option<SecurityConfig>)
+                             -> Connections
+    {
+        Connections {
+            conns: HashMap::new(),
+            state: State::new(),
+            config: Config {
+                rw_timeout: rw_timeout,
+                idle_timeout: idle_timeout,
+                security_config: security,
+            },
+        }
+    }
+
+    pub fn set_idle_timeout(&mut self, idle_timeout: Duration) {
+        self.config.idle_timeout = idle_timeout;
+    }
+
+    pub fn idle_timeout(&self) -> Duration {
+        self.config.idle_timeout
+    }
+
+    pub fn get_conn<'a>(&'a mut self, host: &str, now: Instant)
+                        -> Result<&'a mut KafkaConnection>
+    {
+        if let Some(conn) = self.conns.get_mut(host) {
+            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
+                debug!("Idle timeout reached: {:?}", conn.item);
+                let new_conn = try!(self.config.new_conn(self.state.next_conn_id(), host));
+                let _ = conn.item.shutdown();
+                conn.item = new_conn;
+            }
+            conn.last_checkout = now;
+            let kconn: &mut KafkaConnection = &mut conn.item;
+            // ~ decouple the lifetimes to make the borrowck happy;
+            // this is safe since we're immediatelly returning the
+            // reference and the rest of the code in this method is
+            // not affected
+            return Ok(unsafe { mem::transmute(kconn) });
+        }
+        let cid = self.state.next_conn_id();
+        self.conns.insert(host.to_owned(),
+                          Pooled::new(now, try!(self.config.new_conn(cid, host))));
+        Ok(&mut self.conns.get_mut(host).unwrap().item)
+    }
+
+    pub fn get_conn_any(&mut self, now: Instant) -> Option<&mut KafkaConnection> {
+        for (host, conn) in &mut self.conns {
+            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
+                debug!("Idle timeout reached: {:?}", conn.item);
+                let new_conn_id = self.state.next_conn_id();
+                let new_conn = match self.config.new_conn(new_conn_id, host.as_str()) {
+                    Ok(new_conn) => {
+                        let _ = conn.item.shutdown();
+                        new_conn
+                    }
+                    Err(e) => {
+                        warn!("Failed to establish connection to {}: {:?}", host, e);
+                        continue;
+                    }
+                };
+                conn.item = new_conn;
+            }
+            conn.last_checkout = now;
+            let kconn: &mut KafkaConnection = &mut conn.item;
+            return Some(kconn);
+        }
+        None
+    }
+}
+
+// --------------------------------------------------------------------
+
+trait IsSecured {
+    fn is_secured(&self) -> bool;
+}
+
+#[cfg(not(feature = "security"))]
+type KafkaStream = TcpStream;
+
+#[cfg(not(feature = "security"))]
+impl IsSecured for KafkaStream {
+    fn is_secured(&self) -> bool {
+        false
+    }
+}
+
+#[cfg(feature = "security")]
+use self::openssled::KafkaStream;
+
+#[cfg(feature = "security")]
+mod openssled {
+    use std::io::{self, Read, Write};
+    use std::net::{TcpStream, Shutdown};
+    use std::time::Duration;
+
+    use openssl::ssl::SslStream;
+
+    use super::IsSecured;
+
+    pub enum KafkaStream {
+        Plain(TcpStream),
+        Ssl(SslStream<TcpStream>)
+    }
+
+    impl IsSecured for KafkaStream {
+        fn is_secured(&self) -> bool {
+            match *self {
+                KafkaStream::Ssl(_) => true,
+                _ => false,
+            }
+        }
+    }
+
+    impl KafkaStream {
+        fn get_ref(&self) -> &TcpStream {
+            match *self {
+                KafkaStream::Plain(ref s) => s,
+                KafkaStream::Ssl(ref s) => s.get_ref()
+            }
+        }
+
+        pub fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
+            self.get_ref().set_read_timeout(dur)
+        }
+
+        pub fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
+            self.get_ref().set_write_timeout(dur)
+        }
+
+        pub fn shutdown(&mut self, how: Shutdown) -> io::Result<()> {
+            self.get_ref().shutdown(how)
+        }
+    }
+
+    impl Read for KafkaStream {
+        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+            match *self {
+                KafkaStream::Plain(ref mut s) => s.read(buf),
+                KafkaStream::Ssl(ref mut s) => s.read(buf),
+            }
+        }
+    }
+
+    impl Write for KafkaStream {
+        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+            match *self {
+                KafkaStream::Plain(ref mut s) => s.write(buf),
+                KafkaStream::Ssl(ref mut s) => s.write(buf),
+            }
+        }
+        fn flush(&mut self) -> io::Result<()> {
+            match *self {
+                KafkaStream::Plain(ref mut s) => s.flush(),
+                KafkaStream::Ssl(ref mut s) => s.flush(),
+            }
+        }
+    }
+}
+
+/// A TCP stream to a remote Kafka broker.
+pub struct KafkaConnection {
+    // a surrogate identifier to distinguish between
+    // connections to the same host in debug messages
+    id: u32,
+    // "host:port"
+    host: String,
+    // the (wrapped) tcp stream
+    stream: KafkaStream,
+}
+
+impl fmt::Debug for KafkaConnection {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "KafkaConnection {{ id: {}, secured: {}, host: \"{}\" }}",
+               self.id, self.stream.is_secured(), self.host)
+    }
+}
+
+impl KafkaConnection {
+
+    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
+        let r = self.stream.write(&msg[..]).map_err(From::from);
+        trace!("Sent {} bytes to: {:?} => {:?}", msg.len(), self, r);
+        r
+    }
+
+    pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
+        let r = (&mut self.stream).read_exact(buf).map_err(From::from);
+        trace!("Read {} bytes from: {:?} => {:?}", buf.len(), self, r);
+        r
+    }
+
+    pub fn read_exact_alloc(&mut self, size: u64) -> Result<Vec<u8>> {
+        let size: usize = size as usize;
+        let mut buffer: Vec<u8> = Vec::with_capacity(size);
+        // this is safe actually: we are setting the len to the
+        // buffers capacity and either fully populate it in the
+        // following call to `read_exact` or discard the vector (in
+        // the error case)
+        unsafe { buffer.set_len(size) };
+        try!(self.read_exact(buffer.as_mut_slice()));
+        Ok(buffer)
+    }
+
+    fn shutdown(&mut self) -> Result<()> {
+        let r = self.stream.shutdown(Shutdown::Both);
+        debug!("Shut down: {:?} => {:?}", self, r);
+        r.map_err(From::from)
+    }
+
+    fn from_stream(stream: KafkaStream, id: u32, host: &str, rw_timeout: Option<Duration>)
+                   -> Result<KafkaConnection>
+    {
+        try!(stream.set_read_timeout(rw_timeout));
+        try!(stream.set_write_timeout(rw_timeout));
+        Ok(KafkaConnection { id: id, host: host.to_owned(), stream: stream })
+    }
+
+    #[cfg(not(feature = "security"))]
+    fn new(id: u32, host: &str, rw_timeout: Option<Duration>) -> Result<KafkaConnection> {
+        KafkaConnection::from_stream(try!(TcpStream::connect(host)), id, host, rw_timeout)
+    }
+
+    #[cfg(feature = "security")]
+    fn new(id: u32, host: &str, rw_timeout: Option<Duration>, security: Option<&SslContext>)
+           -> Result<KafkaConnection>
+    {
+        let stream = try!(TcpStream::connect(host));
+        let stream = match security {
+            Some(ctx) => KafkaStream::Ssl(try!(SslStream::connect(ctx, stream))),
+            None => KafkaStream::Plain(stream),
+        };
+        KafkaConnection::from_stream(stream, id, host, rw_timeout)
+    }
+}
diff --git a/src/connection.rs b/src/connection.rs
deleted file mode 100644
index 35d9e0a79..000000000
--- a/src/connection.rs
+++ /dev/null
@@ -1,128 +0,0 @@
-use std::fmt;
-use std::io::{Read, Write};
-use std::net::TcpStream;
-use std::time::Duration;
-
-#[cfg(feature = "security")]
-use openssl::ssl::{SslContext, SslStream};
-
-use error::Result;
-
-#[cfg(not(feature = "security"))]
-type KafkaStream = TcpStream;
-
-#[cfg(feature = "security")]
-use self::openssled::KafkaStream;
-
-#[cfg(feature = "security")]
-mod openssled {
-    use std::io::{self, Read, Write};
-    use std::net::TcpStream;
-    use std::time::Duration;
-
-    use openssl::ssl::SslStream;
-
-    pub enum KafkaStream {
-        Plain(TcpStream),
-        Ssl(SslStream<TcpStream>)
-    }
-
-    impl KafkaStream {
-        fn get_ref(&self) -> &TcpStream {
-            match self {
-                &KafkaStream::Plain(ref s) => s,
-                &KafkaStream::Ssl(ref s) => s.get_ref()
-            }
-        }
-
-        pub fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
-            self.get_ref().set_read_timeout(dur)
-        }
-
-        pub fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
-            self.get_ref().set_write_timeout(dur)
-        }
-    }
-
-    impl Read for KafkaStream {
-        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-            match self {
-                &mut KafkaStream::Plain(ref mut s) => s.read(buf),
-                &mut KafkaStream::Ssl(ref mut s) => s.read(buf),
-            }
-        }
-    }
-
-    impl Write for KafkaStream {
-        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-            match self {
-                &mut KafkaStream::Plain(ref mut s) => s.write(buf),
-                &mut KafkaStream::Ssl(ref mut s) => s.write(buf),
-            }
-        }
-        fn flush(&mut self) -> io::Result<()> {
-            match self {
-                &mut KafkaStream::Plain(ref mut s) => s.flush(),
-                &mut KafkaStream::Ssl(ref mut s) => s.flush(),
-            }
-        }
-    }
-}
-
-pub struct KafkaConnection {
-    host: String,
-    stream: KafkaStream,
-}
-
-impl fmt::Debug for KafkaConnection {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "KafkaConnection {{ host: \"{}\" }}", self.host)
-    }
-}
-
-impl KafkaConnection {
-
-    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
-        self.stream.write(&msg[..]).map_err(From::from)
-    }
-
-    pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
-        (&mut self.stream).read_exact(buf).map_err(From::from)
-    }
-
-    pub fn read_exact_alloc(&mut self, size: u64) -> Result<Vec<u8>> {
-        let size: usize = size as usize;
-        let mut buffer: Vec<u8> = Vec::with_capacity(size);
-        // this is safe actually: we are setting the len to the
-        // buffers capacity and either fully populate it in the
-        // following call to `read_exact` or discard the vector (in
-        // the error case)
-        unsafe { buffer.set_len(size) };
-        try!((&mut self.stream).read_exact(buffer.as_mut_slice()));
-        Ok(buffer)
-    }
-
-    fn from_stream(stream: KafkaStream, host: &str, timeout_secs: i32) -> KafkaConnection {
-        if timeout_secs > 0 {
-            let t = Some(Duration::from_secs(timeout_secs as u64));
-            stream.set_read_timeout(t).expect("Set connection read-timeout");
-            stream.set_write_timeout(t).expect("Set connection write-timeout");
-        }
-        KafkaConnection{host: host.to_owned(), stream: stream}
-    }
-
-    #[cfg(not(feature = "security"))]
-    pub fn new(host: &str, timeout_secs: i32) -> Result<KafkaConnection> {
-        Ok(KafkaConnection::from_stream(try!(TcpStream::connect(host)), host, timeout_secs))
-    }
-
-    #[cfg(feature = "security")]
-    pub fn new(host: &str, timeout_secs: i32, security: Option<&SslContext>) -> Result<KafkaConnection> {
-        let stream = try!(TcpStream::connect(host));
-        let stream = match security {
-            Some(ctx) => KafkaStream::Ssl(try!(SslStream::connect(ctx, stream))),
-            None => KafkaStream::Plain(stream),
-        };
-        Ok(KafkaConnection::from_stream(stream, host, timeout_secs))
-    }
-}
diff --git a/src/consumer/builder.rs b/src/consumer/builder.rs
index 7bdb88e36..a2e8689b4 100644
--- a/src/consumer/builder.rs
+++ b/src/consumer/builder.rs
@@ -30,6 +30,7 @@ pub struct Builder {
     fetch_crc_validation: bool,
     security_config: Option<SecurityConfig>,
     group_offset_storage: GroupOffsetStorage,
+    conn_idle_timeout: u32,
 }
 
 // ~ public only to be shared inside the kafka crate; not supposed to
@@ -48,6 +49,7 @@ pub fn new(client: Option<KafkaClient>, hosts: Vec<String>) -> Builder {
         fallback_offset: DEFAULT_FALLBACK_OFFSET,
         security_config: None,
         group_offset_storage: client::DEFAULT_GROUP_OFFSET_STORAGE,
+        conn_idle_timeout: client::DEFAULT_CONNECTION_IDLE_TIMEOUT,
     };
     if let Some(ref c) = b.client {
         b.fetch_max_wait_time = c.fetch_max_wait_time();
@@ -55,6 +57,7 @@ pub fn new(client: Option<KafkaClient>, hosts: Vec<String>) -> Builder {
         b.fetch_max_bytes_per_partition = c.fetch_max_bytes_per_partition();
         b.fetch_crc_validation = c.fetch_crc_validation();
         b.group_offset_storage = c.group_offset_storage();
+        b.conn_idle_timeout = c.connection_idle_timeout();
     }
     b
 }
@@ -180,6 +183,13 @@ impl Builder {
         self
     }
 
+    /// Specifies the timeout for idle connections.
+    /// See `KafkaClient::set_connection_idle_timeout`.
+    pub fn with_connection_idle_timeout(mut self, timeout: u32) -> Self {
+        self.conn_idle_timeout = timeout;
+        self
+    }
+
     #[cfg(not(feature = "security"))]
     fn new_kafka_client(hosts: Vec<String>, _: Option<SecurityConfig>) -> KafkaClient {
         KafkaClient::new(hosts)
@@ -218,6 +228,7 @@ impl Builder {
         client.set_fetch_min_bytes(self.fetch_min_bytes);
         client.set_fetch_max_bytes_per_partition(self.fetch_max_bytes_per_partition);
         client.set_group_offset_storage(self.group_offset_storage);
+        client.set_connection_idle_timeout(self.conn_idle_timeout);
 
         let config = Config {
             group: self.group,
diff --git a/src/lib.rs b/src/lib.rs
index c2f13786c..e13d96c8c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -35,7 +35,6 @@ pub mod consumer;
 pub mod producer;
 mod utils;
 mod codecs;
-mod connection;
 mod protocol;
 mod compression;
 
diff --git a/src/producer.rs b/src/producer.rs
index 307e02099..dddac7a42 100644
--- a/src/producer.rs
+++ b/src/producer.rs
@@ -324,6 +324,7 @@ pub struct Builder<P = DefaultPartitioner> {
     hosts: Vec<String>,
     compression: Compression,
     ack_timeout: i32,
+    conn_idle_timeout: u32,
     required_acks: RequiredAcks,
     partitioner: P,
     security_config: Option<SecurityConfig>,
@@ -336,12 +337,14 @@ impl Builder {
             hosts: hosts,
             compression: client::DEFAULT_COMPRESSION,
             ack_timeout: DEFAULT_ACK_TIMEOUT,
+            conn_idle_timeout: client::DEFAULT_CONNECTION_IDLE_TIMEOUT,
             required_acks: DEFAULT_REQUIRED_ACKS,
             partitioner: DefaultPartitioner::default(),
             security_config: None,
         };
         if let Some(ref c) = b.client {
             b.compression = c.compression();
+            b.conn_idle_timeout = c.connection_idle_timeout();
         }
         b
     }
@@ -371,6 +374,13 @@ impl Builder {
         self
     }
 
+    /// Specifies the timeout for idle connections.
+    /// See `KafkaClient::set_connection_idle_timeout`.
+    pub fn with_connection_idle_timeout(mut self, timeout: u32) -> Self {
+        self.conn_idle_timeout = timeout;
+        self
+    }
+
     /// Sets how many acknowledgements the kafka brokers should
     /// receive before responding to sent messages.
     ///
@@ -390,13 +400,13 @@ impl<P> Builder<P> {
             hosts: self.hosts,
             compression: self.compression,
             ack_timeout: self.ack_timeout,
+            conn_idle_timeout: self.conn_idle_timeout,
             required_acks: self.required_acks,
             partitioner: partitioner,
             security_config: None,
         }
     }
 
-
     #[cfg(not(feature = "security"))]
     fn new_kafka_client(hosts: Vec<String>, _: Option<SecurityConfig>) -> KafkaClient {
         KafkaClient::new(hosts)
@@ -423,6 +433,7 @@ impl<P> Builder<P> {
             }
         };
         client.set_compression(self.compression);
+        client.set_connection_idle_timeout(self.conn_idle_timeout);
         let state = try!(State::new(&mut client, self.partitioner));
         let config = Config {
             ack_timeout: self.ack_timeout,


Test patch: