Repository: dpkp/kafka-python
Pull request #2504 (2025-02-27 17:41:40)

Improve error handling in client._maybe_connect
Avoid assertion error if broker metadata is not found


Patch:
diff --git a/kafka/client_async.py b/kafka/client_async.py
index 96959d9..9e57efd 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -399,13 +399,23 @@ class KafkaClient(object):
         return False
 
     def _maybe_connect(self, node_id):
-        """Idempotent non-blocking connection attempt to the given node id."""
+        """Idempotent non-blocking connection attempt to the given node id.
+
+        Returns True if connection object exists and is connected / connecting
+        """
         with self._lock:
             conn = self._conns.get(node_id)
 
+            # Check if existing connection should be recreated because host/port changed
+            if conn is not None and self._should_recycle_connection(conn):
+                self._conns.pop(node_id).close()
+                conn = None
+
             if conn is None:
                 broker = self.cluster.broker_metadata(node_id)
-                assert broker, 'Broker id %s not in current metadata' % (node_id,)
+                if broker is None:
+                    log.debug('Broker id %s not in current metadata', node_id)
+                    return False
 
                 log.debug("Initiating connection to node %s at %s:%s",
                           node_id, broker.host, broker.port)
@@ -417,16 +427,11 @@ class KafkaClient(object):
                                         **self.config)
                 self._conns[node_id] = conn
 
-            # Check if existing connection should be recreated because host/port changed
-            elif self._should_recycle_connection(conn):
-                self._conns.pop(node_id)
-                return False
-
             elif conn.connected():
                 return True
 
             conn.connect()
-            return conn.connected()
+            return not conn.disconnected()
 
     def ready(self, node_id, metadata_priority=True):
         """Check whether a node is connected and ok to send more requests.
@@ -621,7 +626,10 @@ class KafkaClient(object):
 
                 # Attempt to complete pending connections
                 for node_id in list(self._connecting):
-                    self._maybe_connect(node_id)
+                    # False return means no more connection progress is possible
+                    # Connected nodes will update _connecting via state_change callback
+                    if not self._maybe_connect(node_id):
+                        self._connecting.remove(node_id)
 
                 # If we got a future that is already done, don't block in _poll
                 if future is not None and future.is_done:
@@ -965,7 +973,12 @@ class KafkaClient(object):
             if try_node is None:
                 self._lock.release()
                 raise Errors.NoBrokersAvailable()
-            self._maybe_connect(try_node)
+            if not self._maybe_connect(try_node):
+                if try_node == node_id:
+                    raise Errors.NodeNotReadyError("Connection failed to %s" % node_id)
+                else:
+                    continue
+
             conn = self._conns[try_node]
 
             # We will intentionally cause socket failures
diff --git a/kafka/coordinator/base.py b/kafka/coordinator/base.py
index 8f588aa..8607e48 100644
--- a/kafka/coordinator/base.py
+++ b/kafka/coordinator/base.py
@@ -453,7 +453,7 @@ class BaseCoordinator(object):
             (protocol, metadata if isinstance(metadata, bytes) else metadata.encode())
             for protocol, metadata in self.group_protocols()
         ]
-        version = self._client.api_version(JoinGroupRequest, max_version=2)
+        version = self._client.api_version(JoinGroupRequest, max_version=3)
         if version == 0:
             request = JoinGroupRequest[version](
                 self.group_id,
@@ -493,6 +493,11 @@ class BaseCoordinator(object):
         future.failure(error)
 
     def _handle_join_group_response(self, future, send_time, response):
+        if response.API_VERSION >= 2:
+            self.sensors.throttle_time.record(response.throttle_time_ms)
+            if response.throttle_time_ms > 0:
+                log.warning("JoinGroupRequest throttled by broker (%d ms)", response.throttle_time_ms)
+
         error_type = Errors.for_code(response.error_code)
         if error_type is Errors.NoError:
             log.debug("Received successful JoinGroup response for group %s: %s",
@@ -554,7 +559,7 @@ class BaseCoordinator(object):
 
     def _on_join_follower(self):
         # send follower's sync group with an empty assignment
-        version = self._client.api_version(SyncGroupRequest, max_version=1)
+        version = self._client.api_version(SyncGroupRequest, max_version=2)
         request = SyncGroupRequest[version](
             self.group_id,
             self._generation.generation_id,
@@ -582,7 +587,7 @@ class BaseCoordinator(object):
         except Exception as e:
             return Future().failure(e)
 
-        version = self._client.api_version(SyncGroupRequest, max_version=1)
+        version = self._client.api_version(SyncGroupRequest, max_version=2)
         request = SyncGroupRequest[version](
             self.group_id,
             self._generation.generation_id,
@@ -614,6 +619,11 @@ class BaseCoordinator(object):
         return future
 
     def _handle_sync_group_response(self, future, send_time, response):
+        if response.API_VERSION >= 1:
+            self.sensors.throttle_time.record(response.throttle_time_ms)
+            if response.throttle_time_ms > 0:
+                log.warning("SyncGroupRequest throttled by broker (%d ms)", response.throttle_time_ms)
+
         error_type = Errors.for_code(response.error_code)
         if error_type is Errors.NoError:
             self.sensors.sync_latency.record((time.time() - send_time) * 1000)
@@ -770,7 +780,7 @@ class BaseCoordinator(object):
                 # this is a minimal effort attempt to leave the group. we do not
                 # attempt any resending if the request fails or times out.
                 log.info('Leaving consumer group (%s).', self.group_id)
-                version = self._client.api_version(LeaveGroupRequest, max_version=1)
+                version = self._client.api_version(LeaveGroupRequest, max_version=2)
                 request = LeaveGroupRequest[version](self.group_id, self._generation.member_id)
                 future = self._client.send(self.coordinator_id, request)
                 future.add_callback(self._handle_leave_group_response)
@@ -780,6 +790,11 @@ class BaseCoordinator(object):
             self.reset_generation()
 
     def _handle_leave_group_response(self, response):
+        if response.API_VERSION >= 1:
+            self.sensors.throttle_time.record(response.throttle_time_ms)
+            if response.throttle_time_ms > 0:
+                log.warning("LeaveGroupRequest throttled by broker (%d ms)", response.throttle_time_ms)
+
         error_type = Errors.for_code(response.error_code)
         if error_type is Errors.NoError:
             log.debug("LeaveGroup request for group %s returned successfully",
@@ -798,7 +813,7 @@ class BaseCoordinator(object):
             e = Errors.NodeNotReadyError(self.coordinator_id)
             return Future().failure(e)
 
-        version = self._client.api_version(HeartbeatRequest, max_version=1)
+        version = self._client.api_version(HeartbeatRequest, max_version=2)
         request = HeartbeatRequest[version](self.group_id,
                                             self._generation.generation_id,
                                             self._generation.member_id)
@@ -811,6 +826,11 @@ class BaseCoordinator(object):
         return future
 
     def _handle_heartbeat_response(self, future, send_time, response):
+        if response.API_VERSION >= 1:
+            self.sensors.throttle_time.record(response.throttle_time_ms)
+            if response.throttle_time_ms > 0:
+                log.warning("HeartbeatRequest throttled by broker (%d ms)", response.throttle_time_ms)
+
         self.sensors.heartbeat_latency.record((time.time() - send_time) * 1000)
         error_type = Errors.for_code(response.error_code)
         if error_type is Errors.NoError:
@@ -899,6 +919,14 @@ class GroupCoordinatorMetrics(object):
             tags), AnonMeasurable(
                 lambda _, now: (now / 1000) - self.heartbeat.last_send))
 
+        self.throttle_time = metrics.sensor('throttle-time')
+        self.throttle_time.add(metrics.metric_name(
+            'throttle-time-avg', self.metric_group_name,
+            'The average throttle time in ms'), Avg())
+        self.throttle_time.add(metrics.metric_name(
+            'throttle-time-max', self.metric_group_name,
+            'The maximum throttle time in ms'), Max())
+
 
 class HeartbeatThread(threading.Thread):
     def __init__(self, coordinator):
diff --git a/kafka/protocol/group.py b/kafka/protocol/group.py
index bcb9655..ee06141 100644
--- a/kafka/protocol/group.py
+++ b/kafka/protocol/group.py
@@ -42,6 +42,12 @@ class JoinGroupResponse_v2(Response):
     )
 
 
+class JoinGroupResponse_v3(Response):
+    API_KEY = 11
+    API_VERSION = 3
+    SCHEMA = JoinGroupResponse_v2.SCHEMA
+
+
 class JoinGroupRequest_v0(Request):
     API_KEY = 11
     API_VERSION = 0
@@ -83,11 +89,19 @@ class JoinGroupRequest_v2(Request):
     UNKNOWN_MEMBER_ID = ''
 
 
+class JoinGroupRequest_v3(Request):
+    API_KEY = 11
+    API_VERSION = 3
+    RESPONSE_TYPE = JoinGroupResponse_v3
+    SCHEMA = JoinGroupRequest_v2.SCHEMA
+    UNKNOWN_MEMBER_ID = ''
+
+
 JoinGroupRequest = [
-    JoinGroupRequest_v0, JoinGroupRequest_v1, JoinGroupRequest_v2
+    JoinGroupRequest_v0, JoinGroupRequest_v1, JoinGroupRequest_v2, JoinGroupRequest_v3
 ]
 JoinGroupResponse = [
-    JoinGroupResponse_v0, JoinGroupResponse_v1, JoinGroupResponse_v2
+    JoinGroupResponse_v0, JoinGroupResponse_v1, JoinGroupResponse_v2, JoinGroupResponse_v3
 ]
 
 
@@ -118,6 +132,12 @@ class SyncGroupResponse_v1(Response):
     )
 
 
+class SyncGroupResponse_v2(Response):
+    API_KEY = 14
+    API_VERSION = 2
+    SCHEMA = SyncGroupResponse_v1.SCHEMA
+
+
 class SyncGroupRequest_v0(Request):
     API_KEY = 14
     API_VERSION = 0
@@ -139,8 +159,15 @@ class SyncGroupRequest_v1(Request):
     SCHEMA = SyncGroupRequest_v0.SCHEMA
 
 
-SyncGroupRequest = [SyncGroupRequest_v0, SyncGroupRequest_v1]
-SyncGroupResponse = [SyncGroupResponse_v0, SyncGroupResponse_v1]
+class SyncGroupRequest_v2(Request):
+    API_KEY = 14
+    API_VERSION = 2
+    RESPONSE_TYPE = SyncGroupResponse_v2
+    SCHEMA = SyncGroupRequest_v1.SCHEMA
+
+
+SyncGroupRequest = [SyncGroupRequest_v0, SyncGroupRequest_v1, SyncGroupRequest_v2]
+SyncGroupResponse = [SyncGroupResponse_v0, SyncGroupResponse_v1, SyncGroupResponse_v2]
 
 
 class MemberAssignment(Struct):
@@ -170,6 +197,12 @@ class HeartbeatResponse_v1(Response):
     )
 
 
+class HeartbeatResponse_v2(Response):
+    API_KEY = 12
+    API_VERSION = 2
+    SCHEMA = HeartbeatResponse_v1.SCHEMA
+
+
 class HeartbeatRequest_v0(Request):
     API_KEY = 12
     API_VERSION = 0
@@ -188,8 +221,15 @@ class HeartbeatRequest_v1(Request):
     SCHEMA = HeartbeatRequest_v0.SCHEMA
 
 
-HeartbeatRequest = [HeartbeatRequest_v0, HeartbeatRequest_v1]
-HeartbeatResponse = [HeartbeatResponse_v0, HeartbeatResponse_v1]
+class HeartbeatRequest_v2(Request):
+    API_KEY = 12
+    API_VERSION = 2
+    RESPONSE_TYPE = HeartbeatResponse_v2
+    SCHEMA = HeartbeatRequest_v1.SCHEMA
+
+
+HeartbeatRequest = [HeartbeatRequest_v0, HeartbeatRequest_v1, HeartbeatRequest_v2]
+HeartbeatResponse = [HeartbeatResponse_v0, HeartbeatResponse_v1, HeartbeatResponse_v2]
 
 
 class LeaveGroupResponse_v0(Response):
@@ -209,6 +249,12 @@ class LeaveGroupResponse_v1(Response):
     )
 
 
+class LeaveGroupResponse_v2(Response):
+    API_KEY = 13
+    API_VERSION = 2
+    SCHEMA = LeaveGroupResponse_v1.SCHEMA
+
+
 class LeaveGroupRequest_v0(Request):
     API_KEY = 13
     API_VERSION = 0
@@ -226,5 +272,12 @@ class LeaveGroupRequest_v1(Request):
     SCHEMA = LeaveGroupRequest_v0.SCHEMA
 
 
-LeaveGroupRequest = [LeaveGroupRequest_v0, LeaveGroupRequest_v1]
-LeaveGroupResponse = [LeaveGroupResponse_v0, LeaveGroupResponse_v1]
+class LeaveGroupRequest_v2(Request):
+    API_KEY = 13
+    API_VERSION = 2
+    RESPONSE_TYPE = LeaveGroupResponse_v2
+    SCHEMA = LeaveGroupRequest_v1.SCHEMA
+
+
+LeaveGroupRequest = [LeaveGroupRequest_v0, LeaveGroupRequest_v1, LeaveGroupRequest_v2]
+LeaveGroupResponse = [LeaveGroupResponse_v0, LeaveGroupResponse_v1, LeaveGroupResponse_v2]


Test patch:
diff --git a/test/test_client_async.py b/test/test_client_async.py
index ccdd570..16ee429 100644
--- a/test/test_client_async.py
+++ b/test/test_client_async.py
@@ -71,19 +71,14 @@ def test_can_connect(cli, conn):
 
 
 def test_maybe_connect(cli, conn):
-    try:
-        # Node not in metadata, raises AssertionError
-        cli._maybe_connect(2)
-    except AssertionError:
-        pass
-    else:
-        assert False, 'Exception not raised'
+    # Node not in metadata, return False
+    assert not cli._maybe_connect(2)
 
     # New node_id creates a conn object
     assert 0 not in cli._conns
     conn.state = ConnectionStates.DISCONNECTED
     conn.connect.side_effect = lambda: conn._set_conn_state(ConnectionStates.CONNECTING)
-    assert cli._maybe_connect(0) is False
+    assert cli._maybe_connect(0) is True
     assert cli._conns[0] is conn