From ef22b84fbe78272a2781ff0fbc50726f2e64c31e Mon Sep 17 00:00:00 2001 From: jearhe <0verme7753@gmail.com> Date: Wed, 26 Aug 2026 13:51:26 +0800 Subject: [PATCH] perf(kafka): batch consumer group subscriptions into one agent round trip --- .../java/com/dbx/agent/kafka/KafkaAgent.java | 172 ++++++++++++++- .../com/dbx/agent/kafka/KafkaAgentTest.java | 59 ++++++ crates/dbx-core/src/mq/adapters/kafka.rs | 196 ++++++++++++------ 3 files changed, 364 insertions(+), 63 deletions(-) diff --git a/agents/drivers/kafka/src/main/java/com/dbx/agent/kafka/KafkaAgent.java b/agents/drivers/kafka/src/main/java/com/dbx/agent/kafka/KafkaAgent.java index 5a4f602485..d1cc17090e 100644 --- a/agents/drivers/kafka/src/main/java/com/dbx/agent/kafka/KafkaAgent.java +++ b/agents/drivers/kafka/src/main/java/com/dbx/agent/kafka/KafkaAgent.java @@ -963,22 +963,184 @@ public final class KafkaAgent { private static Object listConsumerGroups(JsonObject params) throws Exception { AdminClient admin = requireAdmin(); int timeout = requestTimeout(params); + String filterTopic = stringOrEmpty(params, "topic"); + Collection groups = admin.listConsumerGroups( new ListConsumerGroupsOptions().timeoutMs(timeout)) .all().get(timeout, TimeUnit.MILLISECONDS); + List groupIds = groups.stream() + .map(ConsumerGroupListing::groupId) + .sorted() + .toList(); + if (groupIds.isEmpty()) { + return Collections.singletonMap("groups", Collections.emptyList()); + } + + // Batch-describe every group in one Kafka Admin request. Per-group + // failures degrade to an empty member list instead of failing the batch. + Map descriptions = new HashMap<>(); + DescribeConsumerGroupsResult described = admin.describeConsumerGroups( + groupIds, + new DescribeConsumerGroupsOptions().timeoutMs(timeout) + ); + for (String groupId : groupIds) { + try { + descriptions.put( + groupId, + described.describedGroups().get(groupId).get(timeout, TimeUnit.MILLISECONDS) + ); + } catch (Exception error) { + logger().debug("Consumer group details unavailable for {}: {}", groupId, normalizeErrorMessage(error)); + } + } + + // Batch-resolve committed offsets for every group in one Admin request. + Map offsetSpecs = new LinkedHashMap<>(); + for (String groupId : groupIds) { + offsetSpecs.put(groupId, new ListConsumerGroupOffsetsSpec()); + } + Map> committedOffsets = new HashMap<>(); + ListConsumerGroupOffsetsResult offsetsResult = admin.listConsumerGroupOffsets( + offsetSpecs, + new ListConsumerGroupOffsetsOptions().timeoutMs(timeout) + ); + for (String groupId : groupIds) { + try { + Map rows = offsetsResult + .partitionsToOffsetAndMetadata(groupId) + .get(timeout, TimeUnit.MILLISECONDS); + committedOffsets.put(groupId, rows == null ? Collections.emptyMap() : rows); + } catch (Exception error) { + committedOffsets.put(groupId, Collections.emptyMap()); + logger().debug("Committed offsets unavailable for {}: {}", groupId, normalizeErrorMessage(error)); + } + } + + // End offsets: one batched Admin request limited to the requested + // topic's partitions (deduplicated across groups). Without a topic + // filter no end offsets are resolved, keeping the listing cheap. + Map endOffsets = new HashMap<>(); + if (!filterTopic.isEmpty()) { + LinkedHashSet topicPartitions = new LinkedHashSet<>(); + for (Map rows : committedOffsets.values()) { + for (TopicPartition topicPartition : rows.keySet()) { + if (topicPartition.topic().equals(filterTopic)) { + topicPartitions.add(topicPartition); + } + } + } + if (!topicPartitions.isEmpty()) { + Map endOffsetSpecs = new LinkedHashMap<>(); + for (TopicPartition topicPartition : topicPartitions) { + endOffsetSpecs.put(topicPartition, OffsetSpec.latest()); + } + try { + ListOffsetsResult latestOffsets = admin.listOffsets( + endOffsetSpecs, + new ListOffsetsOptions().timeoutMs(timeout) + ); + for (TopicPartition topicPartition : topicPartitions) { + try { + ListOffsetsResult.ListOffsetsResultInfo info = latestOffsets + .partitionResult(topicPartition) + .get(timeout, TimeUnit.MILLISECONDS); + endOffsets.put(topicPartition, info.offset()); + } catch (Exception error) { + // Leave the partition absent so callers can tell an + // unknown end offset (no lag contribution) apart from zero lag. + logger().debug("End offset unavailable for {}: {}", topicPartition, normalizeErrorMessage(error)); + } + } + } catch (Exception error) { + logger().debug("End offset batch unavailable: {}", normalizeErrorMessage(error)); + } + } + } + + Map listingByGroup = groups.stream() + .collect(Collectors.toMap(ConsumerGroupListing::groupId, listing -> listing)); + List> result = new ArrayList<>(); - for (ConsumerGroupListing group : groups) { + for (String groupId : groupIds) { + ConsumerGroupListing listing = listingByGroup.get(groupId); + ConsumerGroupDescription description = descriptions.get(groupId); + Map committed = committedOffsets.getOrDefault(groupId, Collections.emptyMap()); + Map g = new LinkedHashMap<>(); - g.put("groupId", group.groupId()); - g.put("state", group.state().map(Enum::name).orElse("UNKNOWN")); - g.put("simpleGroup", group.isSimpleConsumerGroup()); + g.put("groupId", groupId); + g.put("state", listing == null ? "UNKNOWN" : listing.state().map(Enum::name).orElse("UNKNOWN")); + g.put("simpleGroup", listing != null && listing.isSimpleConsumerGroup()); + g.put("members", memberMaps(description == null ? Collections.emptyList() : description.members())); + g.put("committedOffsets", offsetRows(committed, filterTopic)); + g.put("endOffsets", endOffsetRows(committed, endOffsets, filterTopic)); result.add(g); } - result.sort(Comparator.comparing(m -> (String) m.get("groupId"))); return Collections.singletonMap("groups", result); } + static List> memberMaps(Collection members) { + List> result = new ArrayList<>(); + for (MemberDescription member : members) { + Map m = new LinkedHashMap<>(); + m.put("memberId", member.consumerId()); + m.put("clientId", member.clientId()); + m.put("host", member.host()); + List> assignments = new ArrayList<>(); + for (TopicPartition topicPartition : member.assignment().topicPartitions()) { + Map a = new LinkedHashMap<>(); + a.put("topic", topicPartition.topic()); + a.put("partition", topicPartition.partition()); + assignments.add(a); + } + m.put("assignments", assignments); + result.add(m); + } + return result; + } + + /** Committed offsets as rows, optionally restricted to one topic. */ + static List> offsetRows(Map offsets, String topic) { + List> rows = new ArrayList<>(); + for (Map.Entry entry : offsets.entrySet()) { + if (!topic.isEmpty() && !entry.getKey().topic().equals(topic)) { + continue; + } + Map row = new LinkedHashMap<>(); + row.put("topic", entry.getKey().topic()); + row.put("partition", entry.getKey().partition()); + row.put("offset", entry.getValue().offset()); + rows.add(row); + } + rows.sort(Comparator.comparingInt(row -> (int) row.get("partition"))); + return rows; + } + + /** End offsets for a group's committed partitions, restricted to one topic. */ + static List> endOffsetRows( + Map committed, + Map endOffsets, + String topic + ) { + List> rows = new ArrayList<>(); + for (TopicPartition topicPartition : committed.keySet()) { + if (!topic.isEmpty() && !topicPartition.topic().equals(topic)) { + continue; + } + Long endOffset = endOffsets.get(topicPartition); + if (endOffset == null) { + continue; + } + Map row = new LinkedHashMap<>(); + row.put("topic", topicPartition.topic()); + row.put("partition", topicPartition.partition()); + row.put("offset", endOffset); + rows.add(row); + } + rows.sort(Comparator.comparingInt(row -> (int) row.get("partition"))); + return rows; + } + private static Object getConsumerGroupSnapshot(JsonObject params) throws Exception { AdminClient admin = requireAdmin(); int timeout = requestTimeout(params); diff --git a/agents/drivers/kafka/src/test/java/com/dbx/agent/kafka/KafkaAgentTest.java b/agents/drivers/kafka/src/test/java/com/dbx/agent/kafka/KafkaAgentTest.java index daf6b0edc6..43aba0a5fa 100644 --- a/agents/drivers/kafka/src/test/java/com/dbx/agent/kafka/KafkaAgentTest.java +++ b/agents/drivers/kafka/src/test/java/com/dbx/agent/kafka/KafkaAgentTest.java @@ -142,6 +142,65 @@ class KafkaAgentTest { assertEquals("End offsets unavailable for 1 partition(s)", group.get("error")); } + @Test + void consumerGroupListRowsFilterToTheRequestedTopicAndSortByPartition() { + Map offsets = new HashMap<>(); + offsets.put(new TopicPartition("orders", 2), new OffsetAndMetadata(20)); + offsets.put(new TopicPartition("payments", 0), new OffsetAndMetadata(5)); + offsets.put(new TopicPartition("orders", 0), new OffsetAndMetadata(8)); + + List> rows = KafkaAgent.offsetRows(offsets, "orders"); + + assertEquals(List.of("orders:0", "orders:2"), rows.stream() + .map(row -> row.get("topic") + ":" + row.get("partition")) + .toList()); + assertEquals(8L, rows.get(0).get("offset")); + assertEquals(20L, rows.get(1).get("offset")); + } + + @Test + void consumerGroupListRowsKeepEveryPartitionWithoutATopicFilter() { + Map offsets = new HashMap<>(); + offsets.put(new TopicPartition("orders", 0), new OffsetAndMetadata(8)); + offsets.put(new TopicPartition("payments", 1), new OffsetAndMetadata(3)); + + List> rows = KafkaAgent.offsetRows(offsets, ""); + + assertEquals(2, rows.size()); + } + + @Test + void consumerGroupListEndOffsetRowsSkipPartitionsWithoutResolvedEndOffsets() { + Map committed = new HashMap<>(); + committed.put(new TopicPartition("orders", 0), new OffsetAndMetadata(8)); + committed.put(new TopicPartition("orders", 1), new OffsetAndMetadata(20)); + Map endOffsets = new HashMap<>(); + endOffsets.put(new TopicPartition("orders", 1), 24L); + + List> rows = KafkaAgent.endOffsetRows(committed, endOffsets, "orders"); + + assertEquals(List.of("orders:1"), rows.stream() + .map(row -> row.get("topic") + ":" + row.get("partition")) + .toList()); + assertEquals(24L, rows.get(0).get("offset")); + } + + @Test + void consumerGroupListEndOffsetRowsFilterToTheRequestedTopic() { + Map committed = new HashMap<>(); + committed.put(new TopicPartition("orders", 0), new OffsetAndMetadata(8)); + committed.put(new TopicPartition("payments", 0), new OffsetAndMetadata(4)); + Map endOffsets = new HashMap<>(); + endOffsets.put(new TopicPartition("orders", 0), 10L); + endOffsets.put(new TopicPartition("payments", 0), 9L); + + List> rows = KafkaAgent.endOffsetRows(committed, endOffsets, "orders"); + + assertEquals(List.of("orders:0"), rows.stream() + .map(row -> row.get("topic") + ":" + row.get("partition")) + .toList()); + } + @Test void explicitConsumerGroupOffsetsAcceptOneBoundedPartitionOffset() { JsonObject params = JsonParser.parseString(""" diff --git a/crates/dbx-core/src/mq/adapters/kafka.rs b/crates/dbx-core/src/mq/adapters/kafka.rs index 00936e550e..6a386aecbf 100644 --- a/crates/dbx-core/src/mq/adapters/kafka.rs +++ b/crates/dbx-core/src/mq/adapters/kafka.rs @@ -269,32 +269,20 @@ impl MessageQueueAdmin for KafkaAdmin { // ---- Subscriptions (mapped to consumer groups) ---- async fn list_subscriptions(&self, topic: &TopicRef) -> Result, String> { - // List consumer groups and filter those subscribed to this topic - let result: serde_json::Value = self.call("mq_list_consumer_groups", serde_json::json!({})).await?; + // One batched agent RPC: the agent lists every consumer group, batch-describes + // them, and resolves committed + end offsets for the requested topic with a + // handful of Kafka Admin calls. Subscriptions are filtered in memory below. + // (The previous 1 + 2N serial RPC pattern — describe + lag per group — took + // tens of seconds on remote clusters with many groups, see #7163.) + let result: serde_json::Value = + self.call("mq_list_consumer_groups", serde_json::json!({ "topic": topic.topic })).await?; let groups = result.get("groups").and_then(|v| v.as_array()).cloned().unwrap_or_default(); // For each group, check both active assignments and committed offsets. let mut subs = Vec::new(); for group in groups { let group_id = group.get("groupId").and_then(|v| v.as_str()).unwrap_or_default(); - let desc = match self - .call::("mq_describe_consumer_group", serde_json::json!({ "groupId": group_id })) - .await - { - Ok(desc) => desc, - Err(_) => continue, // Skip groups we can't describe - }; - let lag = self - .call::( - "mq_get_consumer_lag", - serde_json::json!({ - "groupId": group_id, - "topic": topic.topic, - }), - ) - .await - .ok(); - if let Some(sub) = kafka_subscription_for_topic(group_id, &topic.topic, &desc, lag.as_ref()) { + if let Some(sub) = kafka_subscription_from_group_row(group_id, &topic.topic, &group) { subs.push(sub); } } @@ -763,13 +751,12 @@ fn reset_cursor_params(topic: &TopicRef, sub: &str, pos: ResetPosition) -> Resul } } -fn kafka_subscription_for_topic( - group_id: &str, - topic: &str, - desc: &serde_json::Value, - lag: Option<&serde_json::Value>, -) -> Option { - let consumers = desc +/// Build a `SubscriptionInfo` for `topic` from one row of the batched +/// `mq_list_consumer_groups` response. The agent already resolved committed +/// offsets (and end offsets for the requested topic) for every group, so no +/// per-group RPC happens here. +fn kafka_subscription_from_group_row(group_id: &str, topic: &str, row: &serde_json::Value) -> Option { + let consumers = row .get("members") .and_then(|v| v.as_array()) .into_iter() @@ -778,20 +765,58 @@ fn kafka_subscription_for_topic( .map(kafka_consumer_from_member) .collect::>(); let has_active_assignment = !consumers.is_empty(); - let has_committed_offsets = lag - .and_then(|v| v.get("partitions")) + + let committed = row + .get("committedOffsets") .and_then(|v| v.as_array()) - .map(|partitions| !partitions.is_empty()) - .unwrap_or(false); + .into_iter() + .flatten() + .filter(|partition| partition.get("topic").and_then(|v| v.as_str()) == Some(topic)) + .collect::>(); + let has_committed_offsets = !committed.is_empty(); if !has_active_assignment && !has_committed_offsets { return None; } + let end_offsets = row + .get("endOffsets") + .and_then(|v| v.as_array()) + .map(|partitions| { + partitions + .iter() + .filter_map(|partition| { + // Defensive: the agent already filters end offsets to the + // requested topic, but keying by partition id alone would + // let a same-numbered partition of another topic collide. + if partition.get("topic").and_then(|v| v.as_str()) != Some(topic) { + return None; + } + let partition_id = partition.get("partition").and_then(|v| v.as_i64())?; + let offset = partition.get("offset").and_then(|v| v.as_i64())?; + Some((partition_id, offset)) + }) + .collect::>() + }) + .unwrap_or_default(); + + // Sum lag over the topic's committed partitions. A partition without a + // resolved end offset contributes nothing (matches the agent's lag probe: + // unknown end offsets are not reported as zero lag). + let msg_backlog = committed + .iter() + .filter_map(|partition| { + let partition_id = partition.get("partition").and_then(|v| v.as_i64())?; + let committed_offset = partition.get("offset").and_then(|v| v.as_i64())?; + let end_offset = end_offsets.get(&partition_id).copied()?; + Some((end_offset - committed_offset).max(0)) + }) + .sum(); + Some(SubscriptionInfo { name: group_id.to_string(), sub_type: "consumer-group".to_string(), - msg_backlog: lag.and_then(|v| v.get("totalLag")).and_then(|v| v.as_i64()).unwrap_or(0), + msg_backlog, msg_rate_out: 0.0, msg_throughput_out: 0.0, consumers, @@ -1163,31 +1188,37 @@ mod tests { } #[test] - fn kafka_subscription_for_topic_includes_offline_group_with_committed_offsets() { - let desc = serde_json::json!({ + fn kafka_subscription_from_group_row_includes_offline_group_with_committed_offsets() { + let row = serde_json::json!({ "groupId": "orders-service", - "members": [] - }); - let lag = serde_json::json!({ - "totalLag": 7, - "partitions": [ - { "partition": 0, "currentOffset": 3, "endOffset": 10, "lag": 7 } + "state": "EMPTY", + "simpleGroup": false, + "members": [], + "committedOffsets": [ + { "topic": "orders", "partition": 0, "offset": 3 }, + { "topic": "orders", "partition": 1, "offset": 8 } + ], + "endOffsets": [ + { "topic": "orders", "partition": 0, "offset": 10 }, + { "topic": "orders", "partition": 1, "offset": 12 } ] }); - let sub = kafka_subscription_for_topic("orders-service", "orders", &desc, Some(&lag)) + let sub = kafka_subscription_from_group_row("orders-service", "orders", &row) .expect("committed offsets should make an inactive group visible"); assert_eq!(sub.name, "orders-service"); assert_eq!(sub.sub_type, "consumer-group"); - assert_eq!(sub.msg_backlog, 7); + assert_eq!(sub.msg_backlog, 11); // (10-3) + (12-8) assert!(sub.consumers.is_empty()); } #[test] - fn kafka_subscription_for_topic_includes_active_assignment_without_committed_offsets() { - let desc = serde_json::json!({ + fn kafka_subscription_from_group_row_includes_active_assignment_without_committed_offsets() { + let row = serde_json::json!({ "groupId": "live-service", + "state": "STABLE", + "simpleGroup": false, "members": [ { "memberId": "consumer-events", @@ -1203,14 +1234,12 @@ mod tests { { "topic": "audit", "partition": 0 } ] } - ] - }); - let lag = serde_json::json!({ - "totalLag": 0, - "partitions": [] + ], + "committedOffsets": [], + "endOffsets": [] }); - let sub = kafka_subscription_for_topic("live-service", "events", &desc, Some(&lag)) + let sub = kafka_subscription_from_group_row("live-service", "events", &row) .expect("active assignments should make the group visible"); assert_eq!(sub.name, "live-service"); @@ -1221,21 +1250,72 @@ mod tests { } #[test] - fn kafka_subscription_for_topic_ignores_unrelated_group() { - let desc = serde_json::json!({ + fn kafka_subscription_from_group_row_ignores_unrelated_group() { + let row = serde_json::json!({ "groupId": "billing-service", + "state": "EMPTY", + "simpleGroup": false, "members": [{ "assignments": [ { "topic": "billing", "partition": 0 } ] - }] - }); - let lag = serde_json::json!({ - "totalLag": 0, - "partitions": [] + }], + "committedOffsets": [ + { "topic": "billing", "partition": 0, "offset": 2 } + ], + "endOffsets": [ + { "topic": "billing", "partition": 0, "offset": 5 } + ] }); - assert!(kafka_subscription_for_topic("billing-service", "orders", &desc, Some(&lag)).is_none()); + assert!(kafka_subscription_from_group_row("billing-service", "orders", &row).is_none()); + } + + #[test] + fn kafka_subscription_from_group_row_skips_partitions_without_end_offsets() { + let row = serde_json::json!({ + "groupId": "orders-service", + "state": "EMPTY", + "simpleGroup": false, + "members": [], + "committedOffsets": [ + { "topic": "orders", "partition": 0, "offset": 3 }, + { "topic": "orders", "partition": 1, "offset": 8 } + ], + "endOffsets": [ + // partition 1's end offset is unavailable (agent-side failure) + { "topic": "orders", "partition": 0, "offset": 10 } + ] + }); + + let sub = kafka_subscription_from_group_row("orders-service", "orders", &row) + .expect("the group still has visible committed offsets"); + + assert_eq!(sub.msg_backlog, 7); // only partition 0 contributes + } + + #[test] + fn kafka_subscription_from_group_row_filters_committed_offsets_to_the_requested_topic() { + let row = serde_json::json!({ + "groupId": "multi-topic-service", + "state": "EMPTY", + "simpleGroup": false, + "members": [], + "committedOffsets": [ + { "topic": "orders", "partition": 0, "offset": 3 }, + { "topic": "payments", "partition": 0, "offset": 7 } + ], + "endOffsets": [ + { "topic": "orders", "partition": 0, "offset": 8 }, + { "topic": "payments", "partition": 0, "offset": 9 } + ] + }); + + let sub = kafka_subscription_from_group_row("multi-topic-service", "orders", &row) + .expect("committed offsets on the requested topic make the group visible"); + + // Lag totals only the requested topic's partitions. + assert_eq!(sub.msg_backlog, 5); } #[test]