Merge pull request #88 from rdiomar/rdiomar_changes

Various changes/fixes, including:
* Allow customizing socket timeouts
* Read the correct number of bytes from kafka
* Guarantee reading the expected number of bytes from the socket every time
* Remove bufsize from client and conn
* SimpleConsumer flow changes
* Fix some error handling
* Add optional upper limit to consumer fetch buffer size
* Add and fix unit and integration tests

Author: rdiomar

kafka/client.py

@@ -3,11 +3,12 @@
 from functools import partial
 from itertools import count
 import logging
-import socket
 import time
 
-from kafka.common import ErrorMapping, TopicAndPartition
-from kafka.common import ConnectionError, FailedPayloadsException
+from kafka.common import (
+    ErrorMapping, TopicAndPartition, ConnectionError,
+    FailedPayloadsException
+)
 from kafka.conn import KafkaConnection
 from kafka.protocol import KafkaProtocol
 
@@ -19,12 +20,12 @@ class KafkaClient(object):
     CLIENT_ID = "kafka-python"
     ID_GEN = count()
 
-    def __init__(self, host, port, bufsize=4096, client_id=CLIENT_ID):
+    def __init__(self, host, port, client_id=CLIENT_ID, timeout=10):
         # We need one connection to bootstrap
-        self.bufsize = bufsize
         self.client_id = client_id
+        self.timeout = timeout
         self.conns = {               # (host, port) -> KafkaConnection
-            (host, port): KafkaConnection(host, port, bufsize)
+            (host, port): KafkaConnection(host, port, timeout=timeout)
         }
         self.brokers = {}            # broker_id -> BrokerMetadata
         self.topics_to_brokers = {}  # topic_id -> broker_id
@@ -41,7 +42,7 @@ def _get_conn_for_broker(self, broker):
         """
         if (broker.host, broker.port) not in self.conns:
             self.conns[(broker.host, broker.port)] = \
-                KafkaConnection(broker.host, broker.port, self.bufsize)
+                KafkaConnection(broker.host, broker.port, timeout=self.timeout)
 
         return self.conns[(broker.host, broker.port)]
 
@@ -165,14 +166,24 @@ def _send_broker_aware_request(self, payloads, encoder_fn, decoder_fn):
             request = encoder_fn(client_id=self.client_id,
                                  correlation_id=requestId, payloads=payloads)
 
+            failed = False
             # Send the request, recv the response
             try:
                 conn.send(requestId, request)
                 if decoder_fn is None:
                     continue
-                response = conn.recv(requestId)
-            except ConnectionError, e:  # ignore BufferUnderflow for now
-                log.warning("Could not send request [%s] to server %s: %s" % (request, conn, e))
+                try:
+                    response = conn.recv(requestId)
+                except ConnectionError, e:
+                    log.warning("Could not receive response to request [%s] "
+                                "from server %s: %s", request, conn, e)
+                    failed = True
+            except ConnectionError, e:
+                log.warning("Could not send request [%s] to server %s: %s",
+                            request, conn, e)
+                failed = True
+
+            if failed:
                 failed_payloads += payloads
                 self.topics_to_brokers = {} # reset metadata
                 continue

kafka/conn.py

@@ -4,7 +4,6 @@
 import struct
 from threading import local
 
-from kafka.common import BufferUnderflowError
 from kafka.common import ConnectionError
 
 log = logging.getLogger("kafka")
@@ -19,14 +18,14 @@ class KafkaConnection(local):
     we can do something in here to facilitate multiplexed requests/responses
     since the Kafka API includes a correlation id.
     """
-    def __init__(self, host, port, bufsize=4096):
+    def __init__(self, host, port, timeout=10):
         super(KafkaConnection, self).__init__()
         self.host = host
         self.port = port
-        self.bufsize = bufsize
         self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self._sock.connect((host, port))
-        self._sock.settimeout(10)
+        self.timeout = timeout
+        self._sock.settimeout(self.timeout)
         self._dirty = False
 
     def __str__(self):
@@ -36,44 +35,31 @@ def __str__(self):
     #   Private API   #
     ###################
 
-    def _consume_response(self):
-        """
-        Fully consume the response iterator
-        """
-        return "".join(self._consume_response_iter())
-
-    def _consume_response_iter(self):
-        """
-        This method handles the response header and error messages. It
-        then returns an iterator for the chunks of the response
-        """
-        log.debug("Handling response from Kafka")
-
-        # Read the size off of the header
-        resp = self._sock.recv(4)
-        if resp == "":
-            self._raise_connection_error()
-        (size,) = struct.unpack('>i', resp)
-
-        messagesize = size - 4
-        log.debug("About to read %d bytes from Kafka", messagesize)
-
-        # Read the remainder of the response
-        total = 0
-        while total < messagesize:
-            resp = self._sock.recv(self.bufsize)
-            log.debug("Read %d bytes from Kafka", len(resp))
-            if resp == "":
-                raise BufferUnderflowError(
-                    "Not enough data to read this response")
-
-            total += len(resp)
-            yield resp
-
     def _raise_connection_error(self):
         self._dirty = True
         raise ConnectionError("Kafka @ {}:{} went away".format(self.host, self.port))
 
+    def _read_bytes(self, num_bytes):
+        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)
+            except socket.error:
+                log.exception('Unable to receive data from Kafka')
+                self._raise_connection_error()
+            if data == '':
+                log.error("Not enough data to read this response")
+                self._raise_connection_error()
+            bytes_left -= len(data)
+            log.debug("Read %d/%d bytes from Kafka", num_bytes - bytes_left, num_bytes)
+            resp += data
+
+        return resp
+
     ##################
     #   Public API   #
     ##################
@@ -89,7 +75,7 @@ def send(self, request_id, payload):
             sent = self._sock.sendall(payload)
             if sent is not None:
                 self._raise_connection_error()
-        except socket.error:
+        except socket.error, e:
             log.exception('Unable to send payload to Kafka')
             self._raise_connection_error()
 
@@ -98,8 +84,14 @@ def recv(self, request_id):
         Get a response from Kafka
         """
         log.debug("Reading response %d from Kafka" % request_id)
-        self.data = self._consume_response()
-        return self.data
+        # Read the size off of the header
+        resp = self._read_bytes(4)
+
+        (size,) = struct.unpack('>i', resp)
+
+        # Read the remainder of the response
+        resp = self._read_bytes(size)
+        return str(resp)
 
     def copy(self):
         """
@@ -124,5 +116,5 @@ def reinit(self):
         self.close()
         self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self._sock.connect((self.host, self.port))
-        self._sock.settimeout(10)
+        self._sock.settimeout(self.timeout)
         self._dirty = False