Merge pull request #671 from dpkp/disconnects

Improve socket disconnect handling

2 files changed, 25 insertions, 5 deletions

diff --git a/kafka/client_async.py b/kafka/client_async.py
index 6f5d1fe..7719426 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -142,6 +142,7 @@ class KafkaClient(object):
         # Exponential backoff if bootstrap fails
         backoff_ms = self.config['reconnect_backoff_ms'] * 2 ** self._bootstrap_fails
         next_at = self._last_bootstrap + backoff_ms / 1000.0
+        self._refresh_on_disconnects = False
         now = time.time()
         if next_at > now:
             log.debug("Sleeping %0.4f before bootstrapping again", next_at - now)
@@ -180,6 +181,7 @@ class KafkaClient(object):
             log.error('Unable to bootstrap from %s', hosts)
             # Max exponential backoff is 2^12, x4000 (50ms -> 200s)
             self._bootstrap_fails = min(self._bootstrap_fails + 1, 12)
+        self._refresh_on_disconnects = True
 
     def _can_connect(self, node_id):
         if node_id not in self._conns:
@@ -223,7 +225,7 @@ class KafkaClient(object):
             except KeyError:
                 pass
             if self._refresh_on_disconnects:
-                log.warning("Node %s connect failed -- refreshing metadata", node_id)
+                log.warning("Node %s connection failed -- refreshing metadata", node_id)
                 self.cluster.request_update()
 
     def _maybe_connect(self, node_id):
diff --git a/kafka/conn.py b/kafka/conn.py
index 3571e90..b5c7ba0 100644
--- a/kafka/conn.py
+++ b/kafka/conn.py
@@ -381,9 +381,17 @@ class BrokerConnection(object):
         # Not receiving is the state of reading the payload header
         if not self._receiving:
             try:
-                # An extremely small, but non-zero, probability that there are
-                # more than 0 but not yet 4 bytes available to read
-                self._rbuffer.write(self._sock.recv(4 - self._rbuffer.tell()))
+                bytes_to_read = 4 - self._rbuffer.tell()
+                data = self._sock.recv(bytes_to_read)
+                # We expect socket.recv to raise an exception if there is not
+                # enough data to read the full bytes_to_read
+                # but if the socket is disconnected, we will get empty data
+                # without an exception raised
+                if not data:
+                    log.error('%s: socket disconnected', self)
+                    self.close(error=Errors.ConnectionError('socket disconnected'))
+                    return None
+                self._rbuffer.write(data)
             except ssl.SSLWantReadError:
                 return None
             except ConnectionError as e:
@@ -411,7 +419,17 @@ class BrokerConnection(object):
         if self._receiving:
             staged_bytes = self._rbuffer.tell()
             try:
-                self._rbuffer.write(self._sock.recv(self._next_payload_bytes - staged_bytes))
+                bytes_to_read = self._next_payload_bytes - staged_bytes
+                data = self._sock.recv(bytes_to_read)
+                # We expect socket.recv to raise an exception if there is not
+                # enough data to read the full bytes_to_read
+                # but if the socket is disconnected, we will get empty data
+                # without an exception raised
+                if not data:
+                    log.error('%s: socket disconnected', self)
+                    self.close(error=Errors.ConnectionError('socket disconnected'))
+                    return None
+                self._rbuffer.write(data)
             except ssl.SSLWantReadError:
                 return None
             except ConnectionError as e: