feat: add an option to describe lag (#8)

Author: BewareMyPower

README.md

@@ -142,7 +142,19 @@ There are 2 members:
 | 1 | consumer-sub-2 | consumer-sub-2-63b7c688-3007-4650-91eb-404284dfd837 | /127.0.0.1:54213 | [test-2, test-3] |
 ```
 
-The format of each line is `[index] <group-id> <state>`.
+Adding the `--lag` option can describe the lag info for all subscribed topic-partitions:
+
+```bash
+$ time ./build/snctl-cpp groups describe sub --lag
+Group ID: sub
+...
+Offsets info for group 'sub' with 4 topic-partitions:
+| topic-partition | committed offset | end offset | lag |
+| test-0 | 0 | 0 | 0 |
+| test-1 | 1 | 1 | 0 |
+| test-2 | 0 | 0 | 0 |
+| test-3 | 0 | 0 | 0 |
+```
 
 ## Logging
 

include/snctl-cpp/groups.h

@@ -30,6 +30,10 @@ class Groups : public SubCommand {
         .add_argument("group")
         .help("The group id")
         .required();
+    describe_command_.add_argument("--lag")
+        .default_value(false)
+        .implicit_value(true)
+        .help("Show the lag of the group");
 
     add_child(list_command_);
     add_child(describe_command_);
@@ -42,7 +46,8 @@ class Groups : public SubCommand {
       list_groups(rk, rkqu);
     } else if (is_subcommand_used(describe_command_)) {
       auto group = describe_command_.get("group");
-      describe_group(rk, rkqu, group);
+      auto show_lag = describe_command_.get<bool>("lag");
+      describe_group(rk, rkqu, group, show_lag);
     } else {
       fail();
     }

include/snctl-cpp/groups/describe_group.h

@@ -15,12 +15,17 @@
  */
 #pragma once
 
+#include "snctl-cpp/raii_helper.h"
 #include "snctl-cpp/rk_event_wrapper.h"
 #include <cassert>
 #include <cstddef>
+#include <cstdint>
 #include <iostream>
 #include <librdkafka/rdkafka.h>
+#include <map>
 #include <ostream>
+#include <sstream>
+#include <stdexcept>
 #include <string>
 
 static std::ostream &operator<<(std::ostream &os, const rd_kafka_Node_t *node) {
@@ -37,8 +42,96 @@ static std::ostream &operator<<(std::ostream &os,
   return os;
 }
 
+static std::ostream &operator<<(std::ostream &os,
+                                const rd_kafka_topic_partition_t *partition) {
+  return os << *partition;
+}
+
+// Return a map, whose key's format is "<topic>-<partition>" and  value is the
+// committed offset of the topic-partition.
+inline auto
+query_committed_offsets(rd_kafka_t *rk, rd_kafka_queue_t *rkqu,
+                        const std::string &expected_group,
+                        rd_kafka_topic_partition_list_t *rk_topic_partitions) {
+  auto *group_offset = rd_kafka_ListConsumerGroupOffsets_new(
+      expected_group.c_str(), rk_topic_partitions);
+  GUARD(group_offset, rd_kafka_ListConsumerGroupOffsets_destroy);
+
+  rd_kafka_ListConsumerGroupOffsets_t *group_offsets[1];
+  group_offsets[0] = group_offset;
+
+  rd_kafka_ListConsumerGroupOffsets(rk, &group_offsets[0], 1, nullptr, rkqu);
+  auto event = RdKafkaEvent::poll(rkqu);
+  const auto *result =
+      rd_kafka_event_ListConsumerGroupOffsets_result(event.handle());
+  assert(result != nullptr);
+
+  size_t group_count;
+  auto *group_result =
+      rd_kafka_ListConsumerGroupOffsets_result_groups(result, &group_count);
+  if (group_count != 1) {
+    std::ostringstream oss;
+    oss << "Expected exactly one group, but got " << group_count
+        << " in ListConsumerGroupOffsets";
+    throw std::runtime_error(oss.str());
+  }
+
+  if (expected_group !=
+      std::string(rd_kafka_group_result_name(group_result[0]))) {
+    std::ostringstream oss;
+    oss << "Expected group '" << expected_group << "', but got '"
+        << rd_kafka_group_result_name(group_result[0])
+        << "' in ListConsumerGroupOffsets";
+    throw std::runtime_error(oss.str());
+  }
+
+  std::map<std::string, int64_t> offsets;
+  const auto *partitions = rd_kafka_group_result_partitions(group_result[0]);
+  if (partitions == nullptr) {
+    return offsets;
+  }
+  for (size_t i = 0; i < partitions->cnt; i++) {
+    const auto &partition = partitions->elems[i];
+    if (partition.offset == RD_KAFKA_OFFSET_INVALID) {
+      continue; // Skip invalid offsets
+    }
+    std::ostringstream oss;
+    oss << partition;
+    offsets[oss.str()] = partition.offset;
+  }
+  return offsets;
+}
+
+// Return a map whose key is "<topic>-<partition>" and value is the end offset
+inline auto
+query_end_offsets(rd_kafka_t *rk, rd_kafka_queue_t *rkqu,
+                  rd_kafka_topic_partition_list_t *rk_topic_partitions) {
+  for (int i = 0; i < rk_topic_partitions->cnt; i++) {
+    rk_topic_partitions->elems[i].offset = RD_KAFKA_OFFSET_SPEC_LATEST;
+  }
+
+  rd_kafka_ListOffsets(rk, rk_topic_partitions, nullptr, rkqu);
+  auto event = RdKafkaEvent::poll(rkqu);
+  const auto *result = rd_kafka_event_ListOffsets_result(event.handle());
+  assert(result != nullptr);
+
+  size_t num_partitions;
+  const auto *offsets_result =
+      rd_kafka_ListOffsets_result_infos(result, &num_partitions);
+
+  std::map<std::string, int64_t> offsets;
+  for (size_t i = 0; i < num_partitions; i++) {
+    const auto *rk_topic_partition =
+        rd_kafka_ListOffsetsResultInfo_topic_partition(offsets_result[i]);
+    std::ostringstream oss;
+    oss << rk_topic_partition;
+    offsets[oss.str()] = rk_topic_partition->offset;
+  }
+  return offsets;
+}
+
 inline void describe_group(rd_kafka_t *rk, rd_kafka_queue_t *rkqu,
-                           const std::string &group) {
+                           const std::string &group, bool show_lag) {
   const char *groups[1] = {group.c_str()};
   rd_kafka_DescribeConsumerGroups(rk, groups, 1, nullptr, rkqu);
 
@@ -84,11 +177,16 @@ inline void describe_group(rd_kafka_t *rk, rd_kafka_queue_t *rkqu,
 
     if (member_count > 0) {
       std::cout << "There are " << member_count << " members:" << std::endl;
-      std::cout << "| index | client id | consumer id | host | assignments |"
-                << std::endl;
+      std::cout << "| index | client id | consumer id | host | assignments |";
+      if (show_lag) {
+        std::cout << " lag | end offset |";
+      }
+      std::cout << std::endl;
     } else {
       std::cout << "No members" << std::endl;
     }
+
+    std::map<std::string, int> topic_to_partition_cnt;
     for (size_t i = 0; i < member_count; i++) {
       const auto *member = rd_kafka_ConsumerGroupDescription_member(group, i);
       assert(member != nullptr);
@@ -105,11 +203,53 @@ inline void describe_group(rd_kafka_t *rk, rd_kafka_queue_t *rkqu,
         if (j > 0) {
           std::cout << ", ";
         }
-        std::cout << partitions->elems[j];
+        const auto partition = partitions->elems[j];
+        std::cout << partition;
+        if (auto it = topic_to_partition_cnt.find(partition.topic);
+            it != topic_to_partition_cnt.end()) {
+          if (it->second <= partition.partition) {
+            it->second = partition.partition + 1;
+          }
+        } else {
+          topic_to_partition_cnt[partition.topic] = partition.partition + 1;
+        }
       }
       std::cout << "] |" << std::endl;
     }
 
+    if (show_lag) {
+      auto *rk_topic_partitions =
+          rd_kafka_topic_partition_list_new(topic_to_partition_cnt.size());
+      GUARD(rk_topic_partitions, rd_kafka_topic_partition_list_destroy);
+
+      for (auto &&[topic, partition_cnt] : topic_to_partition_cnt) {
+        for (int i = 0; i < partition_cnt; i++) {
+          rd_kafka_topic_partition_list_add(rk_topic_partitions, topic.c_str(),
+                                            i);
+        }
+      }
+
+      const auto committed_offsets =
+          query_committed_offsets(rk, rkqu, group_id, rk_topic_partitions);
+      const auto end_offsets = query_end_offsets(rk, rkqu, rk_topic_partitions);
+      std::cout << "Offsets info for group '" << group_id << "' with "
+                << committed_offsets.size()
+                << " topic-partitions:" << std::endl;
+      std::cout << "| topic-partition | committed offset | end offset | lag |"
+                << std::endl;
+      for (auto &&[topic_partition, committed_offset] : committed_offsets) {
+        auto end_offset_it = end_offsets.find(topic_partition);
+        if (end_offset_it == end_offsets.cend()) {
+          std::cout << "| " << topic_partition << " | " << committed_offset
+                    << " | N/A | N/A |" << std::endl;
+        } else {
+          auto end_offset = end_offset_it->second;
+          std::cout << "| " << topic_partition << " | " << committed_offset
+                    << " | " << end_offset << " | "
+                    << (end_offset - committed_offset) << " |" << std::endl;
+        }
+      }
+    }
   } catch (const std::runtime_error &e) {
     std::cerr << "Failed to describe group '" << group << "': " << e.what()
               << std::endl;