Improve implementation of Kafka metrics type mapping (#3513)

* Improve implementation of Kafka metrics type mapping

Closes gh-3499

* Assert that type-based improved counter detection produces the same set with ".total"-ending detection

Author: izeye

micrometer-core/src/main/java/io/micrometer/core/instrument/binder/kafka/KafkaMetrics.java

@@ -27,6 +27,8 @@
 import io.micrometer.core.util.internal.logging.WarnThenDebugLogger;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.metrics.KafkaMetric;
+import org.apache.kafka.common.metrics.Measurable;
 
 import java.time.Duration;
 import java.util.*;
@@ -69,6 +71,23 @@ class KafkaMetrics implements MeterBinder, AutoCloseable {
     static final String KAFKA_VERSION_TAG_NAME = "kafka.version";
     static final String DEFAULT_VALUE = "unknown";
 
+    private static final Set<Class<?>> counterMeasurableClasses = new HashSet<>();
+
+    static {
+        Set<String> classNames = new HashSet<>();
+        classNames.add("org.apache.kafka.common.metrics.stats.CumulativeSum");
+        classNames.add("org.apache.kafka.common.metrics.stats.CumulativeCount");
+
+        for (String className : classNames) {
+            try {
+                counterMeasurableClasses.add(Class.forName(className));
+            }
+            catch (ClassNotFoundException e) {
+                // Class doesn't exist in this version of kafka client - skip
+            }
+        }
+    }
+
     private final Supplier<Map<MetricName, ? extends Metric>> metricsSupplier;
 
     private final AtomicReference<Map<MetricName, ? extends Metric>> metrics = new AtomicReference<>();
@@ -135,7 +154,7 @@ void prepareToBindMetrics(MeterRegistry registry) {
         this.metrics.set(this.metricsSupplier.get());
         Map<MetricName, ? extends Metric> metrics = this.metrics.get();
         // Collect static metrics and tags
-        MetricName startTime = null;
+        Metric startTimeMetric = null;
 
         for (Map.Entry<MetricName, ? extends Metric> entry : metrics.entrySet()) {
             MetricName name = entry.getKey();
@@ -144,12 +163,13 @@ void prepareToBindMetrics(MeterRegistry registry) {
                     kafkaVersion = (String) entry.getValue().metricValue();
                 }
                 else if (START_TIME_METRIC_NAME.equals(name.name())) {
-                    startTime = entry.getKey();
+                    startTimeMetric = entry.getValue();
                 }
         }
 
-        if (startTime != null) {
-            bindMeter(registry, startTime, meterName(startTime), meterTags(startTime));
+        if (startTimeMetric != null) {
+            MetricName startTimeMetricName = startTimeMetric.metricName();
+            bindMeter(registry, startTimeMetric, meterName(startTimeMetricName), meterTags(startTimeMetricName));
         }
     }
 
@@ -220,7 +240,7 @@ else if (tags.size() == meterTagsWithCommonTags.size())
 
                     List<Tag> tags = meterTags(name);
                     try {
-                        Meter meter = bindMeter(registry, metric.metricName(), meterName, tags);
+                        Meter meter = bindMeter(registry, metric, meterName, tags);
                         List<Meter> meters = registryMetersByNames.computeIfAbsent(meterName, k -> new ArrayList<>());
                         meters.add(meter);
                     }
@@ -242,18 +262,34 @@ else if (tags.size() == meterTagsWithCommonTags.size())
         }
     }
 
-    private Meter bindMeter(MeterRegistry registry, MetricName metricName, String meterName, Iterable<Tag> tags) {
-        Meter meter = registerMeter(registry, metricName, meterName, tags);
+    private Meter bindMeter(MeterRegistry registry, Metric metric, String meterName, Iterable<Tag> tags) {
+        Meter meter = registerMeter(registry, metric, meterName, tags);
         registeredMeterIds.add(meter.getId());
         return meter;
     }
 
-    private Meter registerMeter(MeterRegistry registry, MetricName metricName, String meterName, Iterable<Tag> tags) {
-        if (meterName.endsWith("total")) {
+    private Meter registerMeter(MeterRegistry registry, Metric metric, String meterName, Iterable<Tag> tags) {
+        MetricName metricName = metric.metricName();
+        Class<? extends Measurable> measurableClass = getMeasurableClass(metric);
+        if ((measurableClass == null && meterName.endsWith("total"))
+                || (measurableClass != null && counterMeasurableClasses.contains(measurableClass))) {
             return registerCounter(registry, metricName, meterName, tags);
         }
-        else {
-            return registerGauge(registry, metricName, meterName, tags);
+
+        return registerGauge(registry, metricName, meterName, tags);
+    }
+
+    @Nullable
+    private static Class<? extends Measurable> getMeasurableClass(Metric metric) {
+        if (!(metric instanceof KafkaMetric)) {
+            return null;
+        }
+
+        try {
+            return ((KafkaMetric) metric).measurable().getClass();
+        }
+        catch (IllegalStateException ex) {
+            return null;
         }
     }
 

micrometer-core/src/test/java/io/micrometer/core/instrument/binder/kafka/KafkaClientMetricsIntegrationTest.java

@@ -15,7 +15,9 @@
  */
 package io.micrometer.core.instrument.binder.kafka;
 
+import io.micrometer.core.instrument.FunctionCounter;
 import io.micrometer.core.instrument.Gauge;
+import io.micrometer.core.instrument.Meter;
 import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -35,7 +37,9 @@
 
 import java.time.Duration;
 import java.util.Collections;
+import java.util.List;
 import java.util.Properties;
+import java.util.stream.Collectors;
 
 import static java.lang.System.out;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -117,6 +121,12 @@ void shouldManageProducerAndConsumerMetrics() {
         out.println("All meters from producer and consumer:");
         printMeters(registry);
 
+        List<Meter> metersEndingWithTotal = registry.getMeters().stream()
+                .filter(meter -> meter.getId().getName().endsWith(".total")).collect(Collectors.toList());
+        List<Meter> functionCounters = registry.getMeters().stream().filter(meter -> meter instanceof FunctionCounter)
+                .collect(Collectors.toList());
+        assertThat(metersEndingWithTotal).isEqualTo(functionCounters);
+
         producerKafkaMetrics.close();
         consumerKafkaMetrics.close();
     }