Use kafka.util.Timer

Author: dpkp

kafka/client_async.py

@@ -27,7 +27,7 @@
 from kafka.metrics.stats.rate import TimeUnit
 from kafka.protocol.broker_api_versions import BROKER_API_VERSIONS
 from kafka.protocol.metadata import MetadataRequest
-from kafka.util import Dict, WeakMethod, ensure_valid_topic_name, timeout_ms_fn
+from kafka.util import Dict, Timer, WeakMethod, ensure_valid_topic_name
 # Although this looks unused, it actually monkey-patches socket.socketpair()
 # and should be left in as long as we're using socket.socketpair() in this file
 from kafka.vendor import socketpair # noqa: F401
@@ -646,11 +646,11 @@ def poll(self, timeout_ms=None, future=None):
         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)
+            timer = Timer(timeout_ms)
         else:
-            timeout_at = begin + (self.config['request_timeout_ms'] / 1000)
+            timer = Timer(self.config['request_timeout_ms'])
+
         # Loop for futures, break after first loop if None
         responses = []
         while True:
@@ -675,12 +675,11 @@ 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", user_timeout_ms, metadata_timeout_ms, idle_connection_timeout_ms, request_timeout_ms)
+                    log.debug("Timeouts: user %f, metadata %f, idle connection %f, request %f", timer.timeout_ms, metadata_timeout_ms, idle_connection_timeout_ms, request_timeout_ms)
                     timeout = min(
-                        user_timeout_ms,
+                        timer.timeout_ms,
                         metadata_timeout_ms,
                         idle_connection_timeout_ms,
                         request_timeout_ms)
@@ -698,7 +697,7 @@ def poll(self, timeout_ms=None, future=None):
                 break
             elif future.is_done:
                 break
-            elif timeout_ms is not None and time.time() >= timeout_at:
+            elif timeout_ms is not None and timer.expired:
                 break
 
         return responses
@@ -1175,16 +1174,16 @@ def await_ready(self, node_id, timeout_ms=30000):
         This method is useful for implementing blocking behaviour on top of the non-blocking `NetworkClient`, use it with
         care.
         """
-        inner_timeout_ms = timeout_ms_fn(timeout_ms, None)
+        timer = Timer(timeout_ms)
         self.poll(timeout_ms=0)
         if self.is_ready(node_id):
             return True
 
-        while not self.is_ready(node_id) and inner_timeout_ms() > 0:
+        while not self.is_ready(node_id) and not timer.expired:
             if self.connection_failed(node_id):
                 raise Errors.KafkaConnectionError("Connection to %s failed." % (node_id,))
             self.maybe_connect(node_id)
-            self.poll(timeout_ms=inner_timeout_ms())
+            self.poll(timeout_ms=timer.timeout_ms)
         return self.is_ready(node_id)
 
     def send_and_receive(self, node_id, request):

kafka/consumer/fetcher.py

@@ -19,7 +19,7 @@
 from kafka.record import MemoryRecords
 from kafka.serializer import Deserializer
 from kafka.structs import TopicPartition, OffsetAndMetadata, OffsetAndTimestamp
-from kafka.util import timeout_ms_fn
+from kafka.util import Timer
 
 log = logging.getLogger(__name__)
 
@@ -230,15 +230,15 @@ def _fetch_offsets_by_times(self, timestamps, timeout_ms=None):
         if not timestamps:
             return {}
 
-        inner_timeout_ms = timeout_ms_fn(timeout_ms, 'Timeout fetching offsets')
+        timer = Timer(timeout_ms, "Failed to get offsets by timestamps in %s ms" % (timeout_ms,))
         timestamps = copy.copy(timestamps)
         fetched_offsets = dict()
         while True:
             if not timestamps:
                 return {}
 
             future = self._send_list_offsets_requests(timestamps)
-            self._client.poll(future=future, timeout_ms=inner_timeout_ms())
+            self._client.poll(future=future, timeout_ms=timer.timeout_ms)
 
             # Timeout w/o future completion
             if not future.is_done:
@@ -256,12 +256,17 @@ def _fetch_offsets_by_times(self, timestamps, timeout_ms=None):
 
             if future.exception.invalid_metadata or self._client.cluster.need_update:
                 refresh_future = self._client.cluster.request_update()
-                self._client.poll(future=refresh_future, timeout_ms=inner_timeout_ms())
+                self._client.poll(future=refresh_future, timeout_ms=timer.timeout_ms)
 
                 if not future.is_done:
                     break
             else:
-                time.sleep(inner_timeout_ms(self.config['retry_backoff_ms']) / 1000)
+                if timer.timeout_ms is None or timer.timeout_ms > self.config['retry_backoff_ms']:
+                    time.sleep(self.config['retry_backoff_ms'] / 1000)
+                else:
+                    time.sleep(timer.timeout_ms / 1000)
+
+            timer.maybe_raise()
 
         raise Errors.KafkaTimeoutError(
             "Failed to get offsets by timestamps in %s ms" % (timeout_ms,))

kafka/coordinator/base.py

@@ -16,7 +16,7 @@
 from kafka.metrics.stats import Avg, Count, Max, Rate
 from kafka.protocol.find_coordinator import FindCoordinatorRequest
 from kafka.protocol.group import HeartbeatRequest, JoinGroupRequest, LeaveGroupRequest, SyncGroupRequest, DEFAULT_GENERATION_ID, UNKNOWN_MEMBER_ID
-from kafka.util import timeout_ms_fn, Timer
+from kafka.util import Timer
 
 log = logging.getLogger('kafka.coordinator')
 
@@ -293,7 +293,10 @@ def ensure_coordinator_ready(self, timeout_ms=None):
                             if not metadata_update.is_done:
                                 return False
                         else:
-                            time.sleep(min(timer.timeout_ms, self.config['retry_backoff_ms']) / 1000)
+                            if timeout_ms is None or timer.timeout_ms > self.config['retry_backoff_ms']:
+                                time.sleep(self.config['retry_backoff_ms'] / 1000)
+                            else:
+                                time.sleep(timer.timeout_ms / 1000)
                     else:
                         raise future.exception  # pylint: disable-msg=raising-bad-type
                 if timer.expired:
@@ -458,7 +461,8 @@ def join_group(self, timeout_ms=None):
             while not self.coordinator_unknown():
                 if not self._client.in_flight_request_count(self.coordinator_id):
                     break
-                self._client.poll(timeout_ms=min(timer.timeout_ms, 200))
+                poll_timeout_ms = 200 if timer.timeout_ms is None or timer.timeout_ms > 200 else timer.timeout_ms
+                self._client.poll(timeout_ms=poll_timeout_ms)
                 if timer.expired:
                     return False
             else:
@@ -491,7 +495,10 @@ def join_group(self, timeout_ms=None):
                 elif timer.expired:
                     return False
                 else:
-                    time.sleep(min(timer.timeout_ms, self.config['retry_backoff_ms']) / 1000)
+                    if timer.timeout_ms is None or timer.timeout_ms > self.config['retry_backoff_ms']:
+                        time.sleep(self.config['retry_backoff_ms'] / 1000)
+                    else:
+                        time.sleep(timer.timeout_ms / 1000)
 
     def _send_join_group_request(self):
         """Join the group and return the assignment for the next generation.

kafka/coordinator/consumer.py

@@ -19,7 +19,7 @@
 from kafka.metrics.stats import Avg, Count, Max, Rate
 from kafka.protocol.commit import OffsetCommitRequest, OffsetFetchRequest
 from kafka.structs import OffsetAndMetadata, TopicPartition
-from kafka.util import timeout_ms_fn, Timer, WeakMethod
+from kafka.util import Timer, WeakMethod
 
 
 log = logging.getLogger(__name__)
@@ -405,7 +405,7 @@ def refresh_committed_offsets_if_needed(self, timeout_ms=None):
         except Errors.KafkaTimeoutError:
             return False
         for partition, offset in six.iteritems(offsets):
-            log.debug("Setting offset for partition %s to the committed offset %s", partition, offset.offset);
+            log.debug("Setting offset for partition %s to the committed offset %s", partition, offset.offset)
             self._subscription.seek(partition, offset.offset)
         return True
 
@@ -424,32 +424,35 @@ def fetch_committed_offsets(self, partitions, timeout_ms=None):
         if not partitions:
             return {}
 
-        inner_timeout_ms = timeout_ms_fn(timeout_ms, None)
+        future_key = frozenset(partitions)
+        timer = Timer(timeout_ms)
         while True:
-            self.ensure_coordinator_ready(timeout_ms=inner_timeout_ms())
+            self.ensure_coordinator_ready(timeout_ms=timer.timeout_ms)
 
             # contact coordinator to fetch committed offsets
-            future_key = frozenset(partitions)
             if future_key in self._offset_fetch_futures:
                 future = self._offset_fetch_futures[future_key]
             else:
                 future = self._send_offset_fetch_request(partitions)
                 self._offset_fetch_futures[future_key] = future
 
-            self._client.poll(future=future, timeout_ms=inner_timeout_ms())
+            self._client.poll(future=future, timeout_ms=timer.timeout_ms)
 
-            if not future.is_done:
-                raise Errors.KafkaTimeoutError()
-            else:
+            if future.is_done:
                 del self._offset_fetch_futures[future_key]
 
-            if future.succeeded():
-                return future.value
+                if future.succeeded():
+                    return future.value
 
-            if not future.retriable():
-                raise future.exception # pylint: disable-msg=raising-bad-type
+                elif not future.retriable():
+                    raise future.exception # pylint: disable-msg=raising-bad-type
 
-            time.sleep(inner_timeout_ms(self.config['retry_backoff_ms']) / 1000)
+            # future failed but is retriable, or is not done yet
+            if timer.timeout_ms is None or timer.timeout_ms > self.config['retry_backoff_ms']:
+                time.sleep(self.config['retry_backoff_ms'] / 1000)
+            else:
+                time.sleep(timer.timeout_ms / 1000)
+            timer.maybe_raise()
 
     def close(self, autocommit=True, timeout_ms=None):
         """Close the coordinator, leave the current group,
@@ -540,23 +543,26 @@ def commit_offsets_sync(self, offsets, timeout_ms=None):
         if not offsets:
             return
 
-        inner_timeout_ms = timeout_ms_fn(timeout_ms, 'Timeout in coordinator.poll')
+        timer = Timer(timeout_ms)
         while True:
-            self.ensure_coordinator_ready(timeout_ms=inner_timeout_ms())
+            self.ensure_coordinator_ready(timeout_ms=timer.timeout_ms)
 
             future = self._send_offset_commit_request(offsets)
-            self._client.poll(future=future, timeout_ms=inner_timeout_ms())
-
-            if not future.is_done:
-                raise Errors.KafkaTimeoutError()
+            self._client.poll(future=future, timeout_ms=timer.timeout_ms)
 
-            if future.succeeded():
-                return future.value
+            if future.is_done:
+                if future.succeeded():
+                    return future.value
 
-            if not future.retriable():
-                raise future.exception # pylint: disable-msg=raising-bad-type
+                elif not future.retriable():
+                    raise future.exception # pylint: disable-msg=raising-bad-type
 
-            time.sleep(inner_timeout_ms(self.config['retry_backoff_ms']) / 1000)
+            # future failed but is retriable, or it is still pending
+            if timer.timeout_ms is None or timer.timeout_ms > self.config['retry_backoff_ms']:
+                time.sleep(self.config['retry_backoff_ms'] / 1000)
+            else:
+                time.sleep(timer.timeout_ms / 1000)
+            timer.maybe_raise()
 
     def _maybe_auto_commit_offsets_sync(self, timeout_ms=None):
         if self.config['enable_auto_commit']:

kafka/producer/kafka.py

@@ -5,7 +5,6 @@
 import logging
 import socket
 import threading
-import time
 import warnings
 import weakref
 
@@ -24,7 +23,7 @@
 from kafka.record.legacy_records import LegacyRecordBatchBuilder
 from kafka.serializer import Serializer
 from kafka.structs import TopicPartition
-from kafka.util import ensure_valid_topic_name
+from kafka.util import Timer, ensure_valid_topic_name
 
 
 log = logging.getLogger(__name__)
@@ -664,8 +663,7 @@ def __getattr__(self, name):
 
     def partitions_for(self, topic):
         """Returns set of all known partitions for the topic."""
-        max_wait = self.config['max_block_ms'] / 1000
-        return self._wait_on_metadata(topic, max_wait)
+        return self._wait_on_metadata(topic, self.config['max_block_ms'])
 
     @classmethod
     def max_usable_produce_magic(cls, api_version):
@@ -835,14 +833,11 @@ def send(self, topic, value=None, key=None, headers=None, partition=None, timest
         assert not (value is None and key is None), 'Need at least one: key or value'
         ensure_valid_topic_name(topic)
         key_bytes = value_bytes = None
+        timer = Timer(self.config['max_block_ms'], "Failed to assign partition for message in max_block_ms.")
         try:
             assigned_partition = None
-            elapsed = 0.0
-            begin = time.time()
-            timeout = self.config['max_block_ms'] / 1000
-            while assigned_partition is None and elapsed < timeout:
-                elapsed = time.time() - begin
-                self._wait_on_metadata(topic, timeout - elapsed)
+            while assigned_partition is None and not timer.expired:
+                self._wait_on_metadata(topic, timer.timeout_ms)
 
                 key_bytes = self._serialize(
                     self.config['key_serializer'],
@@ -856,7 +851,7 @@ def send(self, topic, value=None, key=None, headers=None, partition=None, timest
                 assigned_partition = self._partition(topic, partition, key, value,
                                                      key_bytes, value_bytes)
             if assigned_partition is None:
-                raise Errors.KafkaTimeoutError("Failed to assign partition for message after %s secs." % timeout)
+                raise Errors.KafkaTimeoutError("Failed to assign partition for message after %s secs." % timer.elapsed_ms / 1000)
             else:
                 partition = assigned_partition
 
@@ -931,7 +926,7 @@ def _ensure_valid_record_size(self, size):
                 " the maximum request size you have configured with the"
                 " max_request_size configuration" % (size,))
 
-    def _wait_on_metadata(self, topic, max_wait):
+    def _wait_on_metadata(self, topic, max_wait_ms):
         """
         Wait for cluster metadata including partitions for the given topic to
         be available.
@@ -949,36 +944,29 @@ def _wait_on_metadata(self, topic, max_wait):
         """
         # add topic to metadata topic list if it is not there already.
         self._sender.add_topic(topic)
-        begin = time.time()
-        elapsed = 0.0
+        timer = Timer(max_wait_ms, "Failed to update metadata after %.1f secs." % (max_wait_ms * 1000,))
         metadata_event = None
         while True:
             partitions = self._metadata.partitions_for_topic(topic)
             if partitions is not None:
                 return partitions
-
-            if elapsed >= max_wait:
-                raise Errors.KafkaTimeoutError(
-                    "Failed to update metadata after %.1f secs." % (max_wait,))
-
+            timer.maybe_raise()
             if not metadata_event:
                 metadata_event = threading.Event()
 
             log.debug("%s: Requesting metadata update for topic %s", str(self), topic)
-
             metadata_event.clear()
             future = self._metadata.request_update()
             future.add_both(lambda e, *args: e.set(), metadata_event)
             self._sender.wakeup()
-            metadata_event.wait(max_wait - elapsed)
+            metadata_event.wait(timer.timeout_ms / 1000)
             if not metadata_event.is_set():
                 raise Errors.KafkaTimeoutError(
-                    "Failed to update metadata after %.1f secs." % (max_wait,))
+                    "Failed to update metadata after %.1f secs." % (max_wait_ms * 1000,))
             elif topic in self._metadata.unauthorized_topics:
                 raise Errors.TopicAuthorizationFailedError(set([topic]))
             else:
-                elapsed = time.time() - begin
-                log.debug("%s: _wait_on_metadata woke after %s secs.", str(self), elapsed)
+                log.debug("%s: _wait_on_metadata woke after %s secs.", str(self), timer.elapsed_ms / 1000)
 
     def _serialize(self, f, topic, data):
         if not f: