Bind and unbind Kafka metrics as consumers and producers come and go

Fixes gh-21008

Co-authored-by: Andy Wilkinson <awilkinson@pivotal.io>

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/metrics/KafkaMetricsAutoConfiguration.java

@@ -16,19 +16,22 @@
 
 package org.springframework.boot.actuate.autoconfigure.metrics;
 
-import java.util.Collections;
-
 import io.micrometer.core.instrument.MeterRegistry;
 import io.micrometer.core.instrument.binder.kafka.KafkaClientMetrics;
 
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
+import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
-import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
-import org.springframework.boot.autoconfigure.condition.ConditionalOnSingleCandidate;
+import org.springframework.boot.autoconfigure.kafka.DefaultKafkaConsumerFactoryCustomizer;
+import org.springframework.boot.autoconfigure.kafka.DefaultKafkaProducerFactoryCustomizer;
 import org.springframework.boot.autoconfigure.kafka.KafkaAutoConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
+import org.springframework.kafka.core.DefaultKafkaProducerFactory;
+import org.springframework.kafka.core.MicrometerConsumerListener;
+import org.springframework.kafka.core.MicrometerProducerListener;
 import org.springframework.kafka.core.ProducerFactory;
 
 /**
@@ -39,16 +42,28 @@ import org.springframework.kafka.core.ProducerFactory;
  * @since 2.1.0
  */
 @Configuration(proxyBeanMethods = false)
-@AutoConfigureAfter({ MetricsAutoConfiguration.class, KafkaAutoConfiguration.class })
+@AutoConfigureBefore(KafkaAutoConfiguration.class)
+@AutoConfigureAfter(MetricsAutoConfiguration.class)
 @ConditionalOnClass({ KafkaClientMetrics.class, ProducerFactory.class })
 @ConditionalOnBean(MeterRegistry.class)
 public class KafkaMetricsAutoConfiguration {
 
 	@Bean
-	@ConditionalOnMissingBean
-	@ConditionalOnSingleCandidate(ProducerFactory.class)
-	public KafkaClientMetrics kafkaClientMetrics(ProducerFactory<?, ?> producerFactory) {
-		return new KafkaClientMetrics(producerFactory.createProducer(), Collections.emptyList());
+	public DefaultKafkaProducerFactoryCustomizer kafkaProducerMetrics(MeterRegistry meterRegistry) {
+		return (producerFactory) -> addListener(producerFactory, meterRegistry);
+	}
+
+	@Bean
+	public DefaultKafkaConsumerFactoryCustomizer kafkaConsumerMetrics(MeterRegistry meterRegistry) {
+		return (consumerFactory) -> addListener(consumerFactory, meterRegistry);
+	}
+
+	private <K, V> void addListener(DefaultKafkaConsumerFactory<K, V> factory, MeterRegistry meterRegistry) {
+		factory.addListener(new MicrometerConsumerListener<K, V>(meterRegistry));
+	}
+
+	private <K, V> void addListener(DefaultKafkaProducerFactory<K, V> factory, MeterRegistry meterRegistry) {
+		factory.addListener(new MicrometerProducerListener<K, V>(meterRegistry));
 	}
 
 }

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/metrics/KafkaMetricsAutoConfigurationTests.java

@@ -16,18 +16,18 @@
 
 package org.springframework.boot.actuate.autoconfigure.metrics;
 
-import io.micrometer.core.instrument.binder.kafka.KafkaClientMetrics;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.boot.actuate.autoconfigure.metrics.test.MetricsRun;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.autoconfigure.kafka.KafkaAutoConfiguration;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
-import org.springframework.context.annotation.Bean;
-import org.springframework.context.annotation.Configuration;
+import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
+import org.springframework.kafka.core.DefaultKafkaProducerFactory;
+import org.springframework.kafka.core.MicrometerConsumerListener;
+import org.springframework.kafka.core.MicrometerProducerListener;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.mockito.Mockito.mock;
 
 /**
  * Tests for {@link KafkaMetricsAutoConfiguration}.
@@ -37,35 +37,28 @@ import static org.mockito.Mockito.mock;
  */
 class KafkaMetricsAutoConfigurationTests {
 
-	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner().with(MetricsRun.simple())
+	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
 			.withConfiguration(AutoConfigurations.of(KafkaMetricsAutoConfiguration.class));
 
 	@Test
-	void whenThereIsNoProducerFactoryAutoConfigurationBacksOff() {
-		this.contextRunner.run((context) -> assertThat(context).doesNotHaveBean(KafkaClientMetrics.class));
+	void whenThereIsAMeterRegistryThenMetricsListenersAreAdded() {
+		this.contextRunner.with(MetricsRun.simple())
+				.withConfiguration(AutoConfigurations.of(KafkaAutoConfiguration.class)).run((context) -> {
+					assertThat(((DefaultKafkaProducerFactory<?, ?>) context.getBean(DefaultKafkaProducerFactory.class))
+							.getListeners()).hasSize(1).hasOnlyElementsOfTypes(MicrometerProducerListener.class);
+					assertThat(((DefaultKafkaConsumerFactory<?, ?>) context.getBean(DefaultKafkaConsumerFactory.class))
+							.getListeners()).hasSize(1).hasOnlyElementsOfTypes(MicrometerConsumerListener.class);
+				});
 	}
 
 	@Test
-	void whenThereIsAProducerFactoryKafkaClientMetricsIsConfigured() {
-		this.contextRunner.withConfiguration(AutoConfigurations.of(KafkaAutoConfiguration.class))
-				.run((context) -> assertThat(context).hasSingleBean(KafkaClientMetrics.class));
-	}
-
-	@Test
-	void allowsCustomKafkaClientMetricsToBeUsed() {
-		this.contextRunner.withConfiguration(AutoConfigurations.of(KafkaAutoConfiguration.class))
-				.withUserConfiguration(CustomKafkaClientMetricsConfiguration.class).run((context) -> assertThat(context)
-						.hasSingleBean(KafkaClientMetrics.class).hasBean("customKafkaClientMetrics"));
-	}
-
-	@Configuration(proxyBeanMethods = false)
-	static class CustomKafkaClientMetricsConfiguration {
-
-		@Bean
-		KafkaClientMetrics customKafkaClientMetrics() {
-			return mock(KafkaClientMetrics.class);
-		}
-
+	void whenThereIsNoMeterRegistryThenListenerCustomizationBacksOff() {
+		this.contextRunner.withConfiguration(AutoConfigurations.of(KafkaAutoConfiguration.class)).run((context) -> {
+			assertThat(((DefaultKafkaProducerFactory<?, ?>) context.getBean(DefaultKafkaProducerFactory.class))
+					.getListeners()).isEmpty();
+			assertThat(((DefaultKafkaConsumerFactory<?, ?>) context.getBean(DefaultKafkaConsumerFactory.class))
+					.getListeners()).isEmpty();
+		});
 	}
 
 }

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/kafka/DefaultKafkaConsumerFactoryCustomizer.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2012-2020 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.autoconfigure.kafka;
+
+import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
+
+/**
+ * Callback interface for customizing {@code DefaultKafkaConsumerFactory} beans.
+ *
+ * @author Stephane Nicoll
+ * @since 2.3.0
+ */
+@FunctionalInterface
+public interface DefaultKafkaConsumerFactoryCustomizer {
+
+	/**
+	 * Customize the {@link DefaultKafkaConsumerFactory}.
+	 * @param consumerFactory the consumer factory to customize
+	 */
+	void customize(DefaultKafkaConsumerFactory<?, ?> consumerFactory);
+
+}

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/kafka/DefaultKafkaProducerFactoryCustomizer.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2012-2020 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.autoconfigure.kafka;
+
+import org.springframework.kafka.core.DefaultKafkaProducerFactory;
+
+/**
+ * Callback interface for customizing {@code DefaultKafkaProducerFactory} beans.
+ *
+ * @author Stephane Nicoll
+ * @since 2.3.0
+ */
+@FunctionalInterface
+public interface DefaultKafkaProducerFactoryCustomizer {
+
+	/**
+	 * Customize the {@link DefaultKafkaProducerFactory}.
+	 * @param producerFactory the producer factory to customize
+	 */
+	void customize(DefaultKafkaProducerFactory<?, ?> producerFactory);
+
+}

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/kafka/KafkaAutoConfiguration.java

@@ -81,19 +81,25 @@ public class KafkaAutoConfiguration {
 
 	@Bean
 	@ConditionalOnMissingBean(ConsumerFactory.class)
-	public ConsumerFactory<?, ?> kafkaConsumerFactory() {
-		return new DefaultKafkaConsumerFactory<>(this.properties.buildConsumerProperties());
+	public ConsumerFactory<?, ?> kafkaConsumerFactory(
+			ObjectProvider<DefaultKafkaConsumerFactoryCustomizer> customizers) {
+		DefaultKafkaConsumerFactory<Object, Object> factory = new DefaultKafkaConsumerFactory<>(
+				this.properties.buildConsumerProperties());
+		customizers.orderedStream().forEach((customizer) -> customizer.customize(factory));
+		return factory;
 	}
 
 	@Bean
 	@ConditionalOnMissingBean(ProducerFactory.class)
-	public ProducerFactory<?, ?> kafkaProducerFactory() {
+	public ProducerFactory<?, ?> kafkaProducerFactory(
+			ObjectProvider<DefaultKafkaProducerFactoryCustomizer> customizers) {
 		DefaultKafkaProducerFactory<?, ?> factory = new DefaultKafkaProducerFactory<>(
 				this.properties.buildProducerProperties());
 		String transactionIdPrefix = this.properties.getProducer().getTransactionIdPrefix();
 		if (transactionIdPrefix != null) {
 			factory.setTransactionIdPrefix(transactionIdPrefix);
 		}
+		customizers.orderedStream().forEach((customizer) -> customizer.customize(factory));
 		return factory;
 	}