Repository: dpkp/kafka-python
Pull request #2498 (2025-02-26 19:42:12)

Support custom per-request timeouts


Patch:
diff --git a/kafka/client_async.py b/kafka/client_async.py
index 057e8ae..96959d9 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -546,7 +546,7 @@ class KafkaClient(object):
             return False
         return conn.connected() and conn.can_send_more()
 
-    def send(self, node_id, request, wakeup=True):
+    def send(self, node_id, request, wakeup=True, request_timeout_ms=None):
         """Send a request to a specific node. Bytes are placed on an
         internal per-connection send-queue. Actual network I/O will be
         triggered in a subsequent call to .poll()
@@ -554,7 +554,13 @@ class KafkaClient(object):
         Arguments:
             node_id (int): destination node
             request (Struct): request object (not-encoded)
-            wakeup (bool): optional flag to disable thread-wakeup
+
+        Keyword Arguments:
+            wakeup (bool, optional): optional flag to disable thread-wakeup.
+            request_timeout_ms (int, optional): Provide custom timeout in milliseconds.
+                If response is not processed before timeout, client will fail the
+                request and close the connection.
+                Default: None (uses value from client configuration)
 
         Raises:
             AssertionError: if node_id is not in current cluster metadata
@@ -570,7 +576,7 @@ class KafkaClient(object):
         # conn.send will queue the request internally
         # we will need to call send_pending_requests()
         # to trigger network I/O
-        future = conn.send(request, blocking=False)
+        future = conn.send(request, blocking=False, request_timeout_ms=request_timeout_ms)
         if not future.is_done:
             self._sending.add(conn)
 
@@ -729,11 +735,13 @@ class KafkaClient(object):
 
         for conn in six.itervalues(self._conns):
             if conn.requests_timed_out():
+                timed_out = conn.timed_out_ifrs()
+                timeout_ms = (timed_out[0][2] - timed_out[0][1]) * 1000
                 log.warning('%s timed out after %s ms. Closing connection.',
-                            conn, conn.config['request_timeout_ms'])
+                            conn, timeout_ms)
                 conn.close(error=Errors.RequestTimedOutError(
                     'Request timed out after %s ms' %
-                    conn.config['request_timeout_ms']))
+                    timeout_ms))
 
         if self._sensors:
             self._sensors.io_time.record((time.time() - end_select) * 1000000000)
diff --git a/kafka/conn.py b/kafka/conn.py
index 2a4f1df..347e500 100644
--- a/kafka/conn.py
+++ b/kafka/conn.py
@@ -948,7 +948,7 @@ class BrokerConnection(object):
         # drop lock before state change callback and processing futures
         self.config['state_change_callback'](self.node_id, sock, self)
         sock.close()
-        for (_correlation_id, (future, _timestamp)) in ifrs:
+        for (_correlation_id, (future, _timestamp, _timeout)) in ifrs:
             future.failure(error)
 
     def _can_send_recv(self):
@@ -956,8 +956,20 @@ class BrokerConnection(object):
         return self.state in (ConnectionStates.AUTHENTICATING,
                               ConnectionStates.CONNECTED)
 
-    def send(self, request, blocking=True):
-        """Queue request for async network send, return Future()"""
+    def send(self, request, blocking=True, request_timeout_ms=None):
+        """Queue request for async network send, return Future()
+
+        Arguments:
+            request (Request): kafka protocol request object to send.
+
+        Keyword Arguments:
+            blocking (bool, optional): Whether to immediately send via
+                blocking socket I/O. Default: True.
+            request_timeout_ms: Custom timeout in milliseconds for request.
+                Default: None (uses value from connection configuration)
+
+        Returns: future
+        """
         future = Future()
         if self.connecting():
             return future.failure(Errors.NodeNotReadyError(str(self)))
@@ -965,9 +977,9 @@ class BrokerConnection(object):
             return future.failure(Errors.KafkaConnectionError(str(self)))
         elif not self.can_send_more():
             return future.failure(Errors.TooManyInFlightRequests(str(self)))
-        return self._send(request, blocking=blocking)
+        return self._send(request, blocking=blocking, request_timeout_ms=request_timeout_ms)
 
-    def _send(self, request, blocking=True):
+    def _send(self, request, blocking=True, request_timeout_ms=None):
         future = Future()
         with self._lock:
             if not self._can_send_recv():
@@ -980,9 +992,11 @@ class BrokerConnection(object):
 
             log.debug('%s Request %d: %s', self, correlation_id, request)
             if request.expect_response():
-                sent_time = time.time()
                 assert correlation_id not in self.in_flight_requests, 'Correlation ID already in-flight!'
-                self.in_flight_requests[correlation_id] = (future, sent_time)
+                sent_time = time.time()
+                request_timeout_ms = request_timeout_ms or self.config['request_timeout_ms']
+                timeout_at = sent_time + (request_timeout_ms / 1000)
+                self.in_flight_requests[correlation_id] = (future, sent_time, timeout_at)
             else:
                 future.success(None)
 
@@ -1061,18 +1075,20 @@ class BrokerConnection(object):
         """
         responses = self._recv()
         if not responses and self.requests_timed_out():
+            timed_out = self.timed_out_ifrs()
+            timeout_ms = (timed_out[0][2] - timed_out[0][1]) * 1000
             log.warning('%s timed out after %s ms. Closing connection.',
-                        self, self.config['request_timeout_ms'])
+                        self, timeout_ms)
             self.close(error=Errors.RequestTimedOutError(
                 'Request timed out after %s ms' %
-                self.config['request_timeout_ms']))
+                timeout_ms))
             return ()
 
         # augment responses w/ correlation_id, future, and timestamp
         for i, (correlation_id, response) in enumerate(responses):
             try:
                 with self._lock:
-                    (future, timestamp) = self.in_flight_requests.pop(correlation_id)
+                    (future, timestamp, _timeout) = self.in_flight_requests.pop(correlation_id)
             except KeyError:
                 self.close(Errors.KafkaConnectionError('Received unrecognized correlation id'))
                 return ()
@@ -1143,13 +1159,17 @@ class BrokerConnection(object):
     def requests_timed_out(self):
         return self.next_ifr_request_timeout_ms() == 0
 
+    def timed_out_ifrs(self):
+        now = time.time()
+        ifrs = sorted(self.in_flight_requests.values(), reverse=True, key=lambda ifr: ifr[2])
+        return list(filter(lambda ifr: ifr[2] <= now, ifrs))
+
     def next_ifr_request_timeout_ms(self):
         with self._lock:
             if self.in_flight_requests:
-                get_timestamp = lambda v: v[1]
-                oldest_at = min(map(get_timestamp,
-                                    self.in_flight_requests.values()))
-                next_timeout = oldest_at + self.config['request_timeout_ms'] / 1000.0
+                get_timeout = lambda v: v[2]
+                next_timeout = min(map(get_timeout,
+                                   self.in_flight_requests.values()))
                 return max(0, (next_timeout - time.time()) * 1000)
             else:
                 return float('inf')
diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py
index b544e4b..9dd4b84 100644
--- a/kafka/consumer/fetcher.py
+++ b/kafka/consumer/fetcher.py
@@ -456,10 +456,20 @@ class Fetcher(six.Iterator):
             batch = records.next_batch()
             while batch is not None:
 
-                # LegacyRecordBatch cannot access either base_offset or last_offset_delta
+                # Try DefaultsRecordBatch / message log format v2
+                # base_offset, last_offset_delta, and control batches
                 try:
                     self._subscriptions.assignment[tp].last_offset_from_message_batch = batch.base_offset + \
                                                                                         batch.last_offset_delta
+                    # Control batches have a single record indicating whether a transaction
+                    # was aborted or committed.
+                    # When isolation_level is READ_COMMITTED (currently unsupported)
+                    # we should also skip all messages from aborted transactions
+                    # For now we only support READ_UNCOMMITTED and so we ignore the
+                    # abort/commit signal.
+                    if batch.is_control_batch:
+                        batch = records.next_batch()
+                        continue
                 except AttributeError:
                     pass
 
@@ -674,7 +684,7 @@ class Fetcher(six.Iterator):
                 if next_offset_from_batch_header > self._subscriptions.assignment[partition].position:
                     log.debug(
                         "Advance position for partition %s from %s to %s (last message batch location plus one)"
-                        " to correct for deleted compacted messages",
+                        " to correct for deleted compacted messages and/or transactional control records",
                         partition, self._subscriptions.assignment[partition].position, next_offset_from_batch_header)
                     self._subscriptions.assignment[partition].position = next_offset_from_batch_header
 
diff --git a/kafka/coordinator/consumer.py b/kafka/coordinator/consumer.py
index 026fac8..5f62f73 100644
--- a/kafka/coordinator/consumer.py
+++ b/kafka/coordinator/consumer.py
@@ -128,7 +128,10 @@ class ConsumerCoordinator(BaseCoordinator):
 
     def __del__(self):
         if hasattr(self, '_cluster') and self._cluster:
-            self._cluster.remove_listener(WeakMethod(self._handle_metadata_update))
+            try:
+                self._cluster.remove_listener(WeakMethod(self._handle_metadata_update))
+            except TypeError:
+                pass
         super(ConsumerCoordinator, self).__del__()
 
     def protocol_type(self):
diff --git a/kafka/record/default_records.py b/kafka/record/default_records.py
index a098c42..b3a6fd0 100644
--- a/kafka/record/default_records.py
+++ b/kafka/record/default_records.py
@@ -269,8 +269,12 @@ class DefaultRecordBatch(DefaultRecordBase, ABCRecordBatch):
                 "payload, but instead read {}".format(length, pos - start_pos))
         self._pos = pos
 
-        return DefaultRecord(
-            offset, timestamp, self.timestamp_type, key, value, headers)
+        if self.is_control_batch:
+            return ControlRecord(
+                offset, timestamp, self.timestamp_type, key, value, headers)
+        else:
+            return DefaultRecord(
+                offset, timestamp, self.timestamp_type, key, value, headers)
 
     def __iter__(self):
         self._maybe_uncompress()
@@ -362,6 +366,45 @@ class DefaultRecord(ABCRecord):
         )
 
 
+class ControlRecord(DefaultRecord):
+    __slots__ = ("_offset", "_timestamp", "_timestamp_type", "_key", "_value",
+                 "_headers", "_version", "_type")
+
+    KEY_STRUCT = struct.Struct(
+        ">h"  # Current Version => Int16
+        "h"  # Type => Int16 (0 indicates an abort marker, 1 indicates a commit)
+    )
+
+    def __init__(self, offset, timestamp, timestamp_type, key, value, headers):
+        super(ControlRecord, self).__init__(offset, timestamp, timestamp_type, key, value, headers)
+        (self._version, self._type) = self.KEY_STRUCT.unpack(self._key)
+
+    # see https://kafka.apache.org/documentation/#controlbatch
+    @property
+    def version(self):
+        return self._version
+
+    @property
+    def type(self):
+        return self._type
+
+    @property
+    def abort(self):
+        return self._type == 0
+
+    @property
+    def commit(self):
+        return self._type == 1
+
+    def __repr__(self):
+        return (
+            "ControlRecord(offset={!r}, timestamp={!r}, timestamp_type={!r},"
+            " version={!r}, type={!r} <{!s}>)".format(
+                self._offset, self._timestamp, self._timestamp_type,
+                self._version, self._type, "abort" if self.abort else "commit")
+        )
+
+
 class DefaultRecordBatchBuilder(DefaultRecordBase, ABCRecordBatchBuilder):
 
     # excluding key, value and headers:


Test patch:
diff --git a/test/record/test_records.py b/test/record/test_records.py
index 5ed22d8..cab9592 100644
--- a/test/record/test_records.py
+++ b/test/record/test_records.py
@@ -60,6 +60,15 @@ record_batch_data_v0 = [
     b'\x00\xff\xff\xff\xff\x00\x00\x00\x03123'
 ]
 
+# Single record control batch (abort)
+control_batch_data_v2 = [
+    b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00R\x00\x00\x00\x00'
+    b'\x02e\x97\xff\xd0\x00\x20\x00\x00\x00\x00\x00\x00\x00\x00\x00'
+    b'\x98\x96\x7f\x00\x00\x00\x00\x00\x98\x96'
+    b'\x7f\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff'
+    b'\x00\x00\x00\x01@\x00\x00\x00\x08\x00\x00\x00\x00,opaque-control-message\x00'
+]
+
 
 def test_memory_records_v2():
     data_bytes = b"".join(record_batch_data_v2) + b"\x00" * 4
@@ -230,3 +239,18 @@ def test_memory_records_builder_full(magic, compression_type):
         key=None, timestamp=None, value=b"M")
     assert metadata is None
     assert builder.next_offset() == 1
+
+
+def test_control_record_v2():
+    data_bytes = b"".join(control_batch_data_v2)
+    records = MemoryRecords(data_bytes)
+
+    assert records.has_next() is True
+    batch = records.next_batch()
+    assert batch.is_control_batch is True
+    recs = list(batch)
+    assert len(recs) == 1
+    assert recs[0].version == 0
+    assert recs[0].type == 0
+    assert recs[0].abort is True
+    assert recs[0].commit is False
diff --git a/test/test_client_async.py b/test/test_client_async.py
index b9b4150..ccdd570 100644
--- a/test/test_client_async.py
+++ b/test/test_client_async.py
@@ -43,7 +43,7 @@ def test_bootstrap(mocker, conn):
     kwargs.pop('state_change_callback')
     kwargs.pop('node_id')
     assert kwargs == cli.config
-    conn.send.assert_called_once_with(MetadataRequest[0]([]), blocking=False)
+    conn.send.assert_called_once_with(MetadataRequest[0]([]), blocking=False, request_timeout_ms=None)
     assert cli._bootstrap_fails == 0
     assert cli.cluster.brokers() == set([BrokerMetadata(0, 'foo', 12, None),
                                          BrokerMetadata(1, 'bar', 34, None)])
@@ -220,12 +220,12 @@ def test_send(cli, conn):
     request = ProduceRequest[0](0, 0, [])
     assert request.expect_response() is False
     ret = cli.send(0, request)
-    conn.send.assert_called_with(request, blocking=False)
+    conn.send.assert_called_with(request, blocking=False, request_timeout_ms=None)
     assert isinstance(ret, Future)
 
     request = MetadataRequest[0]([])
     cli.send(0, request)
-    conn.send.assert_called_with(request, blocking=False)
+    conn.send.assert_called_with(request, blocking=False, request_timeout_ms=None)
 
 
 def test_poll(mocker):
diff --git a/test/test_conn.py b/test/test_conn.py
index fb41728..f41153f 100644
--- a/test/test_conn.py
+++ b/test/test_conn.py
@@ -347,14 +347,14 @@ def test_requests_timed_out(conn):
         # No in-flight requests, not timed out
         assert not conn.requests_timed_out()
 
-        # Single request, timestamp = now (0)
-        conn.in_flight_requests[0] = ('foo', 0)
+        # Single request, timeout_at > now (0)
+        conn.in_flight_requests[0] = ('foo', 0, 1)
         assert not conn.requests_timed_out()
 
         # Add another request w/ timestamp > request_timeout ago
         request_timeout = conn.config['request_timeout_ms']
         expired_timestamp = 0 - request_timeout - 1
-        conn.in_flight_requests[1] = ('bar', expired_timestamp)
+        conn.in_flight_requests[1] = ('bar', 0, expired_timestamp)
         assert conn.requests_timed_out()
 
         # Drop the expired request and we should be good to go again