BE: Metrics: Fix JMX metrics collection (2) (#1285)

Co-authored-by: iliax <ilya.kuramshin@almatech.dev>

Author: iliax

api/src/main/java/io/kafbat/ui/service/metrics/RawMetric.java

@@ -6,8 +6,10 @@
 import java.math.BigDecimal;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Stream;
 
 public interface RawMetric {
@@ -26,24 +28,27 @@ static RawMetric create(String name, Map<String, String> labels, BigDecimal valu
 
   static Stream<MetricSnapshot> groupIntoSnapshot(Collection<RawMetric> rawMetrics) {
     Map<String, Gauge> map = new LinkedHashMap<>();
+    Map<String, String[]> gaugeLabels = new HashMap<>();
     for (RawMetric m : rawMetrics) {
       var lbls = m.labels().keySet()
           .stream()
           .map(PrometheusNaming::sanitizeLabelName)
           .toArray(String[]::new);
-      var lblVals = Arrays.stream(lbls)
+      var lblVals = m.labels().keySet()
+          .stream()
           .map(l -> m.labels().get(l))
           .toArray(String[]::new);
       var sanitizedName = PrometheusNaming.sanitizeMetricName(m.name());
       var gauge = map.computeIfAbsent(
-          sanitizedName,
-          n -> Gauge.builder()
-              .name(sanitizedName)
-              .help(sanitizedName)
-              .labelNames(lbls)
-              .build()
+          sanitizedName, n -> {
+            gaugeLabels.put(n, lbls);
+            return Gauge.builder().name(n).help(n).labelNames(lbls).build();
+          }
       );
-      gauge.labelValues(lblVals).set(m.value().doubleValue());
+      if (Arrays.equals(lbls, gaugeLabels.get(sanitizedName))) {
+        //using labels of first registered gauge, if not fit - skipping
+        gauge.labelValues(lblVals).set(m.value().doubleValue());
+      }
     }
     return map.values().stream().map(Gauge::collect);
   }

api/src/test/java/io/kafbat/ui/service/metrics/RawMetricTest.java

@@ -0,0 +1,35 @@
+package io.kafbat.ui.service.metrics;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.prometheus.metrics.model.snapshots.DataPointSnapshot;
+import io.prometheus.metrics.model.snapshots.Labels;
+import io.prometheus.metrics.model.snapshots.MetricSnapshot;
+import java.math.BigDecimal;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+class RawMetricTest {
+
+  @Test
+  void groupIntoSnapshotChoosesFirstGaugeLabels() {
+    List<MetricSnapshot> list = RawMetric.groupIntoSnapshot(
+        List.of(
+            RawMetric.create("name", Map.of("l1", "v1"), BigDecimal.ONE),
+            RawMetric.create("name", Map.of("l1", "v11"), BigDecimal.TWO),
+            RawMetric.create("name", Map.of("l1", "v1", "l2", "v2"), BigDecimal.TEN)
+        )
+    ).toList();
+
+    assertThat(list)
+        .hasSize(1)
+        .element(0)
+        .satisfies(snap ->
+            assertThat(snap.getDataPoints())
+                .map(DataPointSnapshot::getLabels)
+                .containsExactly(Labels.of("l1", "v1"), Labels.of("l1", "v11"))
+        );
+  }
+
+}