kafka_consumer: cap total broker_timestamps history by a per-cluster memory budget (DataDog#24508)

* kafka_consumer: cap broker_timestamps retained history by a per-cluster budget

The DSM broker_timestamps cache grows unbounded with partition count (up to
timestamp_history_size samples for every partition), reaching multiple GiB on
large clusters. Bound it to a fixed entry budget (~0.5 GiB/cluster): the
per-partition history depth is scaled down as partition count grows
(budget // num_partitions, clamped to [2, timestamp_history_size]). Small
clusters keep full history; large clusters shrink to a shallow depth, so total
memory stays bounded.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add changelog entry

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: piochelepiotr

kafka_consumer/changelog.d/24508.fixed

@@ -0,0 +1 @@
+Cap the total DSM broker timestamps history per cluster with a memory budget that scales the per-partition depth by partition count, bounding cluster-check-runner memory on large clusters.

kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py

@@ -21,6 +21,10 @@
 
 MAX_TIMESTAMPS = 1000
 
+# Total broker-timestamp entries retained per cluster, ~0.5 GiB at ~89 bytes/entry. The
+# per-partition history is scaled down from this budget as the partition count grows.
+MAX_TIMESTAMP_ENTRIES = 6_000_000
+
 LAG_EXTRAPOLATION_LIMIT_SECONDS = 600
 
 
@@ -297,20 +301,27 @@ def _earliest_consumer_offsets(self, consumer_offsets):
                     earliest[topic_partition] = offset
         return earliest
 
+    def _max_history(self, num_partitions):
+        """Per-partition timestamp cap, scaled so total retained entries stay within the budget."""
+        if num_partitions <= 0:
+            return self._max_timestamps
+        return max(2, min(self._max_timestamps, MAX_TIMESTAMP_ENTRIES // num_partitions))
+
     def _add_broker_timestamps(self, broker_timestamps, highwater_offsets, prune_floors=None):
         prune_floors = prune_floors or {}
+        max_history = self._max_history(len(highwater_offsets))
         for (topic, partition), highwater_offset in highwater_offsets.items():
             timestamps = broker_timestamps["{}_{}".format(topic, partition)]
             # Reset detected: clear the whole cache. Low-offset survivors are from the
             # previous generation and VW pins the minimum endpoint, so they'd never age out.
             if any(o > highwater_offset for o in timestamps):
                 timestamps.clear()
             timestamps[highwater_offset] = time()
-            if len(timestamps) >= self._max_timestamps:
+            if len(timestamps) >= max_history:
                 prune_floor = prune_floors.get((topic, partition))
                 if prune_floor is not None:
                     _prune_below_anchor(timestamps, prune_floor)
-                _visvalingam_whyatt(timestamps, max(2, self._max_timestamps // 2))
+                _visvalingam_whyatt(timestamps, max(2, max_history // 2))
 
     def _save_broker_timestamps(self, broker_timestamps, persistent_cache_key):
         """Saves broker timestamps to persistent cache."""

kafka_consumer/tests/test_unit.py

@@ -5,6 +5,7 @@
 import json
 import logging
 import marshal
+from collections import defaultdict
 from contextlib import nullcontext as does_not_raise
 
 import mock
@@ -528,6 +529,30 @@ def test_check_clears_cache_on_partial_reset(kafka_instance, check, dd_run_check
     assert 100 in timestamps
 
 
+def test_max_history_scales_with_partition_count(check, kafka_instance):
+    kafka_consumer_check = check(kafka_instance)  # default timestamp_history_size = 1000
+    assert kafka_consumer_check._max_history(0) == 1000
+    assert kafka_consumer_check._max_history(1) == 1000
+    assert kafka_consumer_check._max_history(6000) == 1000  # 6M / 6000 = 1000 (breakeven)
+    assert kafka_consumer_check._max_history(60000) == 100  # 6M / 60000
+    assert kafka_consumer_check._max_history(6_000_000) == 2  # floored at 2
+
+
+def test_add_broker_timestamps_caps_total_by_budget(check, kafka_instance, monkeypatch):
+    import datadog_checks.kafka_consumer.kafka_consumer as kc
+
+    monkeypatch.setattr(kc, 'MAX_TIMESTAMP_ENTRIES', 1000)
+    kafka_consumer_check = check(kafka_instance)
+    # 100 partitions against a 1000-entry budget -> per-partition cap of 10.
+    highwater_offsets = {("topic1", p): 500 for p in range(100)}
+    broker_timestamps = defaultdict(dict)
+    broker_timestamps["topic1_0"] = {i: float(i) for i in range(60)}
+
+    kafka_consumer_check._add_broker_timestamps(broker_timestamps, highwater_offsets)
+
+    assert len(broker_timestamps["topic1_0"]) <= 10
+
+
 @pytest.mark.parametrize(
     'timestamp_history_size, initial_cache, highwater_offset, '
     'highwater_time, consumer_offset, expected_cache_size, expected_lag',