Fixes #4976 - added optional custom scheduler to Consumer KafkaClientMetrics

Author: vasiliy-sarzhynskyi

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

@@ -24,6 +24,8 @@
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.common.Metric;
 
+import java.util.concurrent.ScheduledExecutorService;
+
 /**
  * Kafka Client metrics binder. This should be closed on application shutdown to clean up
  * resources.
@@ -60,6 +62,16 @@ public KafkaClientMetrics(Producer<?, ?> kafkaProducer) {
         super(kafkaProducer::metrics);
     }
 
+    /**
+     * Kafka {@link Consumer} metrics binder
+     * @param kafkaConsumer consumer instance to be instrumented
+     * @param tags additional tags
+     * @param scheduler scheduler to check and bind metrics
+     */
+    public KafkaClientMetrics(Consumer<?, ?> kafkaConsumer, Iterable<Tag> tags, ScheduledExecutorService scheduler) {
+        super(kafkaConsumer::metrics, tags, scheduler);
+    }
+
     /**
      * Kafka {@link Consumer} metrics binder
      * @param kafkaConsumer consumer instance to be instrumented

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

@@ -71,6 +71,8 @@ class KafkaMetrics implements MeterBinder, AutoCloseable {
     static final String KAFKA_VERSION_TAG_NAME = "kafka.version";
     static final String DEFAULT_VALUE = "unknown";
 
+    private static final String DEFAULT_SCHEDULER_THREAD_NAME_PREFIX = "micrometer-kafka-metrics";
+
     private static final Set<Class<?>> counterMeasurableClasses = new HashSet<>();
 
     static {
@@ -96,8 +98,9 @@ class KafkaMetrics implements MeterBinder, AutoCloseable {
 
     private final Duration refreshInterval;
 
-    private final ScheduledExecutorService scheduler = Executors
-        .newSingleThreadScheduledExecutor(new NamedThreadFactory("micrometer-kafka-metrics"));
+    private final ScheduledExecutorService scheduler;
+
+    private final boolean schedulerExternallyManaged;
 
     @Nullable
     private Iterable<Tag> commonTags;
@@ -122,11 +125,23 @@ class KafkaMetrics implements MeterBinder, AutoCloseable {
         this(metricsSupplier, extraTags, DEFAULT_REFRESH_INTERVAL);
     }
 
+    KafkaMetrics(Supplier<Map<MetricName, ? extends Metric>> metricsSupplier, Iterable<Tag> extraTags,
+            ScheduledExecutorService scheduler) {
+        this(metricsSupplier, extraTags, DEFAULT_REFRESH_INTERVAL, scheduler, true);
+    }
+
     KafkaMetrics(Supplier<Map<MetricName, ? extends Metric>> metricsSupplier, Iterable<Tag> extraTags,
             Duration refreshInterval) {
+        this(metricsSupplier, extraTags, refreshInterval, createDefaultScheduler(), false);
+    }
+
+    KafkaMetrics(Supplier<Map<MetricName, ? extends Metric>> metricsSupplier, Iterable<Tag> extraTags,
+            Duration refreshInterval, ScheduledExecutorService scheduler, boolean schedulerExternallyManaged) {
         this.metricsSupplier = metricsSupplier;
         this.extraTags = extraTags;
         this.refreshInterval = refreshInterval;
+        this.scheduler = scheduler;
+        this.schedulerExternallyManaged = schedulerExternallyManaged;
     }
 
     @Override
@@ -295,6 +310,10 @@ private static Class<? extends Measurable> getMeasurableClass(Metric metric) {
         }
     }
 
+    private static ScheduledExecutorService createDefaultScheduler() {
+        return Executors.newSingleThreadScheduledExecutor(new NamedThreadFactory(DEFAULT_SCHEDULER_THREAD_NAME_PREFIX));
+    }
+
     private Gauge registerGauge(MeterRegistry registry, MetricName metricName, String meterName, Iterable<Tag> tags) {
         return Gauge.builder(meterName, this.metrics, toMetricValue(metricName))
             .tags(tags)
@@ -344,7 +363,9 @@ private Meter.Id meterIdForComparison(MetricName metricName) {
 
     @Override
     public void close() {
-        this.scheduler.shutdownNow();
+        if (!schedulerExternallyManaged) {
+            this.scheduler.shutdownNow();
+        }
 
         for (Meter.Id id : registeredMeterIds) {
             registry.remove(id);

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

@@ -25,6 +25,8 @@
 import org.junit.jupiter.api.Test;
 
 import java.util.Properties;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
 
 import static io.micrometer.core.instrument.binder.kafka.KafkaClientMetrics.METRIC_NAME_PREFIX;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.*;
@@ -34,7 +36,7 @@ class KafkaClientMetricsConsumerTest {
 
     private static final String BOOTSTRAP_SERVERS = "localhost:9092";
 
-    private Tags tags = Tags.of("app", "myapp", "version", "1");
+    private final Tags tags = Tags.of("app", "myapp", "version", "1");
 
     KafkaClientMetrics metrics;
 
@@ -71,6 +73,27 @@ void shouldCreateMetersWithTags() {
         }
     }
 
+    @Test
+    void shouldCreateMetersWithTagsAndCustomScheduler() {
+        try (Consumer<String, String> consumer = createConsumer()) {
+            ScheduledExecutorService customScheduler = Executors.newScheduledThreadPool(2);
+            metrics = new KafkaClientMetrics(consumer, tags, customScheduler);
+            MeterRegistry registry = new SimpleMeterRegistry();
+
+            metrics.bindTo(registry);
+
+            assertThat(registry.getMeters()).hasSizeGreaterThan(0)
+                .extracting(meter -> meter.getId().getTag("app"))
+                .allMatch(s -> s.equals("myapp"));
+
+            metrics.close();
+            assertThat(customScheduler.isShutdown()).isFalse();
+
+            customScheduler.shutdownNow();
+            assertThat(customScheduler.isShutdown()).isTrue();
+        }
+    }
+
     private Consumer<String, String> createConsumer() {
         Properties consumerConfig = new Properties();
         consumerConfig.put(BOOTSTRAP_SERVERS_CONFIG, BOOTSTRAP_SERVERS);

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

@@ -34,10 +34,13 @@
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
 
 class KafkaMetricsTest {
 
@@ -68,7 +71,7 @@ void shouldKeepMetersWhenMetricsDoNotChange() {
     }
 
     @Test
-    void closeShouldRemoveAllMeters() {
+    void closeShouldRemoveAllMetersAndShutdownDefaultScheduler() {
         // Given
         Supplier<Map<MetricName, ? extends Metric>> supplier = () -> {
             MetricName metricName = new MetricName("a", "b", "c", new LinkedHashMap<>());
@@ -80,9 +83,35 @@ void closeShouldRemoveAllMeters() {
 
         kafkaMetrics.bindTo(registry);
         assertThat(registry.getMeters()).hasSize(1);
+        assertThat(isDefaultMetricsSchedulerThreadAlive()).isTrue();
 
         kafkaMetrics.close();
         assertThat(registry.getMeters()).isEmpty();
+        await().until(() -> !isDefaultMetricsSchedulerThreadAlive());
+    }
+
+    @Test
+    void closeShouldRemoveAllMetersAndNotShutdownCustomScheduler() {
+        // Given
+        Supplier<Map<MetricName, ? extends Metric>> supplier = () -> {
+            MetricName metricName = new MetricName("a", "b", "c", new LinkedHashMap<>());
+            KafkaMetric metric = new KafkaMetric(this, metricName, new Value(), new MetricConfig(), Time.SYSTEM);
+            return Collections.singletonMap(metricName, metric);
+        };
+        ScheduledExecutorService customScheduler = Executors.newScheduledThreadPool(2);
+        kafkaMetrics = new KafkaMetrics(supplier, Collections.emptyList(), customScheduler);
+        MeterRegistry registry = new SimpleMeterRegistry();
+
+        kafkaMetrics.bindTo(registry);
+        assertThat(registry.getMeters()).hasSize(1);
+        await().until(() -> !isDefaultMetricsSchedulerThreadAlive());
+
+        kafkaMetrics.close();
+        assertThat(registry.getMeters()).isEmpty();
+        assertThat(customScheduler.isShutdown()).isFalse();
+
+        customScheduler.shutdownNow();
+        assertThat(customScheduler.isShutdown()).isTrue();
     }
 
     @Test
@@ -552,4 +581,13 @@ private KafkaMetric createKafkaMetric(MetricName metricName) {
         return new KafkaMetric(this, metricName, new Value(), new MetricConfig(), Time.SYSTEM);
     }
 
+    private static boolean isDefaultMetricsSchedulerThreadAlive() {
+        return Thread.getAllStackTraces()
+            .keySet()
+            .stream()
+            .filter(Thread::isAlive)
+            .map(Thread::getName)
+            .anyMatch(name -> name.startsWith("micrometer-kafka-metrics"));
+    }
+
 }