Reorganize init_fetches calls during iteration

Dana Powers · Dana Powers · commit 458bdb50f62a · 2016-01-10T23:00:10.000-08:00
Generally should not init_fetches while the generator has pending
  messages; this revision adds an explicit check / noop to the
  public interface, and uses a private method internally to
  attempt to pipeline fetch requests.
diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py
@@ -86,9 +86,21 @@ def __init__(self, client, subscriptions, **configs):
     def init_fetches(self):
         """Send FetchRequests asynchronously for all assigned partitions.
 
+        Note: noop if there are unconsumed records internal to the fetcher
+
         Returns:
             List of Futures: each future resolves to a FetchResponse
         """
+        # We need to be careful when creating fetch records during iteration
+        # so we verify that there are no records in the deque, or in an
+        # iterator
+        if self._records or self._iterator:
+            log.debug('Skipping init_fetches because there are unconsumed'
+                      ' records internally')
+            return []
+        return self._init_fetches()
+
+    def _init_fetches(self):
         futures = []
         for node_id, request in six.iteritems(self._create_fetch_requests()):
             if self._client.ready(node_id):
@@ -339,6 +351,11 @@ def _message_generator(self):
             self._raise_if_unauthorized_topics()
             self._raise_if_record_too_large()
 
+            # Send additional FetchRequests when the internal queue is low
+            # this should enable moderate pipelining
+            if len(self._records) == 1:
+                self._init_fetches()
+
             (fetch_offset, tp, messages) = self._records.popleft()
 
             if not self._subscriptions.is_assigned(tp):
@@ -378,11 +395,6 @@ def _message_generator(self):
                 log.debug("Ignoring fetched records for %s at offset %s",
                           tp, fetch_offset)
 
-            # Send any additional FetchRequests that we can now
-            # this will likely fetch each partition individually, rather than
-            # fetch multiple partitions in bulk when they are on the same broker
-            self.init_fetches()
-
     def __iter__(self):  # pylint: disable=non-iterator-returned
         return self
 
diff --git a/kafka/consumer/group.py b/kafka/consumer/group.py
@@ -626,8 +626,6 @@ def _message_generator(self):
                 partitions = self._subscription.missing_fetch_positions()
                 self._update_fetch_positions(partitions)
 
-            # init any new fetches (won't resend pending fetches)
-            self._fetcher.init_fetches()
             self._client.poll(
                 max(0, self._consumer_timeout - time.time()) * 1000)
 
@@ -641,13 +639,16 @@ def _message_generator(self):
                 if time.time() > timeout_at:
                     log.debug("internal iterator timeout - breaking for poll")
                     break
+            else:
+                self._fetcher.init_fetches()
 
     def __iter__(self):  # pylint: disable=non-iterator-returned
         return self
 
     def __next__(self):
         if not self._iterator:
             self._iterator = self._message_generator()
+            self._fetcher.init_fetches()
 
         # consumer_timeout_ms can be used to stop iteration early
         if self.config['consumer_timeout_ms'] >= 0:
