Improve kafka consumer highwater offset collection time (#20716)

* Improve kafka consumer highwater offset collection time

* update seed mock client

* fix the failing test_oauth_config

* fetch highwater offset if there are actual partitions to check

* linter errors

Author: arshdeeptinna

kafka_consumer/changelog.d/20716.fixed

@@ -0,0 +1 @@
+Improve kafka consumer highwater offset collection time

kafka_consumer/datadog_checks/kafka_consumer/client.py

@@ -13,6 +13,7 @@ def __init__(self, config, log) -> None:
         self.log = log
         self._kafka_client = None
         self._consumer = None
+        self._cluster_metadata = None
 
     @property
     def kafka_client(self):
@@ -106,18 +107,48 @@ def consumer_offsets_for_times(self, partitions):
             )
         ]
 
+    def _list_topics(self):
+        if self._cluster_metadata:
+            return self._cluster_metadata
+
+        try:
+            self.request_metadata_update()
+
+        except KafkaException as e:
+            self.log.error("Received exception when listing topics: %s", e)
+
+        return self._cluster_metadata
+
+    def get_topic_partitions(self):
+        topic_partitions = {}
+        try:
+            cluster_metadata = self._list_topics()
+            for topic in cluster_metadata.topics:
+                topic_metadata = cluster_metadata.topics[topic]
+                partitions = list(topic_metadata.partitions)
+                topic_partitions[topic] = partitions
+
+        except KafkaException as e:
+            self.log.error("Received exception when listing topics: %s", e)
+
+        return topic_partitions
+
     def get_partitions_for_topic(self, topic):
+        partitions = []
         try:
-            cluster_metadata = self.kafka_client.list_topics(topic, timeout=self.config._request_timeout)
+            cluster_metadata = self._list_topics()
         except KafkaException as e:
             self.log.error("Received exception when getting partitions for topic %s: %s", topic, e)
             return []
-        topic_metadata = cluster_metadata.topics[topic]
-        return list(topic_metadata.partitions)
+
+        if topic in cluster_metadata.topics:
+            topic_metadata = cluster_metadata.topics[topic]
+            partitions = list(topic_metadata.partitions)
+        return partitions
 
     def request_metadata_update(self):
         # https://github.com/confluentinc/confluent-kafka-python/issues/594
-        self.kafka_client.list_topics(None, timeout=self.config._request_timeout)
+        self._cluster_metadata = self.kafka_client.list_topics(None, timeout=self.config._request_timeout)
 
     def list_consumer_groups(self):
         groups = []

kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py

@@ -343,60 +343,43 @@ def get_highwater_offsets(self, consumer_offsets):
         self.log.debug('Getting highwater offsets')
 
         cluster_id = ""
+        dd_consumer_group = "datadog-agent"
         highwater_offsets = {}
-        topics_with_consumer_offset = set()
-        topic_partition_with_consumer_offset = set()
+        topic_partitions_to_check = set()
 
-        for consumer_group, offsets in consumer_offsets.items():
-            self.log.debug('CONSUMER GROUP: %s', consumer_group)
-            topic_partitions_for_highwater_offsets = set()
-            self.client.open_consumer(consumer_group)
-            cluster_id, topics = self.client.consumer_get_cluster_id_and_list_topics(consumer_group)
-
-            if not self.config._monitor_all_broker_highwatermarks:
-                for topic, partition in offsets:
-                    topics_with_consumer_offset.add(topic)
-                    topic_partition_with_consumer_offset.add((topic, partition))
-
-            for topic, partitions in topics:
+        if self.config._monitor_all_broker_highwatermarks:
+            all_topic_partitions = self.client.get_topic_partitions()
+            for topic in all_topic_partitions:
                 if topic in KAFKA_INTERNAL_TOPICS:
                     self.log.debug("Skipping internal topic %s", topic)
                     continue
-                if not self.config._monitor_all_broker_highwatermarks and topic not in topics_with_consumer_offset:
-                    self.log.debug("Skipping non-relevant topic %s", topic)
-                    continue
 
-                for partition in partitions:
-                    if (topic, partition) in highwater_offsets:
-                        self.log.debug(
-                            'Highwater offset already collected for topic %s with partition %s', topic, partition
-                        )
+                for partition in all_topic_partitions[topic]:
+                    topic_partitions_to_check.add((topic, partition))
+
+        else:
+            for _, offsets in consumer_offsets.items():
+                for topic, partition in offsets:
+                    if topic in KAFKA_INTERNAL_TOPICS:
+                        self.log.debug("Skipping internal topic %s", topic)
                         continue
-                    if (
-                        self.config._monitor_all_broker_highwatermarks
-                        or (topic, partition) in topic_partition_with_consumer_offset
-                    ):
-                        topic_partitions_for_highwater_offsets.add((topic, partition))
-                        self.log.debug('TOPIC: %s', topic)
-                        self.log.debug('PARTITION: %s', partition)
-                    else:
-                        self.log.debug("Skipping non-relevant partition %s of topic %s", partition, topic)
 
-            if topic_partitions_for_highwater_offsets:
-                self.log.debug(
-                    'Querying %s highwater offsets for consumer group %s',
-                    len(topic_partitions_for_highwater_offsets),
-                    consumer_group,
-                )
-                for topic, partition, offset in self.client.consumer_offsets_for_times(
-                    partitions=topic_partitions_for_highwater_offsets
-                ):
-                    highwater_offsets[(topic, partition)] = offset
-            else:
-                self.log.debug('No new highwater offsets to query for consumer group %s', consumer_group)
+                    topic_partitions_to_check.add((topic, partition))
 
-            self.client.close_consumer()
+        self.client.open_consumer(dd_consumer_group)
+        cluster_id, _ = self.client.consumer_get_cluster_id_and_list_topics(dd_consumer_group)
+        self.log.debug(
+            'Querying %s highwater offsets for consumer group %s',
+            len(topic_partitions_to_check),
+            dd_consumer_group,
+        )
+        if topic_partitions_to_check:
+            for topic, partition, offset in self.client.consumer_offsets_for_times(
+                partitions=topic_partitions_to_check
+            ):
+                highwater_offsets[(topic, partition)] = offset
 
+        self.client.close_consumer()
         self.log.debug('Got %s highwater offsets', len(highwater_offsets))
         return highwater_offsets, cluster_id
 

kafka_consumer/tests/test_unit.py

@@ -36,7 +36,7 @@ def fake_consumer_offsets_for_times(partitions):
 def seed_mock_client(cluster_id="cluster_id"):
     """Set some common defaults for the mock client to kafka."""
     client = mock.create_autospec(KafkaClient)
-    client.list_consumer_groups.return_value = ["consumer_group1"]
+    client.list_consumer_groups.return_value = ["consumer_group1", "datadog-agent"]
     client.get_partitions_for_topic.return_value = ['partition1']
     client.list_consumer_group_offsets.return_value = [("consumer_group1", [("topic1", "partition1", 2)])]
     client.describe_consumer_group.return_value = 'STABLE'
@@ -133,6 +133,20 @@ def test_tls_verify_is_string(tls_verify, expected, check, kafka_instance):
 
 mock_client = mock.MagicMock()
 mock_client.get_highwater_offsets.return_value = ({}, "")
+mock_client.consumer_get_cluster_id_and_list_topics.return_value = (
+    "cluster_id",
+    # topics
+    [
+        # Used in unit tets
+        ('topic1', ["partition1"]),
+        ('topic2', ["partition2"]),
+        # Copied from integration tests
+        ('dc', [0, 1]),
+        ('unconsumed_topic', [0, 1]),
+        ('marvel', [0, 1]),
+        ('__consumer_offsets', [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
+    ],
+)
 
 
 @pytest.mark.parametrize(
@@ -468,7 +482,7 @@ def test_load_broker_timestamps_empty(
 
 def test_client_init(kafka_instance, check, dd_run_check):
     """
-    We only open a connection to a consumer once per consumer group.
+    We only open a connection to datadog-agent consumer once.
 
     Doing so more often degrades performance, as described in this issue:
     https://github.com/DataDog/integrations-core/issues/19564
@@ -478,7 +492,7 @@ def test_client_init(kafka_instance, check, dd_run_check):
     check.client = mock_client
     dd_run_check(check)
 
-    assert check.client.open_consumer.mock_calls == [mock.call("consumer_group1")]
+    assert check.client.open_consumer.mock_calls == [mock.call("datadog-agent")]
 
 
 def test_resolve_start_offsets():