From 733ffdab12e7c1648f81717075252575f59e6cdc Mon Sep 17 00:00:00 2001 From: Piotr WOLSKI Date: Wed, 24 Jun 2026 11:03:29 +0200 Subject: [PATCH 1/3] kafka_consumer: always fetch highwater offsets when cluster monitoring is enabled (#24149) * kafka_consumer: always fetch highwater offsets when cluster monitoring is enabled When enable_cluster_monitoring is true the consumer context count can easily exceed the default max_partition_contexts (500), causing the check to skip highwater offset collection entirely. This silently zeros out kafka.topic.message_rate and stops kafka.broker_offset from being emitted, because _collect_topic_metadata receives an empty highwater_offsets dict. Bypass the context limit guard when cluster monitoring is active so that highwater offsets are always fetched; the existing per-metric context caps in report_highwater_offsets and report_consumer_offsets_and_lag still apply. Co-Authored-By: Claude Sonnet 4.6 * kafka_consumer: add changelog entry for PR #24149 Co-Authored-By: Claude Sonnet 4.6 * kafka_consumer: bypass context reporting limit when cluster monitoring is enabled When enable_cluster_monitoring is true, report all consumer lag and highwater offset metrics without capping at max_partition_contexts. Cluster monitoring users need full cluster visibility by design, so capping metric reporting makes no sense in that mode. Uses float('inf') as the reporting limit, which works correctly with the existing int comparisons in report_highwater_offsets and report_consumer_offsets_and_lag (int == inf is False, int >= inf is False, int < inf is True). Also suppresses the misleading "narrow your target" warning in cluster monitoring mode. Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Sonnet 4.6 --- kafka_consumer/changelog.d/24149.fixed | 1 + .../datadog_checks/kafka_consumer/kafka_consumer.py | 13 +++++++++---- 2 files changed, 10 insertions(+), 4 deletions(-) create mode 100644 kafka_consumer/changelog.d/24149.fixed diff --git a/kafka_consumer/changelog.d/24149.fixed b/kafka_consumer/changelog.d/24149.fixed new file mode 100644 index 0000000000000..933adbb7ea7e4 --- /dev/null +++ b/kafka_consumer/changelog.d/24149.fixed @@ -0,0 +1 @@ +Fix kafka.broker_offset and kafka.topic.message_rate not being collected when enable_cluster_monitoring is true and the consumer context count exceeds max_partition_contexts. diff --git a/kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py b/kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py index 6dea7e472c240..0ac2e213a5748 100644 --- a/kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py +++ b/kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py @@ -67,7 +67,9 @@ def check(self, _): persistent_cache_key = "broker_timestamps_" consumer_contexts_count = self.count_consumer_contexts(consumer_offsets) try: - if consumer_contexts_count < self._context_limit: + # Cluster monitoring always requires highwater offsets (for topic.message_rate and other + # cluster metadata metrics), so bypass the consumer context limit in that case. + if consumer_contexts_count < self._context_limit or self.config._cluster_monitoring_enabled: # Fetch highwater offsets # Build partitions list or use all if configured # If cluster monitoring is enabled, always fetch all broker highwater marks @@ -100,7 +102,10 @@ def check(self, _): consumer_offsets, highwater_offsets, ) - if total_contexts >= self._context_limit: + # When cluster monitoring is enabled, all offsets and lag metrics are reported regardless + # of context count so that the full cluster picture is always available. + reporting_limit = float('inf') if self.config._cluster_monitoring_enabled else self._context_limit + if total_contexts >= self._context_limit and not self.config._cluster_monitoring_enabled: self.warning( """Discovered %s metric contexts - this exceeds the maximum number of %s contexts permitted by the check. Please narrow your target by specifying in your kafka_consumer.yaml the consumer groups, topics @@ -113,11 +118,11 @@ def check(self, _): if self.config._kafka_cluster_id_override: cluster_id = self.config._kafka_cluster_id_override - self.report_highwater_offsets(highwater_offsets, self._context_limit, cluster_id) + self.report_highwater_offsets(highwater_offsets, reporting_limit, cluster_id) self.report_consumer_offsets_and_lag( consumer_offsets, highwater_offsets, - self._context_limit - len(highwater_offsets), + reporting_limit - len(highwater_offsets), broker_timestamps, cluster_id, ) From 534afff2b0c16bd0392643187c9dd280e1cf69b6 Mon Sep 17 00:00:00 2001 From: Piotr WOLSKI Date: Wed, 24 Jun 2026 14:23:42 +0200 Subject: [PATCH 2/3] Add consumer-group rebalance, empty-group, and metadata signals to kafka_consumer (#23915) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add consumer-group rebalance, empty-group, and metadata signals Enrich the cluster-monitoring consumer-group collection with signals that cannot be derived from existing tagged metrics: - kafka.consumer_group.rebalancing (1/0): detected via group state (PreparingRebalance/CompletingRebalance) for classic groups and via assignment != target_assignment for KIP-848 consumer-protocol groups. - kafka.consumer_group.empty (1/0): 1 when the group is in the EMPTY state (committed offsets but no active members). Dimensional metadata is added as tags on existing gauges rather than new metrics: partition_assignor, consumer_group_type, and is_simple_consumer_group on consumer_group.members, and group_instance_id (static membership) on consumer_group.member.partitions. Co-Authored-By: Claude Opus 4.8 (1M context) * Add changelog entry for PR #23915 * Address review feedback on consumer-group signals - Extract _build_group_meta_tags helper from the collection loop. - Use `is not None` guards for partition_assignor and group_instance_id so empty-string values are not silently dropped. - Emit consumer_group.rebalancing and consumer_group.empty with the same group_meta_tags as consumer_group.members so the sibling gauges share a tag set and can be correlated in dashboards. - Reduce test mock duplication: _collect_groups now reuses seed_mock_kafka_client and a shared _stub_consumer_groups helper. - Add tests for the dimensional-tag omission path and the no-target-assignment rebalance-skip branch. - Name the new tag keys in the README. Co-Authored-By: Claude Opus 4.8 (1M context) * Address round-2 review feedback on consumer-group signals - Revert partition_assignor guard to `if assignor:` so KIP-848 and EMPTY-state groups (which report an empty assignor) don't emit a blank-value partition_assignor: tag. Parametrize the absent-tags test to cover both None and "". - Type-hint state_name on _is_group_rebalancing. - Add comments documenting the member-level vs group-level tag-set choice and the EMPTY-state basis for consumer_group.empty. Co-Authored-By: Claude Opus 4.8 (1M context) * Add kafka.consumer_group.membership_changes count metric Caches a hash of sorted member IDs per consumer group after each check run. Emits consumer_group.membership_changes (+1 count) whenever the hash differs from the previous run, catching rebalances that complete between two polling intervals and would otherwise be invisible to the rebalancing gauge. Co-Authored-By: Claude Sonnet 4.6 * Remove kafka.consumer_group.empty — redundant with consumer_group_state tag The EMPTY state is already visible as consumer_group_state:EMPTY on the rebalancing metric, making a dedicated gauge unnecessary. Co-Authored-By: Claude Sonnet 4.6 * Fix ruff formatting in cluster_metadata.py Co-Authored-By: Claude Sonnet 4.6 * Address review feedback on consumer-group signals - Extract _load_member_hashes_cache / _save_member_hashes_cache helpers to match the _load_*/_save_* pattern used by every other cache in the class - Use getattr for m.member_id to avoid TypeError on missing/null values - Fix README and changelog: replace "empty-group detection" with accurate description of rebalance detection and membership-change counting; note that empty groups are visible via the consumer_group_state:EMPTY tag - Add three unit tests covering membership_changes: no prior cache (no emit), unchanged members (no emit), changed members (emit once) Co-Authored-By: Claude Sonnet 4.6 * Remove consumer-group signals demo compose file Co-Authored-By: Claude Sonnet 4.6 * Update kafka_consumer/README.md Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/README.md Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/README.md Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/README.md Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/metadata.csv Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Update kafka_consumer/metadata.csv Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> * Address review comments on consumer-group signals PR - Fix _is_group_rebalancing to return True when assignment is None but target_assignment is present (member has target but no current assignment is unambiguous drift); add test covering this case. - Replace comma-joined member-ID hash with json.dumps to prevent delimiter collisions (e.g. ['a,b','c'] and ['a','b,c'] now produce distinct hashes); update expected hash in existing test. - Validate cache decoded from JSON is a dict; discard and continue if not to prevent AttributeError on .get(); add test seeding a list-shaped cache and asserting group gauges still emit. - Update changelog entry to name both new DSM-only metrics explicitly. Co-Authored-By: Claude Sonnet 4.6 * Fix ruff formatting in cluster_metadata.py Co-Authored-By: Claude Sonnet 4.6 * Update changelog entry wording Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: domalessi <111786334+domalessi@users.noreply.github.com> --- kafka_consumer/README.md | 8 +- kafka_consumer/changelog.d/23915.added | 1 + .../kafka_consumer/cluster_metadata.py | 94 +++++- kafka_consumer/metadata.csv | 2 + kafka_consumer/tests/test_cluster_metadata.py | 313 ++++++++++++++++++ 5 files changed, 406 insertions(+), 12 deletions(-) create mode 100644 kafka_consumer/changelog.d/23915.added diff --git a/kafka_consumer/README.md b/kafka_consumer/README.md index b4d7ca358978d..681ee26a33c1a 100644 --- a/kafka_consumer/README.md +++ b/kafka_consumer/README.md @@ -62,10 +62,10 @@ instances: When `enable_cluster_monitoring` is enabled, the integration collects cluster-wide metrics for [Data Streams Monitoring][18] in addition to consumer lag: -- **Brokers**: Configuration and health metrics -- **Topics and partitions**: Sizes, offsets, and replication status -- **Consumer groups**: Member details and group state -- **Schema registry**: Schema metadata (requires `schema_registry_url`) +- **Brokers**: Configuration and health metrics. +- **Topics and partitions**: Sizes, offsets, and replication status. +- **Consumer groups**: Member details, group state, rebalance detection, membership-change counting, and metadata exposed as tags (`partition_assignor`, `consumer_group_type`, `is_simple_consumer_group`, and `group_instance_id`). Empty groups are visible through the `consumer_group_state:EMPTY` tag on `kafka.consumer_group.members`. +- **Schema registry**: Schema metadata (requires `schema_registry_url`). #### Batched collection diff --git a/kafka_consumer/changelog.d/23915.added b/kafka_consumer/changelog.d/23915.added new file mode 100644 index 0000000000000..c459489ec863e --- /dev/null +++ b/kafka_consumer/changelog.d/23915.added @@ -0,0 +1 @@ +Add consumer group rebalance detection and membership-change counting metrics, plus partition assignor, group type, simple-group, and static-membership tags, when cluster monitoring is enabled. (DSM only) diff --git a/kafka_consumer/datadog_checks/kafka_consumer/cluster_metadata.py b/kafka_consumer/datadog_checks/kafka_consumer/cluster_metadata.py index 9f282e68fb3b9..684dc5b02c002 100644 --- a/kafka_consumer/datadog_checks/kafka_consumer/cluster_metadata.py +++ b/kafka_consumer/datadog_checks/kafka_consumer/cluster_metadata.py @@ -18,6 +18,8 @@ from datadog_checks.kafka_consumer.constants import KAFKA_INTERNAL_TOPICS +CONSUMER_GROUP_REBALANCING_STATES = frozenset({'PREPARING_REBALANCING', 'COMPLETING_REBALANCING'}) + class SchemaDefinition(TypedDict): schema: str @@ -75,6 +77,7 @@ def __init__(self, check, client, config, log): 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' + self.CONSUMER_GROUP_MEMBERS_CACHE_KEY = 'kafka_consumer_group_members_cache' self.SCHEMA_CACHE_KEY = 'kafka_schema_cache' self.SCHEMA_VERSION_CHECK_CACHE_KEY = 'kafka_schema_version_check_cache' self.SCHEMA_COMPATIBILITY_FETCH_CACHE_KEY = 'kafka_schema_compatibility_fetch_cache' @@ -848,6 +851,9 @@ def _collect_consumer_group_metadata(self, metadata): except Exception as e: self.log.warning("Error getting consumer group details for %s: %s", group_id, e) + prev_member_hashes = self._load_member_hashes_cache() + current_member_hashes = {} + for group_id, group_info in group_id_to_info.items(): group_tags = self._get_tags(cluster_id) + [f'consumer_group:{group_id}'] state = group_info.state @@ -858,29 +864,101 @@ def _collect_consumer_group_metadata(self, metadata): if coordinator: state_tags.append(f'coordinator:{coordinator.id}') - self.check.gauge('consumer_group.members', len(members), tags=state_tags) + # All group-level gauges share the same tag set so they can be correlated in dashboards. + group_meta_tags = self._build_group_meta_tags(state_tags, group_info) + + self.check.gauge('consumer_group.members', len(members), tags=group_meta_tags) + self.check.gauge( + 'consumer_group.rebalancing', + 1 if self._is_group_rebalancing(state_name, members) else 0, + tags=group_meta_tags, + ) + + member_ids = sorted(getattr(m, 'member_id', '') or '' for m in members) + member_hash = hashlib.sha256(json.dumps(member_ids, separators=(',', ':')).encode()).hexdigest() + current_member_hashes[group_id] = member_hash - member_info = [] - topics_for_group = set() + if prev_member_hashes is not None: + prev_hash = prev_member_hashes.get(group_id) + if prev_hash is not None and prev_hash != member_hash: + self.check.count('consumer_group.membership_changes', 1, tags=group_meta_tags) for member in members: - member_id = member.member_id client_id = member.client_id host = member.host - member_info.append({'member_id': member_id, 'client_id': client_id, 'host': host}) - if hasattr(member, 'assignment') and member.assignment: partition_count = len(member.assignment.topic_partitions) + # Member-level gauges deliberately use state_tags, not group_meta_tags: the + # group-level dimensional tags are omitted here to keep per-member cardinality bounded. member_tags = state_tags + [ f'client_id:{client_id}', f'member_host:{host}', ] + group_instance_id = getattr(member, 'group_instance_id', None) + if group_instance_id is not None: + member_tags.append(f'group_instance_id:{group_instance_id}') self.check.gauge('consumer_group.member.partitions', partition_count, tags=member_tags) - for tp in member.assignment.topic_partitions: - topics_for_group.add(tp.topic) + self._save_member_hashes_cache(current_member_hashes) + + def _load_member_hashes_cache(self) -> dict[str, str] | None: + """Return the previous member-hash map, or None if unreadable.""" + try: + cached = self.check.read_persistent_cache(self.CONSUMER_GROUP_MEMBERS_CACHE_KEY) + if not cached: + return None + result = json.loads(cached) + if not isinstance(result, dict): + self.log.debug("Consumer group members cache has unexpected shape; discarding") + return None + return result + except Exception as e: + self.log.debug("Could not read consumer group members cache: %s", e) + return None + + def _save_member_hashes_cache(self, hashes: dict[str, str]) -> None: + """Persist the current member-hash map.""" + try: + self.check.write_persistent_cache(self.CONSUMER_GROUP_MEMBERS_CACHE_KEY, json.dumps(hashes)) + except Exception as e: + self.log.debug("Could not write consumer group members cache: %s", e) + + def _build_group_meta_tags(self, state_tags: list[str], group_info) -> list[str]: + """Build the group-level tag list, appending dimensional metadata when the broker provides it.""" + tags = list(state_tags) + assignor = getattr(group_info, 'partition_assignor', None) + # KIP-848 and EMPTY-state groups report an empty assignor; skip it to avoid a blank-value tag. + if assignor: + tags.append(f'partition_assignor:{assignor}') + group_type = getattr(group_info, 'type', None) + if group_type is not None: + type_name = group_type.name if hasattr(group_type, 'name') else str(group_type) + tags.append(f'consumer_group_type:{type_name}') + is_simple = getattr(group_info, 'is_simple_consumer_group', None) + if is_simple is not None: + tags.append(f'is_simple_consumer_group:{str(bool(is_simple)).lower()}') + return tags + + def _is_group_rebalancing(self, state_name: str, members) -> bool: + """Detect an in-progress rebalance via group state (classic) or assignment drift (KIP-848).""" + if state_name in CONSUMER_GROUP_REBALANCING_STATES: + return True + for member in members: + target = getattr(member, 'target_assignment', None) + if target is None: + # Classic-protocol member: no KIP-848 target, skip. + continue + assignment = getattr(member, 'assignment', None) + if assignment is None: + # Member has a target but no current assignment — unambiguous drift. + return True + current_tps = {(tp.topic, tp.partition) for tp in assignment.topic_partitions} + target_tps = {(tp.topic, tp.partition) for tp in target.topic_partitions} + if current_tps != target_tps: + return True + return False def _load_schema_id_cache(self) -> dict[str, SchemaDefinition]: """Load the permanent schema ID cache from persistent storage. diff --git a/kafka_consumer/metadata.csv b/kafka_consumer/metadata.csv index 7ecfe19e3d3d4..6e7b10e0561df 100644 --- a/kafka_consumer/metadata.csv +++ b/kafka_consumer/metadata.csv @@ -15,6 +15,8 @@ kafka.cluster.controller_id,gauge,,instance,,ID of the broker acting as the clus kafka.consumer_group.count,gauge,,item,,Total number of consumer groups. (DSM only),0,kafka_consumer,consumer groups,, kafka.consumer_group.member.partitions,gauge,,item,,Number of partitions assigned to this consumer group member. (DSM only),0,kafka_consumer,member partitions,, kafka.consumer_group.members,gauge,,item,,Number of members in the consumer group. (DSM only),0,kafka_consumer,group members,, +kafka.consumer_group.membership_changes,count,,,,Number of times the consumer group membership changed between check runs. (DSM only),-1,kafka_consumer,membership changes,, +kafka.consumer_group.rebalancing,gauge,,,,Whether the consumer group is rebalancing (1) or stable (0). (DSM only),-1,kafka_consumer,rebalancing,, kafka.consumer_lag,gauge,,message,,Lag in messages between consumer and broker.,-1,kafka_consumer,consumer lag,, kafka.consumer_offset,gauge,,offset,,Current message offset on consumer.,0,kafka_consumer,consumer offset,, kafka.estimated_consumer_lag,gauge,,second,,Lag in seconds between consumer and broker. This metric is provided through Data Streams Monitoring. Additional charges may apply.,-1,kafka_consumer,consumer time lag,, diff --git a/kafka_consumer/tests/test_cluster_metadata.py b/kafka_consumer/tests/test_cluster_metadata.py index ea2071bf1c9a5..878625d41b868 100644 --- a/kafka_consumer/tests/test_cluster_metadata.py +++ b/kafka_consumer/tests/test_cluster_metadata.py @@ -120,11 +120,20 @@ def mock_describe_configs(resources): state_mock.name = 'STABLE' describe_result.state = state_mock + # Group-level metadata used for dimensional tags + describe_result.partition_assignor = 'range' + describe_result.is_simple_consumer_group = False + type_mock = mock.MagicMock() + type_mock.name = 'CLASSIC' + describe_result.type = type_mock + # Mock member member = mock.MagicMock() member.member_id = 'm1' member.client_id = 'c1' member.host = 'h1' + member.group_instance_id = None + member.target_assignment = None # Mock assignment with topic_partitions assignment = mock.MagicMock() @@ -507,6 +516,24 @@ def mocked_read_cache(key): 'consumer_group:test-group', 'consumer_group_state:STABLE', 'coordinator:1', + 'partition_assignor:range', + 'consumer_group_type:CLASSIC', + 'is_simple_consumer_group:false', + ], + ) + + aggregator.assert_metric( + 'kafka.consumer_group.rebalancing', + value=0, + tags=[ + 'test_tag:test_value', + 'kafka_cluster_id:test-cluster-id', + 'consumer_group:test-group', + 'consumer_group_state:STABLE', + 'coordinator:1', + 'partition_assignor:range', + 'consumer_group_type:CLASSIC', + 'is_simple_consumer_group:false', ], ) @@ -1547,3 +1574,289 @@ def test_schema_registry_none_compat_in_cache_omits_field(check, dd_run_check, a assert len(schema_events) == 1 assert 'compatibility' not in schema_events[0] assert 'global_compatibility' not in schema_events[0] + + +def _tp(topic, partition): + tp = mock.MagicMock() + tp.topic = topic + tp.partition = partition + return tp + + +def _make_assignment(tps): + if tps is None: + return None + assignment = mock.MagicMock() + assignment.topic_partitions = [_tp(t, p) for t, p in tps] + return assignment + + +def _make_member( + client_id='c1', host='h1', assignment_tps=(('test-topic', 0),), target_tps=None, group_instance_id=None +): + member = mock.MagicMock() + member.member_id = f'm-{client_id}' + member.client_id = client_id + member.host = host + member.group_instance_id = group_instance_id + member.assignment = _make_assignment(assignment_tps) + member.target_assignment = _make_assignment(target_tps) + return member + + +def _make_group_describe( + state_name='STABLE', assignor='range', is_simple=False, group_type='CONSUMER', members=(), coordinator_id=1 +): + describe_result = mock.MagicMock() + state_mock = mock.MagicMock() + state_mock.name = state_name + describe_result.state = state_mock + describe_result.partition_assignor = assignor + describe_result.is_simple_consumer_group = is_simple + if group_type is None: + describe_result.type = None + else: + type_mock = mock.MagicMock() + type_mock.name = group_type + describe_result.type = type_mock + coordinator_mock = mock.MagicMock() + coordinator_mock.id = coordinator_id + describe_result.coordinator = coordinator_mock + describe_result.members = list(members) + return describe_result + + +def _stub_consumer_groups(admin, describe_by_group): + """Wire list_consumer_groups + describe_consumer_groups futures on a mock admin client.""" + list_result = mock.MagicMock() + list_result.errors = [] + list_result.valid = [mock.MagicMock(group_id=gid) for gid in describe_by_group] + list_future = mock.MagicMock() + list_future.result.return_value = list_result + admin.list_consumer_groups.return_value = list_future + + futures = {} + for gid, describe_result in describe_by_group.items(): + future = mock.MagicMock() + future.result.return_value = describe_result + futures[gid] = future + admin.describe_consumer_groups.return_value = futures + + +def _collect_groups(check, describe_result, group_id='test-group'): + """Run _collect_consumer_group_metadata against a single mocked consumer group. + + Reuses the shared seed_mock_kafka_client wiring and only swaps in the + consumer-group futures, so the admin-client mock setup is not duplicated. + """ + instance = {'kafka_connect_str': 'localhost:9092', 'enable_cluster_monitoring': True} + kafka_consumer_check = check(instance) + + mock_client = seed_mock_kafka_client() + _stub_consumer_groups(mock_client.kafka_client, {group_id: describe_result}) + kafka_consumer_check.metadata_collector.client = mock_client + + metadata = mock.MagicMock() + metadata.cluster_id = 'test-cluster-id' + kafka_consumer_check.metadata_collector._collect_consumer_group_metadata(metadata) + return kafka_consumer_check + + +def _collect_groups_with_cache(check, describe_result, seed=None, group_id='test-group'): + """Like _collect_groups but wires the persistent cache so membership-change logic is exercised.""" + instance = {'kafka_connect_str': 'localhost:9092', 'enable_cluster_monitoring': True} + kafka_consumer_check = check(instance) + + mock_client = seed_mock_kafka_client() + _stub_consumer_groups(mock_client.kafka_client, {group_id: describe_result}) + kafka_consumer_check.metadata_collector.client = mock_client + _wire_cache(kafka_consumer_check, seed) + + metadata = mock.MagicMock() + metadata.cluster_id = 'test-cluster-id' + kafka_consumer_check.metadata_collector._collect_consumer_group_metadata(metadata) + return kafka_consumer_check + + +def test_consumer_group_rebalancing_state_based(check, aggregator): + """A group in a rebalancing state reports rebalancing=1 (classic protocol).""" + describe_result = _make_group_describe(state_name='PREPARING_REBALANCING', members=[_make_member()]) + _collect_groups(check, describe_result) + aggregator.assert_metric( + 'kafka.consumer_group.rebalancing', + value=1, + tags=[ + 'kafka_cluster_id:test-cluster-id', + 'consumer_group:test-group', + 'consumer_group_state:PREPARING_REBALANCING', + 'coordinator:1', + 'partition_assignor:range', + 'consumer_group_type:CONSUMER', + 'is_simple_consumer_group:false', + ], + ) + + +def test_consumer_group_rebalancing_target_assignment(check, aggregator): + """A stable group whose assignment != target_assignment reports rebalancing=1 (KIP-848).""" + member = _make_member(assignment_tps=[('test-topic', 0)], target_tps=[('test-topic', 0), ('test-topic', 1)]) + describe_result = _make_group_describe(state_name='STABLE', members=[member]) + _collect_groups(check, describe_result) + aggregator.assert_metric('kafka.consumer_group.rebalancing', value=1) + + +def test_consumer_group_not_rebalancing_when_assignment_matches_target(check, aggregator): + """A stable group whose assignment == target_assignment reports rebalancing=0.""" + member = _make_member(assignment_tps=[('test-topic', 0)], target_tps=[('test-topic', 0)]) + describe_result = _make_group_describe(state_name='STABLE', members=[member]) + _collect_groups(check, describe_result) + aggregator.assert_metric('kafka.consumer_group.rebalancing', value=0) + + +def test_consumer_group_not_rebalancing_when_no_target_assignment(check, aggregator): + """A stable classic-protocol member (no target_assignment) is skipped, reporting rebalancing=0.""" + member = _make_member(assignment_tps=[('test-topic', 0)], target_tps=None) + describe_result = _make_group_describe(state_name='STABLE', members=[member]) + _collect_groups(check, describe_result) + aggregator.assert_metric('kafka.consumer_group.rebalancing', value=0) + + +def test_consumer_group_dimensional_tags(check, aggregator): + """Group-level metadata is attached as tags on consumer_group.members.""" + describe_result = _make_group_describe( + state_name='STABLE', + assignor='cooperative-sticky', + is_simple=True, + group_type='CONSUMER', + members=[_make_member()], + ) + _collect_groups(check, describe_result) + aggregator.assert_metric( + 'kafka.consumer_group.members', + value=1, + tags=[ + 'kafka_cluster_id:test-cluster-id', + 'consumer_group:test-group', + 'consumer_group_state:STABLE', + 'coordinator:1', + 'partition_assignor:cooperative-sticky', + 'consumer_group_type:CONSUMER', + 'is_simple_consumer_group:true', + ], + ) + + +@pytest.mark.parametrize('assignor', [None, ''], ids=['none', 'empty_string']) +def test_consumer_group_dimensional_tags_absent_when_unset(check, aggregator, assignor): + """When the broker reports no assignor (None or empty for KIP-848 groups), no dimensional tags are attached.""" + describe_result = _make_group_describe( + state_name='STABLE', assignor=assignor, is_simple=None, group_type=None, members=[_make_member()] + ) + _collect_groups(check, describe_result) + aggregator.assert_metric( + 'kafka.consumer_group.members', + value=1, + tags=[ + 'kafka_cluster_id:test-cluster-id', + 'consumer_group:test-group', + 'consumer_group_state:STABLE', + 'coordinator:1', + ], + ) + + +def test_consumer_group_member_static_membership_tag(check, aggregator): + """A member with a group_instance_id is tagged as a static member.""" + member = _make_member(group_instance_id='static-1') + describe_result = _make_group_describe(state_name='STABLE', members=[member]) + _collect_groups(check, describe_result) + aggregator.assert_metric( + 'kafka.consumer_group.member.partitions', + value=1, + tags=[ + 'kafka_cluster_id:test-cluster-id', + 'consumer_group:test-group', + 'consumer_group_state:STABLE', + 'coordinator:1', + 'client_id:c1', + 'member_host:h1', + 'group_instance_id:static-1', + ], + ) + + +def test_membership_changes_not_emitted_on_first_run(check, aggregator): + """No membership_changes on first run — no prior cache to compare against.""" + describe_result = _make_group_describe(members=[_make_member()]) + _collect_groups_with_cache(check, describe_result) + aggregator.assert_metric('kafka.consumer_group.membership_changes', count=0) + + +def test_membership_changes_not_emitted_when_members_unchanged(check, aggregator): + """No membership_changes when the member set is identical to the previous run.""" + member = _make_member(client_id='c1') + prev_hash = hashlib.sha256(b'["m-c1"]').hexdigest() + cache_key = 'kafka_consumer_group_members_cache' + describe_result = _make_group_describe(members=[member]) + _collect_groups_with_cache( + check, + describe_result, + seed={cache_key: json.dumps({'test-group': prev_hash})}, + ) + aggregator.assert_metric('kafka.consumer_group.membership_changes', count=0) + + +def test_membership_changes_emitted_when_members_differ(check, aggregator): + """membership_changes fires exactly once when the member set differs from the prior run.""" + cache_key = 'kafka_consumer_group_members_cache' + old_hash = hashlib.sha256(b'["m-old"]').hexdigest() + describe_result = _make_group_describe(members=[_make_member(client_id='new')]) + _collect_groups_with_cache( + check, + describe_result, + seed={cache_key: json.dumps({'test-group': old_hash})}, + ) + aggregator.assert_metric( + 'kafka.consumer_group.membership_changes', + value=1, + count=1, + tags=[ + 'kafka_cluster_id:test-cluster-id', + 'consumer_group:test-group', + 'consumer_group_state:STABLE', + 'coordinator:1', + 'partition_assignor:range', + 'consumer_group_type:CONSUMER', + 'is_simple_consumer_group:false', + ], + ) + + +def test_consumer_group_rebalancing_when_assignment_none_but_target_present(check, aggregator): + """A KIP-848 member with no current assignment but a non-empty target reports rebalancing=1.""" + member = _make_member(assignment_tps=None, target_tps=[('orders', 0)]) + describe_result = _make_group_describe(state_name='STABLE', members=[member]) + _collect_groups(check, describe_result) + aggregator.assert_metric('kafka.consumer_group.rebalancing', value=1) + + +def test_membership_hash_delimiter_collision(check, aggregator): + """Member IDs that share characters with the delimiter produce distinct hashes.""" + ids_a = ['a,b', 'c'] + ids_b = ['a', 'b,c'] + hash_a = hashlib.sha256(json.dumps(sorted(ids_a), separators=(',', ':')).encode()).hexdigest() + hash_b = hashlib.sha256(json.dumps(sorted(ids_b), separators=(',', ':')).encode()).hexdigest() + assert hash_a != hash_b + + +def test_malformed_cache_does_not_abort_collection(check, aggregator): + """A non-dict cache value is silently discarded and group gauges are still emitted.""" + cache_key = 'kafka_consumer_group_members_cache' + describe_result = _make_group_describe(members=[_make_member()]) + _collect_groups_with_cache( + check, + describe_result, + seed={cache_key: json.dumps([])}, # list instead of dict + ) + aggregator.assert_metric('kafka.consumer_group.members', count=1) + aggregator.assert_metric('kafka.consumer_group.membership_changes', count=0) From bd332d9e51cefb4ab88068c6fa657d2f19a7140e Mon Sep 17 00:00:00 2001 From: Piotr WOLSKI Date: Wed, 24 Jun 2026 15:40:51 +0200 Subject: [PATCH 3/3] [kafka_actions] Bound message reads to a start-of-check snapshot (#24162) * [kafka_actions] Bound message reads to a start-of-check snapshot read_messages could hang until its global timeout whenever a selective filter matched fewer messages than n_messages_retrieved: once the consumer drained the existing backlog it kept polling the live head, and because a continuously-produced topic almost always delivers a message within the poll window, the "no more messages" (poll == None) exit never fired. Fix consumption to a snapshot of the log taken when the check starts: - Capture each partition's high watermark up front and never yield a message at or beyond it, so messages produced after the check began are excluded and live-tailing is impossible. - Enable enable.partition.eof and stop a partition on its EOF event or when its captured watermark is reached; return as soon as all are drained. - Reduce the default timeout from 20s to 5s (now only a safety net) and surface a hit_timeout stat so a truncated read is distinguishable from a complete one. Verified against a live 10-partition topic: the previously-hanging filtered read now returns in ~0.3s instead of 20s. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Piotr Wolski * [kafka_actions] Add changelog entry Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Piotr Wolski * [kafka_actions] Fix import grouping for ruff isort Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Piotr Wolski --------- Signed-off-by: Piotr Wolski Co-authored-by: Claude Opus 4.8 (1M context) --- kafka_actions/assets/configuration/spec.yaml | 6 +- kafka_actions/changelog.d/24162.fixed | 1 + .../datadog_checks/kafka_actions/check.py | 13 +- .../kafka_actions/data/conf.yaml.example | 6 +- .../kafka_actions/kafka_client.py | 127 +++++++---- kafka_actions/tests/test_unit.py | 197 +++++++++++------- 6 files changed, 237 insertions(+), 113 deletions(-) create mode 100644 kafka_actions/changelog.d/24162.fixed diff --git a/kafka_actions/assets/configuration/spec.yaml b/kafka_actions/assets/configuration/spec.yaml index 348e5abc8e4eb..b43f0955f03b9 100644 --- a/kafka_actions/assets/configuration/spec.yaml +++ b/kafka_actions/assets/configuration/spec.yaml @@ -205,8 +205,10 @@ files: - name: read_messages description: | Configuration for reading messages from Kafka topics. - Messages are streamed in real-time and sent to Datadog as they arrive. - The check has a 20-second timeout for the entire operation. + Only messages already present in the log when the check starts are read; the + per-partition high watermark is captured up front and consumption stops once every + partition is drained, so messages produced after the check began are never returned. + A 5-second timeout bounds the entire operation as a safety net. Supports JSON, BSON, Protobuf, and Avro with optional Schema Registry integration. Filtering is applied after deserialization using jq-style expressions. fleet_configurable: true diff --git a/kafka_actions/changelog.d/24162.fixed b/kafka_actions/changelog.d/24162.fixed new file mode 100644 index 0000000000000..15b66c7e54006 --- /dev/null +++ b/kafka_actions/changelog.d/24162.fixed @@ -0,0 +1 @@ +Fix `read_messages` hanging until the global timeout when a filter matched fewer messages than `n_messages_retrieved`. Consumption is now bounded to a snapshot of the log taken when the check starts (per-partition high watermark + `enable.partition.eof`), the default timeout is reduced from 20s to 5s, and a `hit_timeout` stat distinguishes a truncated read from a complete one. diff --git a/kafka_actions/datadog_checks/kafka_actions/check.py b/kafka_actions/datadog_checks/kafka_actions/check.py index 36749930be796..e5aee37daa984 100644 --- a/kafka_actions/datadog_checks/kafka_actions/check.py +++ b/kafka_actions/datadog_checks/kafka_actions/check.py @@ -267,7 +267,7 @@ def _action_read_messages(self): start_timestamp = config.get('start_timestamp') n_messages_retrieved = config.get('n_messages_retrieved', 10) max_scanned_messages = config.get('max_scanned_messages', 1000) - timeout_ms = config.get('timeout_ms', 20000) + timeout_ms = config.get('timeout_ms', 5000) filter_expression = config.get('filter', '') consumer_group_id = config.get('consumer_group_id') or f"datadog-agent-{self.remote_config_id}" @@ -326,6 +326,8 @@ def _action_read_messages(self): if scanned_count >= max_scanned_messages and sent_count < n_messages_retrieved: hit_scan_limit = True + hit_timeout = self.kafka_client.hit_timeout and not hit_retrieved_limit and not hit_scan_limit + elapsed_time = time.time() - start_time stats = { @@ -337,6 +339,7 @@ def _action_read_messages(self): 'messages_filtered_out': filtered_out_count, 'hit_scan_limit': hit_scan_limit, 'hit_retrieved_limit': hit_retrieved_limit, + 'hit_timeout': hit_timeout, 'elapsed_time_seconds': round(elapsed_time, 3), 'n_messages_retrieved': n_messages_retrieved, 'max_scanned_messages': max_scanned_messages, @@ -358,6 +361,14 @@ def _action_read_messages(self): sent_count, ) + if hit_timeout: + self.log.warning( + "Hit the %dms timeout after scanning %d messages and retrieving %d. Result may be incomplete.", + timeout_ms, + scanned_count, + sent_count, + ) + return stats def _evaluate_filter(self, filter_expression: str, deserialized_msg: DeserializedMessage) -> bool: diff --git a/kafka_actions/datadog_checks/kafka_actions/data/conf.yaml.example b/kafka_actions/datadog_checks/kafka_actions/data/conf.yaml.example index 0836927e93555..a2528c7c0d42d 100644 --- a/kafka_actions/datadog_checks/kafka_actions/data/conf.yaml.example +++ b/kafka_actions/datadog_checks/kafka_actions/data/conf.yaml.example @@ -197,8 +197,10 @@ instances: ## @param read_messages - mapping - optional ## Configuration for reading messages from Kafka topics. - ## Messages are streamed in real-time and sent to Datadog as they arrive. - ## The check has a 20-second timeout for the entire operation. + ## Only messages already present in the log when the check starts are read; the + ## per-partition high watermark is captured up front and consumption stops once every + ## partition is drained, so messages produced after the check began are never returned. + ## A 5-second timeout bounds the entire operation as a safety net. ## Supports JSON, BSON, Protobuf, and Avro with optional Schema Registry integration. ## Filtering is applied after deserialization using jq-style expressions. # diff --git a/kafka_actions/datadog_checks/kafka_actions/kafka_client.py b/kafka_actions/datadog_checks/kafka_actions/kafka_client.py index 4b25b1c100600..0ab39818747af 100644 --- a/kafka_actions/datadog_checks/kafka_actions/kafka_client.py +++ b/kafka_actions/datadog_checks/kafka_actions/kafka_client.py @@ -32,6 +32,8 @@ def __init__(self, config: KafkaActionsConfig, log): self.consumer = None self.producer = None self.admin_client = None + # True when consume_messages stopped on the timeout rather than draining all partitions. + self.hit_timeout = False def _get_authentication_config(self) -> dict[str, Any]: """Build authentication configuration for librdkafka.""" @@ -134,6 +136,8 @@ def get_consumer(self, group_id: str = 'kafka_actions') -> Consumer: 'group.id': group_id, 'auto.offset.reset': 'earliest', 'enable.auto.commit': False, + # Signal end-of-partition via a _PARTITION_EOF event so we stop once drained. + 'enable.partition.eof': True, } ) self.consumer = Consumer(config) @@ -193,13 +197,16 @@ def consume_messages( start_offset: int = -2, start_timestamp: int | None = None, max_messages: int = 1000, - timeout_ms: int = 30000, + timeout_ms: int = 5000, group_id: str = 'kafka_actions', ): - """Consume messages from a Kafka topic, yielding them as they arrive. + """Consume the messages already present in a topic, yielding them as they are read. - This is a generator that yields messages in real-time as they're consumed, - allowing for immediate processing and sending to Datadog. + The per-partition high watermark is captured before consumption begins and no message + at or beyond it is yielded, so messages produced after the check starts are never + returned and the generator can't tail a live topic. A partition stops on EOF or when its + captured watermark is reached; the generator returns once all are drained. ``timeout_ms`` + is only a safety net. Args: topic: Topic name @@ -207,15 +214,17 @@ def consume_messages( start_offset: Starting offset (-1 for latest, -2 for earliest) start_timestamp: Starting timestamp in milliseconds since epoch. When set, start_offset is ignored. max_messages: Maximum messages to consume - timeout_ms: Global timeout in milliseconds for the entire consumption + timeout_ms: Safety-net timeout in milliseconds for the entire consumption group_id: Consumer group ID Yields: - Kafka messages as they arrive + Kafka messages that existed in the log when consumption began """ consumer = self.get_consumer(group_id) + admin = self.get_admin_client() start_time = time.time() global_timeout_s = timeout_ms / 1000.0 + self.hit_timeout = False try: if partition == -1: @@ -227,44 +236,42 @@ def consume_messages( else: partition_ids = [partition] - if start_timestamp is not None: - # Resolve timestamp to per-partition offsets using offsets_for_times. - timestamp_partitions = [TopicPartition(topic, p, start_timestamp) for p in partition_ids] - partitions = consumer.offsets_for_times(timestamp_partitions, timeout=10) - for tp in partitions: - if tp.offset != -1: - self.log.debug( - "Partition %d: timestamp %d resolved to offset %d", - tp.partition, - start_timestamp, - tp.offset, - ) - elif start_offset == -1: - # For "latest" offset, seek back from the high watermark to read the last N existing messages. - # Use AdminClient.list_offsets to fetch all high watermarks in a single batched call. - admin = self.get_admin_client() - offset_request = {TopicPartition(topic, p): OffsetSpec.latest() for p in partition_ids} - futures = admin.list_offsets(offset_request, request_timeout=10) - - partitions = [] - for tp, future in futures.items(): - result = future.result() - seek_offset = max(0, result.offset - max_messages) - partitions.append(TopicPartition(topic, tp.partition, seek_offset)) - self.log.debug("Partition %d: high=%d, seeking to %d", tp.partition, result.offset, seek_offset) - else: - partitions = [TopicPartition(topic, p, start_offset) for p in partition_ids] + # Snapshot each partition's high watermark; we never read at or beyond it. + end_request = {TopicPartition(topic, p): OffsetSpec.latest() for p in partition_ids} + end_futures = admin.list_offsets(end_request, request_timeout=10) + end_offsets = {tp.partition: future.result().offset for tp, future in end_futures.items()} - self.log.debug("Assigning partitions: %s", partitions) + start_offsets = self._resolve_start_offsets( + consumer, admin, topic, partition_ids, start_offset, start_timestamp, max_messages, end_offsets + ) + + # Assign only partitions that have messages in [start, high_watermark). + partitions = [] + active = set() + for p in partition_ids: + start = start_offsets.get(p, 0) + end = end_offsets.get(p, 0) + if start < end: + partitions.append(TopicPartition(topic, p, start)) + active.add(p) + else: + self.log.debug("Partition %d: nothing to read (start=%d, high=%d)", p, start, end) + + if not partitions: + self.log.debug("No messages to read for topic %s in [start, high-watermark)", topic) + return + + self.log.debug("Assigning partitions: %s (high watermarks: %s)", partitions, end_offsets) consumer.assign(partitions) consumed = 0 - while consumed < max_messages: + while consumed < max_messages and active: elapsed = time.time() - start_time remaining_timeout = global_timeout_s - elapsed if remaining_timeout <= 0: + self.hit_timeout = True self.log.debug("Global timeout reached after %d messages", consumed) break @@ -272,19 +279,28 @@ def consume_messages( msg = consumer.poll(timeout=poll_timeout) if msg is None: - self.log.debug("Poll returned None (no more messages available), stopping consumption") - break + # End-of-data arrives as an EOF event, not None; keep polling until drained. + continue if msg.error(): if msg.error().code() == KafkaError._PARTITION_EOF: - self.log.debug("Reached end of partition") + active.discard(msg.partition()) continue else: raise KafkaException(msg.error()) + p = msg.partition() + # Never surface a message at or beyond the captured high watermark. + if p not in active or msg.offset() >= end_offsets.get(p, 0): + active.discard(p) + continue + yield msg consumed += 1 + if msg.offset() >= end_offsets[p] - 1: + active.discard(p) + self.log.debug("Consumed %d messages from topic %s in %.2fs", consumed, topic, time.time() - start_time) finally: @@ -292,6 +308,41 @@ def consume_messages( consumer.close() self.consumer = None + def _resolve_start_offsets( + self, + consumer, + admin, + topic: str, + partition_ids: list[int], + start_offset: int, + start_timestamp: int | None, + max_messages: int, + end_offsets: dict[int, int], + ) -> dict[int, int]: + """Return a {partition: start_offset} map. A start at or beyond the high watermark + means there is nothing to read for that partition.""" + if start_timestamp is not None: + # An offset < 0 means the timestamp is past the end of the log: nothing to read. + timestamp_partitions = [TopicPartition(topic, p, start_timestamp) for p in partition_ids] + resolved = consumer.offsets_for_times(timestamp_partitions, timeout=10) + start_offsets = {} + for tp in resolved: + end = end_offsets.get(tp.partition, 0) + start_offsets[tp.partition] = tp.offset if tp.offset is not None and tp.offset >= 0 else end + return start_offsets + + if start_offset == -1: + # "latest": seek back from the high watermark to read the last N existing messages. + return {p: max(0, end_offsets.get(p, 0) - max_messages) for p in partition_ids} + + if start_offset == -2: + # "earliest": use the low watermark as the numeric start. + low_request = {TopicPartition(topic, p): OffsetSpec.earliest() for p in partition_ids} + low_futures = admin.list_offsets(low_request, request_timeout=10) + return {tp.partition: future.result().offset for tp, future in low_futures.items()} + + return dict.fromkeys(partition_ids, start_offset) + def produce_message( self, topic: str, diff --git a/kafka_actions/tests/test_unit.py b/kafka_actions/tests/test_unit.py index 76dfbc870ce77..46ba38d4a3a1f 100644 --- a/kafka_actions/tests/test_unit.py +++ b/kafka_actions/tests/test_unit.py @@ -4,15 +4,40 @@ import base64 import json -from unittest.mock import patch +import logging +from unittest.mock import MagicMock, patch import pytest +from confluent_kafka import KafkaError from datadog_checks.kafka_actions import KafkaActionsCheck +from datadog_checks.kafka_actions.kafka_client import KafkaActionsClient pytestmark = [pytest.mark.unit] +def _futures(offsets): + """Build a {topic-partition: future} map mimicking AdminClient.list_offsets results.""" + out = {} + for p, off in offsets.items(): + fut = MagicMock() + fut.result.return_value = MagicMock(offset=off) + out[MagicMock(partition=p)] = fut + return out + + +def _eof(partition): + """A mock poll() result representing a _PARTITION_EOF event.""" + msg = MagicMock() + msg.error.return_value = MagicMock(code=MagicMock(return_value=KafkaError._PARTITION_EOF)) + msg.partition.return_value = partition + return msg + + +def _client(): + return KafkaActionsClient({'kafka_connect_str': 'localhost:9092'}, logging.getLogger('test')) + + class MockKafkaMessage: """Mock confluent_kafka.Message for testing.""" @@ -385,107 +410,113 @@ def test_read_messages_nested_field_filter(self, aggregator, dd_run_check): class TestConsumeMessagesLatestOffset: - """Test that start_offset=-1 seeks back from high watermark instead of waiting at end.""" + """Test that start_offset=-1 seeks back from the captured high watermark.""" def test_latest_offset_seeks_back_from_high_watermark(self): - from unittest.mock import MagicMock - consumer = MagicMock() - consumer.poll.return_value = None - metadata = MagicMock() metadata.topics = {'t': MagicMock(partitions={0: MagicMock(), 1: MagicMock()})} consumer.list_topics.return_value = metadata + consumer.poll.side_effect = [_eof(0), _eof(1)] mock_admin = MagicMock() - tp0, tp1 = MagicMock(partition=0), MagicMock(partition=1) - f0, f1 = MagicMock(), MagicMock() - f0.result.return_value = MagicMock(offset=100) - f1.result.return_value = MagicMock(offset=200) - mock_admin.list_offsets.return_value = {tp0: f0, tp1: f1} - - import logging - - from datadog_checks.kafka_actions.kafka_client import KafkaActionsClient - - client = KafkaActionsClient({'kafka_connect_str': 'localhost:9092'}, logging.getLogger('test')) + mock_admin.list_offsets.return_value = _futures({0: 100, 1: 200}) + client = _client() with ( patch.object(client, 'get_consumer', return_value=consumer), patch.object(client, 'get_admin_client', return_value=mock_admin), ): list(client.consume_messages(topic='t', start_offset=-1, max_messages=10, timeout_ms=500)) - mock_admin.list_offsets.assert_called_once() assigned = {tp.partition: tp.offset for tp in consumer.assign.call_args[0][0]} assert assigned[0] == 90 # max(0, 100 - 10) assert assigned[1] == 190 # max(0, 200 - 10) +class TestConsumeMessagesSnapshotBound: + """Test that consumption never crosses the high watermark captured at the start.""" + + def test_does_not_yield_messages_at_or_beyond_high_watermark(self): + consumer = MagicMock() + metadata = MagicMock() + metadata.topics = {'t': MagicMock(partitions={0: MagicMock()})} + consumer.list_topics.return_value = metadata + # High watermark is 5; a message at offset 5 arrives after the snapshot and must be dropped. + consumer.poll.side_effect = [ + MockKafkaMessage(key=b'k', value=b'v', partition=0, offset=0), + MockKafkaMessage(key=b'k', value=b'v', partition=0, offset=1), + MockKafkaMessage(key=b'k', value=b'v', partition=0, offset=5), + ] + + mock_admin = MagicMock() + mock_admin.list_offsets.return_value = _futures({0: 5}) + + client = _client() + with ( + patch.object(client, 'get_consumer', return_value=consumer), + patch.object(client, 'get_admin_client', return_value=mock_admin), + ): + result = list(client.consume_messages(topic='t', start_offset=0, max_messages=1000, timeout_ms=500)) + + assert [m.offset() for m in result] == [0, 1] + + class TestConsumeMessagesStartTimestamp: """Test that start_timestamp resolves to per-partition offsets via offsets_for_times.""" def test_start_timestamp_resolves_offsets(self): - from unittest.mock import MagicMock - consumer = MagicMock() - consumer.poll.return_value = None - metadata = MagicMock() metadata.topics = {'t': MagicMock(partitions={0: MagicMock(), 1: MagicMock()})} consumer.list_topics.return_value = metadata + consumer.offsets_for_times.return_value = [ + MagicMock(partition=0, offset=50), + MagicMock(partition=1, offset=120), + ] + consumer.poll.side_effect = [_eof(0), _eof(1)] - # offsets_for_times returns TopicPartitions with resolved offsets - resolved_tp0 = MagicMock(partition=0, offset=50) - resolved_tp1 = MagicMock(partition=1, offset=120) - consumer.offsets_for_times.return_value = [resolved_tp0, resolved_tp1] - - import logging - - from datadog_checks.kafka_actions.kafka_client import KafkaActionsClient - - client = KafkaActionsClient({'kafka_connect_str': 'localhost:9092'}, logging.getLogger('test')) + mock_admin = MagicMock() + mock_admin.list_offsets.return_value = _futures({0: 100, 1: 200}) - with patch.object(client, 'get_consumer', return_value=consumer): + client = _client() + with ( + patch.object(client, 'get_consumer', return_value=consumer), + patch.object(client, 'get_admin_client', return_value=mock_admin), + ): list(client.consume_messages(topic='t', start_timestamp=1700000000000, max_messages=10, timeout_ms=500)) consumer.offsets_for_times.assert_called_once() - call_args = consumer.offsets_for_times.call_args[0][0] - assert len(call_args) == 2 - # Verify each TopicPartition was created with the timestamp - for tp in call_args: + for tp in consumer.offsets_for_times.call_args[0][0]: assert tp.offset == 1700000000000 + assigned = {tp.partition: tp.offset for tp in consumer.assign.call_args[0][0]} + assert assigned == {0: 50, 1: 120} - assigned = consumer.assign.call_args[0][0] - assert len(assigned) == 2 - - def test_start_timestamp_includes_partitions_at_end(self): - from unittest.mock import MagicMock - + def test_start_timestamp_skips_partitions_past_end(self): consumer = MagicMock() - consumer.poll.return_value = None - metadata = MagicMock() metadata.topics = {'t': MagicMock(partitions={0: MagicMock(), 1: MagicMock()})} consumer.list_topics.return_value = metadata + # Partition 0 has no message at/after the timestamp (offset=-1); it must not be assigned. + consumer.offsets_for_times.return_value = [ + MagicMock(partition=0, offset=-1), + MagicMock(partition=1, offset=120), + ] + consumer.poll.side_effect = [_eof(1)] - # Partition 0 has no messages at the timestamp (offset=-1 = OFFSET_END), partition 1 does - resolved_tp0 = MagicMock(partition=0, offset=-1) - resolved_tp1 = MagicMock(partition=1, offset=120) - consumer.offsets_for_times.return_value = [resolved_tp0, resolved_tp1] - - import logging - - from datadog_checks.kafka_actions.kafka_client import KafkaActionsClient - - client = KafkaActionsClient({'kafka_connect_str': 'localhost:9092'}, logging.getLogger('test')) + mock_admin = MagicMock() + mock_admin.list_offsets.return_value = _futures({0: 100, 1: 200}) - with patch.object(client, 'get_consumer', return_value=consumer): + client = _client() + with ( + patch.object(client, 'get_consumer', return_value=consumer), + patch.object(client, 'get_admin_client', return_value=mock_admin), + ): list(client.consume_messages(topic='t', start_timestamp=1700000000000, max_messages=10, timeout_ms=500)) - # Both partitions should be assigned; partition 0 waits at end for new messages assigned = consumer.assign.call_args[0][0] - assert len(assigned) == 2 + assert len(assigned) == 1 + assert assigned[0].partition == 1 def test_start_timestamp_overrides_start_offset(self, dd_run_check, aggregator): """Test that start_timestamp is passed through from check config.""" @@ -525,27 +556,53 @@ def test_start_timestamp_overrides_start_offset(self, dd_run_check, aggregator): class TestConsumeMessagesEarlyReturn: - """Test that consume_messages returns early on None poll or partition EOF.""" - - def test_none_poll_breaks_immediately(self): - from unittest.mock import MagicMock + """Test that consume_messages stops when all partitions are drained (EOF), not on None.""" + def test_eof_drains_and_stops(self): consumer = MagicMock() metadata = MagicMock() metadata.topics = {'t': MagicMock(partitions={0: MagicMock()})} consumer.list_topics.return_value = metadata - consumer.poll.side_effect = [MockKafkaMessage(key=b'k', value=b'v', partition=0, offset=0), None] - - import logging + # A None poll must NOT end consumption; only EOF (or the watermark) does. + consumer.poll.side_effect = [ + MockKafkaMessage(key=b'k', value=b'v', partition=0, offset=0), + None, + _eof(0), + ] - from datadog_checks.kafka_actions.kafka_client import KafkaActionsClient + mock_admin = MagicMock() + mock_admin.list_offsets.return_value = _futures({0: 100}) - client = KafkaActionsClient({'kafka_connect_str': 'localhost:9092'}, logging.getLogger('test')) - with patch.object(client, 'get_consumer', return_value=consumer): - result = list(client.consume_messages(topic='t', start_offset=0, max_messages=1000, timeout_ms=30000)) + client = _client() + with ( + patch.object(client, 'get_consumer', return_value=consumer), + patch.object(client, 'get_admin_client', return_value=mock_admin), + ): + result = list(client.consume_messages(topic='t', start_offset=0, max_messages=1000, timeout_ms=500)) assert len(result) == 1 - assert consumer.poll.call_count == 2 + assert consumer.poll.call_count == 3 + + def test_empty_range_returns_without_polling(self): + consumer = MagicMock() + metadata = MagicMock() + metadata.topics = {'t': MagicMock(partitions={0: MagicMock()})} + consumer.list_topics.return_value = metadata + + mock_admin = MagicMock() + # start (10) is already at the high watermark (10): nothing to read. + mock_admin.list_offsets.return_value = _futures({0: 10}) + + client = _client() + with ( + patch.object(client, 'get_consumer', return_value=consumer), + patch.object(client, 'get_admin_client', return_value=mock_admin), + ): + result = list(client.consume_messages(topic='t', start_offset=10, max_messages=1000, timeout_ms=500)) + + assert result == [] + consumer.assign.assert_not_called() + consumer.poll.assert_not_called() if __name__ == '__main__':