Handle dirty flag in conn.recv()

* If the connection is dirty, reinit
* If we get a BufferUnderflowError, the server could have gone away, so mark it dirty

1 files changed, 3 insertions, 1 deletions

diff --git a/kafka/conn.py b/kafka/conn.py
index a8d6b1f..4749bdd 100644
--- a/kafka/conn.py
+++ b/kafka/conn.py
@@ -44,7 +44,8 @@ class KafkaConnection(local):
         bytes_left = num_bytes
         resp = ''
         log.debug("About to read %d bytes from Kafka", num_bytes)
-
+        if self._dirty:
+            self.reinit()
         while bytes_left:
             try:
                 data = self._sock.recv(bytes_left)
@@ -52,6 +53,7 @@ class KafkaConnection(local):
                 log.error('Unable to receive data from Kafka: %s', e)
                 self._raise_connection_error()
             if data == '':
+                self._dirty = True
                 raise BufferUnderflowError("Not enough data to read this response")
             bytes_left -= len(data)
             log.debug("Read %d/%d bytes from Kafka", num_bytes - bytes_left, num_bytes)