Repository: dpkp/kafka-python
Pull request #2652 (2025-06-23 22:55:21)

Fix coordinator lock contention during close()
Fixes heartbeat thread clean shutdown during `close()`. Were holding `coordinator._lock` in outer function so that inner `notify()` did not actually unblock heartbeat thread.


Patch:
diff --git a/kafka/consumer/group.py b/kafka/consumer/group.py
index 4eb9e2a..5ac14eb 100644
--- a/kafka/consumer/group.py
+++ b/kafka/consumer/group.py
@@ -5,7 +5,7 @@ import logging
 import socket
 import time
 
-from kafka.errors import KafkaConfigurationError, KafkaTimeoutError, UnsupportedVersionError
+from kafka.errors import KafkaConfigurationError, UnsupportedVersionError
 
 from kafka.vendor import six
 
diff --git a/kafka/coordinator/base.py b/kafka/coordinator/base.py
index 0eb7f0e..5e1f726 100644
--- a/kafka/coordinator/base.py
+++ b/kafka/coordinator/base.py
@@ -857,14 +857,12 @@ class BaseCoordinator(object):
                 self._heartbeat_thread.disable()
 
     def _close_heartbeat_thread(self, timeout_ms=None):
-        with self._lock:
-            if self._heartbeat_thread is not None:
-                heartbeat_log.info('Stopping heartbeat thread')
-                try:
-                    self._heartbeat_thread.close(timeout_ms=timeout_ms)
-                except ReferenceError:
-                    pass
-                self._heartbeat_thread = None
+        if self._heartbeat_thread is not None:
+            try:
+                self._heartbeat_thread.close(timeout_ms=timeout_ms)
+            except ReferenceError:
+                pass
+            self._heartbeat_thread = None
 
     def __del__(self):
         try:
@@ -1047,17 +1045,20 @@ class HeartbeatThread(threading.Thread):
             self.enabled = False
 
     def close(self, timeout_ms=None):
-        if self.closed:
-            return
-        self.closed = True
+        with self.coordinator._lock:
+            if self.closed:
+                return
 
-        # Generally this should not happen - close() is triggered
-        # by the coordinator. But in some cases GC may close the coordinator
-        # from within the heartbeat thread.
-        if threading.current_thread() == self:
-            return
+            heartbeat_log.info('Stopping heartbeat thread')
+            self.closed = True
 
-        with self.coordinator._lock:
+            # Generally this should not happen - close() is triggered
+            # by the coordinator. But in some cases GC may close the coordinator
+            # from within the heartbeat thread.
+            if threading.current_thread() == self:
+                return
+
+            # Notify coordinator lock to wake thread from sleep/lock.wait
             self.coordinator._lock.notify()
 
         if self.is_alive():
diff --git a/kafka/protocol/parser.py b/kafka/protocol/parser.py
index 4bc4273..8683a95 100644
--- a/kafka/protocol/parser.py
+++ b/kafka/protocol/parser.py
@@ -6,7 +6,7 @@ import logging
 import kafka.errors as Errors
 from kafka.protocol.find_coordinator import FindCoordinatorResponse
 from kafka.protocol.frame import KafkaBytes
-from kafka.protocol.types import Int32, TaggedFields
+from kafka.protocol.types import Int32
 from kafka.version import __version__
 
 log = logging.getLogger(__name__)


Test patch:
diff --git a/test/integration/test_admin_integration.py b/test/integration/test_admin_integration.py
index 1c38400..5292080 100644
--- a/test/integration/test_admin_integration.py
+++ b/test/integration/test_admin_integration.py
@@ -9,7 +9,7 @@ from time import time, sleep
 from kafka.admin import (
     ACLFilter, ACLOperation, ACLPermissionType, ResourcePattern, ResourceType, ACL, ConfigResource, ConfigResourceType)
 from kafka.errors import (
-        BrokerResponseError, KafkaError, NoError, CoordinatorNotAvailableError, NonEmptyGroupError,
+        BrokerResponseError, NoError, CoordinatorNotAvailableError, NonEmptyGroupError,
         GroupIdNotFoundError, OffsetOutOfRangeError, UnknownTopicOrPartitionError)
 
 
diff --git a/test/test_assignors.py b/test/test_assignors.py
index 858ef42..91d7157 100644
--- a/test/test_assignors.py
+++ b/test/test_assignors.py
@@ -9,8 +9,8 @@ import pytest
 from kafka.structs import TopicPartition
 from kafka.coordinator.assignors.range import RangePartitionAssignor
 from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
-from kafka.coordinator.assignors.sticky.sticky_assignor import StickyPartitionAssignor, StickyAssignorUserDataV1
-from kafka.coordinator.protocol import ConsumerProtocolMemberAssignment, ConsumerProtocolMemberMetadata
+from kafka.coordinator.assignors.sticky.sticky_assignor import StickyPartitionAssignor
+from kafka.coordinator.protocol import ConsumerProtocolMemberAssignment
 from kafka.vendor import six