kafka_consumer: cache earliest offsets across collection intervals (DataDog#24515)

* kafka_consumer: cache earliest offsets across collection intervals

Log start offsets only move via the broker's log-cleaner cycle
(log.retention.check.interval.ms), so refetching them on every check
run is unnecessary broker load. Cache the result with a TTL derived
from that broker config, clamped to a sane range.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* Add changelog entry

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* kafka_consumer: trim redundant comments, add earliest-offsets cache test

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* Address review feedback: avoid an undocumented metric and refetch missing partitions on a fresh cache

log.retention.check.interval.ms is now read from broker config data purely to
derive the earliest-offsets cache TTL, without being added to the metric
emission list. fetch_earliest_offsets also now refetches from the broker when
a fresh cache doesn't cover every requested partition (e.g. newly added
partitions), while keeping the cache's original expiration.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

Author: piochelepiotr

kafka_consumer/changelog.d/24515.fixed

@@ -0,0 +1 @@
+Cache the earliest (log-start) offset used for `kafka.topic.size` and `kafka.partition.size` across collection intervals instead of refetching it from the broker on every run.

kafka_consumer/datadog_checks/kafka_consumer/cluster_metadata.py

@@ -68,8 +68,14 @@ def __init__(self, check, client, config, log):
         self.SCHEMA_COMPATIBILITY_FETCH_CACHE_MAX_SIZE = 20_000
         self.SCHEMA_ID_CACHE_MAX_SIZE = 20_000
 
+        self.EARLIEST_OFFSETS_DEFAULT_TTL = 300  # 5 minutes, matches Kafka's own broker default
+        self.EARLIEST_OFFSETS_MIN_TTL = 60
+        self.EARLIEST_OFFSETS_MAX_TTL = 1800
+        self._log_retention_check_interval_s: float | None = None
+
         self.BROKER_CONFIG_CACHE_KEY = 'kafka_broker_config_cache'
         self.BROKER_CONFIG_FETCH_CACHE_KEY = 'kafka_broker_config_fetch_cache'
+        self.EARLIEST_OFFSETS_CACHE_KEY = 'kafka_earliest_offsets_cache'
         self.TOPIC_CONFIG_CACHE_KEY = 'kafka_topic_config_cache'
         self.TOPIC_CONFIG_FETCH_CACHE_KEY = 'kafka_topic_config_fetch_cache'
         self.TOPIC_HWM_SUM_CACHE_KEY = 'kafka_topic_hwm_sum_cache'
@@ -348,6 +354,21 @@ def _collect_broker_metadata(self, metadata=None):
                                 config_data[config_name],
                             )
 
+                retention_check_interval_ms = config_data.get('log.retention.check.interval.ms')
+                if retention_check_interval_ms:
+                    try:
+                        interval_s = float(retention_check_interval_ms) / 1000
+                        if interval_s > 0:
+                            self._log_retention_check_interval_s = min(
+                                interval_s, self._log_retention_check_interval_s or interval_s
+                            )
+                    except (ValueError, TypeError):
+                        self.log.debug(
+                            "Could not convert broker %s config log.retention.check.interval.ms value %r to float",
+                            broker_id_str,
+                            retention_check_interval_ms,
+                        )
+
                 truncated_config = self._truncate_config_for_event(config_data, max_configs=50)
                 event_text = json.dumps(truncated_config, indent=2, sort_keys=True)
 
@@ -386,7 +407,66 @@ def _collect_broker_metadata(self, metadata=None):
                 "data-streams-message",
             )
 
+    def _topic_partition_pairs(self, topic_partitions):
+        return {
+            (topic, partition)
+            for topic, partitions in topic_partitions.items()
+            if topic not in KAFKA_INTERNAL_TOPICS
+            for partition in partitions
+        }
+
+    def _earliest_offsets_ttl(self) -> float:
+        """TTL for the earliest-offsets cache, derived from the broker's log-cleaner cycle."""
+        ttl = self._log_retention_check_interval_s or self.EARLIEST_OFFSETS_DEFAULT_TTL
+        return max(self.EARLIEST_OFFSETS_MIN_TTL, min(self.EARLIEST_OFFSETS_MAX_TTL, ttl))
+
+    def _load_earliest_offsets_cache(self) -> dict[str, Any] | None:
+        try:
+            cached_str = self.check.read_persistent_cache(self.EARLIEST_OFFSETS_CACHE_KEY)
+            if not cached_str:
+                return None
+            data = json.loads(cached_str)
+            return {
+                'expire_at': data['expire_at'],
+                'offsets': {(topic, partition): offset for topic, partition, offset in data['offsets']},
+            }
+        except Exception as e:
+            self.log.debug("Could not read earliest offsets cache: %s", e)
+            return None
+
+    def _save_earliest_offsets_cache(self, offsets: dict[tuple[str, int], int], expire_at: float | None = None) -> None:
+        try:
+            payload = {
+                'expire_at': expire_at if expire_at is not None else time.time() + self._earliest_offsets_ttl(),
+                'offsets': [[topic, partition, offset] for (topic, partition), offset in offsets.items()],
+            }
+            self.check.write_persistent_cache(self.EARLIEST_OFFSETS_CACHE_KEY, json.dumps(payload))
+        except Exception as e:
+            self.log.debug("Could not write earliest offsets cache: %s", e)
+
     def fetch_earliest_offsets(self, topic_partitions):
+        """Return cached log-start offsets, refetching from the broker only once the TTL expires."""
+        requested = self._topic_partition_pairs(topic_partitions)
+        if not requested:
+            return {}
+
+        cached = self._load_earliest_offsets_cache()
+        if cached is not None and time.time() < cached['expire_at']:
+            cached_offsets = {tp: offset for tp, offset in cached['offsets'].items() if tp in requested}
+            if cached_offsets.keys() == requested:
+                return cached_offsets
+
+            result = self._fetch_earliest_offsets_from_broker(topic_partitions)
+            if result:
+                self._save_earliest_offsets_cache(result, expire_at=cached['expire_at'])
+            return result
+
+        result = self._fetch_earliest_offsets_from_broker(topic_partitions)
+        if result:
+            self._save_earliest_offsets_cache(result)
+        return result
+
+    def _fetch_earliest_offsets_from_broker(self, topic_partitions):
         """Batch-fetch log-start offsets via AdminClient.list_offsets(earliest).
 
         Uses ListOffsets with the EARLIEST_TIMESTAMP sentinel, which the broker
@@ -397,9 +477,7 @@ def fetch_earliest_offsets(self, topic_partitions):
         """
         requests = {
             TopicPartition(topic, partition): OffsetSpec.earliest()
-            for topic, partitions in topic_partitions.items()
-            if topic not in KAFKA_INTERNAL_TOPICS
-            for partition in partitions
+            for topic, partition in self._topic_partition_pairs(topic_partitions)
         }
         if not requests:
             return {}

kafka_consumer/tests/test_cluster_metadata.py

@@ -1135,6 +1135,57 @@ def test_kafka_configs_refresh_interval(check, interval, expected_interval, expe
     assert collector.cache.refresh_jitter == expected_jitter
 
 
+def test_fetch_earliest_offsets_cached_across_calls(check):
+    """fetch_earliest_offsets should hit the broker once, then serve later calls from cache."""
+    instance = {
+        'kafka_connect_str': 'localhost:9092',
+        'enable_cluster_monitoring': True,
+    }
+    kafka_consumer_check = check(instance)
+    mock_kafka_client = seed_mock_kafka_client()
+    kafka_consumer_check.client = mock_kafka_client
+    kafka_consumer_check.metadata_collector.client = mock_kafka_client
+
+    _wire_cache(kafka_consumer_check)
+
+    collector = kafka_consumer_check.metadata_collector
+    topic_partitions = {'test-topic': [0, 1]}
+
+    first = collector.fetch_earliest_offsets(topic_partitions)
+    second = collector.fetch_earliest_offsets(topic_partitions)
+
+    expected = {('test-topic', 0): 10, ('test-topic', 1): 20}
+    assert first == expected
+    assert second == expected
+    assert mock_kafka_client.kafka_client.list_offsets.call_count == 1
+
+
+def test_fetch_earliest_offsets_refetches_when_cache_missing_partitions(check):
+    """A fresh cache that doesn't cover every requested partition triggers a full refetch, keeping the same TTL."""
+    instance = {
+        'kafka_connect_str': 'localhost:9092',
+        'enable_cluster_monitoring': True,
+    }
+    kafka_consumer_check = check(instance)
+    mock_kafka_client = seed_mock_kafka_client()
+    kafka_consumer_check.client = mock_kafka_client
+    kafka_consumer_check.metadata_collector.client = mock_kafka_client
+
+    collector = kafka_consumer_check.metadata_collector
+    expire_at = time.time() + 300
+    seed_payload = json.dumps({'expire_at': expire_at, 'offsets': [['test-topic', 0, 10]]})
+    _wire_cache(kafka_consumer_check, seed={collector.EARLIEST_OFFSETS_CACHE_KEY: seed_payload})
+
+    topic_partitions = {'test-topic': [0, 1]}
+    result = collector.fetch_earliest_offsets(topic_partitions)
+
+    assert result == {('test-topic', 0): 10, ('test-topic', 1): 20}
+    assert mock_kafka_client.kafka_client.list_offsets.call_count == 1
+
+    saved = json.loads(kafka_consumer_check.write_persistent_cache.call_args[0][1])
+    assert saved['expire_at'] == expire_at
+
+
 def test_schema_registry_oauth_oidc_token(check, dd_run_check, aggregator):
     """Test that OIDC OAuth token is fetched and passed as Bearer header for Schema Registry."""
     instance = {