Support client.poll with future and timeout_ms

Author: dpkp

kafka/client_async.py

@@ -636,11 +636,14 @@ def poll(self, timeout_ms=None, future=None):
         Returns:
             list: responses received (can be empty)
         """
-        if timeout_ms is None:
-            timeout_ms = self.config['request_timeout_ms']
-        elif not isinstance(timeout_ms, (int, float)):
+        if not isinstance(timeout_ms, (int, float, type(None))):
             raise TypeError('Invalid type for timeout: %s' % type(timeout_ms))
 
+        begin = time.time()
+        if timeout_ms is not None:
+            timeout_at = begin + (timeout_ms / 1000)
+        else:
+            timeout_at = begin + (self.config['request_timeout_ms'] / 1000)
         # Loop for futures, break after first loop if None
         responses = []
         while True:
@@ -665,11 +668,12 @@ def poll(self, timeout_ms=None, future=None):
                 if future is not None and future.is_done:
                     timeout = 0
                 else:
+                    user_timeout_ms = 1000 * max(0, timeout_at - time.time())
                     idle_connection_timeout_ms = self._idle_expiry_manager.next_check_ms()
                     request_timeout_ms = self._next_ifr_request_timeout_ms()
-                    log.debug("Timeouts: user %f, metadata %f, idle connection %f, request %f", timeout_ms, metadata_timeout_ms, idle_connection_timeout_ms, request_timeout_ms)
+                    log.debug("Timeouts: user %f, metadata %f, idle connection %f, request %f", user_timeout_ms, metadata_timeout_ms, idle_connection_timeout_ms, request_timeout_ms)
                     timeout = min(
-                        timeout_ms,
+                        user_timeout_ms,
                         metadata_timeout_ms,
                         idle_connection_timeout_ms,
                         request_timeout_ms)
@@ -683,7 +687,11 @@ def poll(self, timeout_ms=None, future=None):
 
             # If all we had was a timeout (future is None) - only do one poll
             # If we do have a future, we keep looping until it is done
-            if future is None or future.is_done:
+            if future is None:
+                break
+            elif future.is_done:
+                break
+            elif timeout_ms is not None and time.time() >= timeout_at:
                 break
 
         return responses

kafka/consumer/fetcher.py

@@ -276,6 +276,10 @@ def _retrieve_offsets(self, timestamps, timeout_ms=None):
             future = self._send_list_offsets_requests(timestamps)
             self._client.poll(future=future, timeout_ms=inner_timeout_ms())
 
+            # Timeout w/o future completion
+            if not future.is_done:
+                break
+
             if future.succeeded():
                 return future.value
             if not future.retriable():

kafka/coordinator/base.py

@@ -260,12 +260,17 @@ def ensure_coordinator_ready(self, timeout_ms=None):
                 future = self.lookup_coordinator()
                 self._client.poll(future=future, timeout_ms=inner_timeout_ms())
 
+                if not future.is_done:
+                    raise Errors.KafkaTimeoutError()
+
                 if future.failed():
                     if future.retriable():
                         if getattr(future.exception, 'invalid_metadata', False):
                             log.debug('Requesting metadata for group coordinator request: %s', future.exception)
                             metadata_update = self._client.cluster.request_update()
                             self._client.poll(future=metadata_update, timeout_ms=inner_timeout_ms())
+                            if not metadata_update.is_done:
+                                raise Errors.KafkaTimeoutError()
                         else:
                             time.sleep(inner_timeout_ms(self.config['retry_backoff_ms']) / 1000)
                     else:
@@ -416,6 +421,9 @@ def ensure_active_group(self, timeout_ms=None):
 
                 self._client.poll(future=future, timeout_ms=inner_timeout_ms())
 
+                if not future.is_done:
+                    raise Errors.KafkaTimeoutError()
+
                 if future.succeeded():
                     self._on_join_complete(self._generation.generation_id,
                                            self._generation.member_id,

test/fixtures.py

@@ -14,7 +14,7 @@
 from kafka.vendor.six.moves.urllib.parse import urlparse  # pylint: disable=E0611,F0401
 
 from kafka import errors, KafkaAdminClient, KafkaClient, KafkaConsumer, KafkaProducer
-from kafka.errors import InvalidReplicationFactorError
+from kafka.errors import InvalidReplicationFactorError, KafkaTimeoutError
 from kafka.protocol.admin import CreateTopicsRequest
 from kafka.protocol.metadata import MetadataRequest
 from test.testutil import env_kafka_version, random_string
@@ -555,6 +555,8 @@ def _failure(error):
                 future.error_on_callbacks = True
                 future.add_errback(_failure)
                 self._client.poll(future=future, timeout_ms=timeout)
+                if not future.is_done:
+                    raise KafkaTimeoutError()
                 return future.value
             except Exception as exc:
                 time.sleep(1)

test/test_client_async.py

@@ -228,6 +228,9 @@ def test_poll(mocker):
     ifr_request_timeout = mocker.patch.object(KafkaClient, '_next_ifr_request_timeout_ms')
     _poll = mocker.patch.object(KafkaClient, '_poll')
     cli = KafkaClient(api_version=(0, 9))
+    now = time.time()
+    t = mocker.patch('time.time')
+    t.return_value = now
 
     # metadata timeout wins
     ifr_request_timeout.return_value = float('inf')
@@ -346,17 +349,16 @@ def test_maybe_refresh_metadata_cant_send(mocker, client):
     t.return_value = now
 
     # first poll attempts connection
-    client.poll(timeout_ms=12345678)
-    client._poll.assert_called_with(12345.678)
+    client.poll()
+    client._poll.assert_called()
     client._init_connect.assert_called_once_with('foobar')
 
     # poll while connecting should not attempt a new connection
     client._connecting.add('foobar')
     client._can_connect.reset_mock()
-    client.poll(timeout_ms=12345678)
-    client._poll.assert_called_with(12345.678)
+    client.poll()
+    client._poll.assert_called()
     assert not client._can_connect.called
-
     assert not client._metadata_refresh_in_progress