From e63e550a6b1f78a7299e322751fe008c5d79e209 Mon Sep 17 00:00:00 2001
From: Pascal Schmid <psschmid@gmx.de>
Date: Sat, 2 Mar 2024 21:23:29 +0100
Subject: [PATCH] Get lags per partition

---
 .../consumer_offset_fetcher.ex                | 44 ++++++++++++-------
 lib/kafkaex_lag_exporter/utils.ex             |  9 +++-
 2 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/lib/kafkaex_lag_exporter/consumer_offset_fetcher.ex b/lib/kafkaex_lag_exporter/consumer_offset_fetcher.ex
index ee929c1..10d35c4 100644
--- a/lib/kafkaex_lag_exporter/consumer_offset_fetcher.ex
+++ b/lib/kafkaex_lag_exporter/consumer_offset_fetcher.ex
@@ -30,11 +30,20 @@ defmodule KafkaexLagExporter.ConsumerOffsetFetcher do
 
     consumer_group_names = get_consumer_group_names(endpoint)
 
-    consumer_lags =
+    topic_names_for_consumer_groups =
       :brod.describe_groups(endpoint, [], consumer_group_names)
-      |> get_consumer_lags
+      |> get_topic_names_for_consumer_groups
 
-    KafkaexLagExporter.Metrics.kafka_metrics(endpoint, consumer_lags)
+    consumer_lags =
+      topic_names_for_consumer_groups
+      |> Enum.map(fn [consumer_group, topics] ->
+        [consumer_group, get_lag_for_consumer(consumer_group, topics)]
+      end)
+
+    consumer_lag_sum = get_lag_for_consumer_sum(consumer_lags)
+
+    KafkaexLagExporter.Metrics.group_lag_per_partition(endpoint, consumer_lags)
+    KafkaexLagExporter.Metrics.group_sum_lag(endpoint, consumer_lag_sum)
 
     Process.send_after(self(), :tick, @interval)
 
@@ -49,32 +58,33 @@ defmodule KafkaexLagExporter.ConsumerOffsetFetcher do
     |> Enum.map(fn {_, group_name, "consumer"} -> group_name end)
   end
 
-  defp get_consumer_lags({:ok, group_descriptions}) do
+  defp get_topic_names_for_consumer_groups({:ok, group_descriptions}) do
     group_descriptions
-    |> Enum.map(fn %{group_id: consumer_group, members: members} ->
-      [consumer_group, members]
-    end)
+    |> Enum.map(fn %{group_id: consumer_group, members: members} -> [consumer_group, members] end)
     |> Enum.map(fn [consumer_group, members] -> [consumer_group, get_topic_names(members)] end)
-    |> Enum.map(fn [consumer_group, topics] ->
-      [consumer_group, get_lag_for_consumer(consumer_group, topics)]
-
-      # credo:disable-for-next-line
-      # TODO: [consumer_group, topic, get_lag_for_consumer(consumer_group, topic)]
-    end)
   end
 
-  defp get_consumer_lags({_, _}), do: []
-
   defp get_topic_names(members) do
     Enum.flat_map(members, fn member ->
       KafkaexLagExporter.TopicNameParser.parse_topic_names(member.member_assignment)
     end)
   end
 
+  # TODO: test method for multiple topics
   defp get_lag_for_consumer(consumer_group, topics) do
     topics
-    |> Enum.reduce(0, fn topic, acc ->
-      acc + KafkaexLagExporter.KafkaUtils.lag_total(topic, consumer_group, :client1)
+    |> Enum.flat_map(fn topic ->
+      KafkaexLagExporter.KafkaUtils.lag(topic, consumer_group, :client1)
     end)
   end
+
+  # TODO: test method for multiple topics
+  defp get_lag_for_consumer_sum(lags_per_consumer_group) do
+    lags_per_consumer_group
+    |> Enum.map(fn [topic, lag_per_partition] -> [topic, sum_topic_lag(lag_per_partition)] end)
+  end
+
+  defp sum_topic_lag(item, acc \\ 0)
+  defp sum_topic_lag([], acc), do: acc
+  defp sum_topic_lag([h | t], acc), do: sum_topic_lag(t, acc + elem(h, 1))
 end
diff --git a/lib/kafkaex_lag_exporter/utils.ex b/lib/kafkaex_lag_exporter/utils.ex
index f6447c6..acf218c 100644
--- a/lib/kafkaex_lag_exporter/utils.ex
+++ b/lib/kafkaex_lag_exporter/utils.ex
@@ -53,8 +53,13 @@ defmodule KafkaexLagExporter.KafkaUtils do
 
   @spec lag(binary(), binary(), atom()) :: list({non_neg_integer(), integer()})
   def lag(topic, consumer_group, client) do
-    offsets = resolve_offsets(topic, :latest, client)
-    committed_offsets = fetch_committed_offsets(topic, consumer_group, client)
+    offsets =
+      resolve_offsets(topic, :latest, client)
+      |> Enum.sort_by(fn {key, _value} -> key end)
+
+    committed_offsets =
+      fetch_committed_offsets(topic, consumer_group, client)
+      |> Enum.sort_by(fn {key, _value} -> key end)
 
     for {{part, current}, {_part2, committed}} <- Enum.zip(offsets, committed_offsets) do
       {part, current - committed}