GH-1520: Add Producer/Consumer Post Processors

Resolves #1520

To assist with Sleuth instrumentation.

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/core/ConsumerFactory.java

@@ -16,6 +16,8 @@
 
 package org.springframework.kafka.core;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 
@@ -136,6 +138,72 @@ default Deserializer<V> getValueDeserializer() {
 		return null;
 	}
 
+	/**
+	 * Remove a listener.
+	 * @param listener the listener.
+	 * @return true if removed.
+	 * @since 2.5.3
+	 */
+	default boolean removeListener(Listener<K, V> listener) {
+		return false;
+	}
+
+	/**
+	 * Add a listener at a specific index.
+	 * @param index the index (list position).
+	 * @param listener the listener.
+	 * @since 2.5.3
+	 */
+	default void addListener(int index, Listener<K, V> listener) {
+
+	}
+
+	/**
+	 * Add a listener.
+	 * @param listener the listener.
+	 * @since 2.5.3
+	 */
+	default void addListener(Listener<K, V> listener) {
+
+	}
+
+	/**
+	 * Get the current list of listeners.
+	 * @return the listeners.
+	 * @since 2.5.3
+	 */
+	default List<Listener<K, V>> getListeners() {
+		return Collections.emptyList();
+	}
+
+	/**
+	 * Add a post processor.
+	 * @param postProcessor the post processor.
+	 * @since 2.5.3
+	 */
+	default void addPostProcessor(ConsumerPostProcessor<K, V> postProcessor) {
+
+	}
+
+	/**
+	 * Remove a post processor.
+	 * @param postProcessor the post processor.
+	 * @return true if removed.
+	 * @since 2.5.3
+	 */
+	default boolean removePostProcessor(ConsumerPostProcessor<K, V> postProcessor) {
+		return false;
+	}
+
+	/**
+	 * Get the current list of post processors.
+	 * @return the post processor.
+	 * @since 2.5.3
+	 */
+	default List<ConsumerPostProcessor<K, V>> getPostProcessors() {
+		return Collections.emptyList();
+	}
+
 	/**
 	 * Called whenever a consumer is added or removed.
 	 *

spring-kafka/src/main/java/org/springframework/kafka/core/ConsumerPostProcessor.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 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.kafka.core;
+
+import java.util.function.Function;
+
+import org.apache.kafka.clients.consumer.Consumer;
+
+/**
+ * Called by consumer factories to perform post processing on newly created consumers.
+ *
+ * @param <K> the key type.
+ * @param <V> the value type
+ *
+ * @author Gary Russell
+ * @since 2.5.3
+ *
+ */
+public interface ConsumerPostProcessor<K, V> extends Function<Consumer<K, V>, Consumer<K, V>> {
+
+}

spring-kafka/src/main/java/org/springframework/kafka/core/DefaultKafkaConsumerFactory.java

@@ -78,6 +78,8 @@ public class DefaultKafkaConsumerFactory<K, V> extends KafkaResourceFactory
 
 	private final List<Listener<K, V>> listeners = new ArrayList<>();
 
+	private final List<ConsumerPostProcessor<K, V>> postProcessors = new ArrayList<>();
+
 	private Supplier<Deserializer<K>> keyDeserializerSupplier;
 
 	private Supplier<Deserializer<V>> valueDeserializerSupplier;
@@ -165,15 +167,22 @@ public Deserializer<V> getValueDeserializer() {
 	 * @return the listeners.
 	 * @since 2.5
 	 */
+	@Override
 	public List<Listener<K, V>> getListeners() {
 		return Collections.unmodifiableList(this.listeners);
 	}
 
+	@Override
+	public List<ConsumerPostProcessor<K, V>> getPostProcessors() {
+		return Collections.unmodifiableList(this.postProcessors);
+	}
+
 	/**
 	 * Add a listener.
 	 * @param listener the listener.
 	 * @since 2.5
 	 */
+	@Override
 	public void addListener(Listener<K, V> listener) {
 		Assert.notNull(listener, "'listener' cannot be null");
 		this.listeners.add(listener);
@@ -185,6 +194,7 @@ public void addListener(Listener<K, V> listener) {
 	 * @param listener the listener.
 	 * @since 2.5
 	 */
+	@Override
 	public void addListener(int index, Listener<K, V> listener) {
 		Assert.notNull(listener, "'listener' cannot be null");
 		if (index >= this.listeners.size()) {
@@ -195,12 +205,24 @@ public void addListener(int index, Listener<K, V> listener) {
 		}
 	}
 
+	@Override
+	public void addPostProcessor(ConsumerPostProcessor<K, V> postProcessor) {
+		Assert.notNull(postProcessor, "'postProcessor' cannot be null");
+		this.postProcessors.add(postProcessor);
+	}
+
+	@Override
+	public boolean removePostProcessor(ConsumerPostProcessor<K, V> postProcessor) {
+		return this.postProcessors.remove(postProcessor);
+	}
+
 	/**
 	 * Remove a listener.
 	 * @param listener the listener.
 	 * @return true if removed.
 	 * @since 2.5
 	 */
+	@Override
 	public boolean removeListener(Listener<K, V> listener) {
 		return this.listeners.remove(listener);
 	}
@@ -301,6 +323,9 @@ protected Consumer<K, V> createKafkaConsumer(Map<String, Object> configProps) {
 				listener.consumerAdded(id, kafkaConsumer);
 			}
 		}
+		for (ConsumerPostProcessor<K, V> pp : this.postProcessors) {
+			pp.apply(kafkaConsumer);
+		}
 		return kafkaConsumer;
 	}
 

spring-kafka/src/main/java/org/springframework/kafka/core/DefaultKafkaProducerFactory.java

@@ -132,6 +132,8 @@ public class DefaultKafkaProducerFactory<K, V> extends KafkaResourceFactory
 
 	private final List<Listener<K, V>> listeners = new ArrayList<>();
 
+	private final List<ProducerPostProcessor<K, V>> postProcessors = new ArrayList<>();
+
 	private Supplier<Serializer<K>> keySerializerSupplier;
 
 	private Supplier<Serializer<V>> valueSerializerSupplier;
@@ -337,15 +339,22 @@ public Map<String, Object> getConfigurationProperties() {
 	 * @return the listeners.
 	 * @since 2.5
 	 */
+	@Override
 	public List<Listener<K, V>> getListeners() {
 		return Collections.unmodifiableList(this.listeners);
 	}
 
+	@Override
+	public List<ProducerPostProcessor<K, V>> getPostProcessors() {
+		return Collections.unmodifiableList(this.postProcessors);
+	}
+
 	/**
 	 * Add a listener.
 	 * @param listener the listener.
 	 * @since 2.5
 	 */
+	@Override
 	public void addListener(Listener<K, V> listener) {
 		Assert.notNull(listener, "'listener' cannot be null");
 		this.listeners.add(listener);
@@ -357,6 +366,7 @@ public void addListener(Listener<K, V> listener) {
 	 * @param listener the listener.
 	 * @since 2.5
 	 */
+	@Override
 	public void addListener(int index, Listener<K, V> listener) {
 		Assert.notNull(listener, "'listener' cannot be null");
 		if (index >= this.listeners.size()) {
@@ -373,10 +383,22 @@ public void addListener(int index, Listener<K, V> listener) {
 	 * @return true if removed.
 	 * @since 2.5
 	 */
+	@Override
 	public boolean removeListener(Listener<K, V> listener) {
 		return this.listeners.remove(listener);
 	}
 
+	@Override
+	public void addPostProcessor(ProducerPostProcessor<K, V> postProcessor) {
+		Assert.notNull(postProcessor, "'postProcessor' cannot be null");
+		this.postProcessors.add(postProcessor);
+	}
+
+	@Override
+	public boolean removePostProcessor(ProducerPostProcessor<K, V> postProcessor) {
+		return this.postProcessors.remove(postProcessor);
+	}
+
 	/**
 	 * When set to 'true', the producer will ensure that exactly one copy of each message is written in the stream.
 	 */
@@ -642,7 +664,10 @@ private CloseSafeProducer<K, V> doCreateTxProducer(String prefix, String suffix,
 	}
 
 	protected Producer<K, V> createRawProducer(Map<String, Object> rawConfigs) {
-		return new KafkaProducer<>(rawConfigs, this.keySerializerSupplier.get(), this.valueSerializerSupplier.get());
+		KafkaProducer<K, V> kafkaProducer =
+				new KafkaProducer<>(rawConfigs, this.keySerializerSupplier.get(), this.valueSerializerSupplier.get());
+		this.postProcessors.forEach(pp -> pp.apply(kafkaProducer));
+		return kafkaProducer;
 	}
 
 	@Nullable

spring-kafka/src/main/java/org/springframework/kafka/core/ProducerFactory.java

@@ -17,6 +17,8 @@
 package org.springframework.kafka.core;
 
 import java.time.Duration;
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
 
@@ -168,6 +170,72 @@ default Duration getPhysicalCloseTimeout() {
 		return DEFAULT_PHYSICAL_CLOSE_TIMEOUT;
 	}
 
+	/**
+	 * Add a listener.
+	 * @param listener the listener.
+	 * @since 2.5.3
+	 */
+	default void addListener(Listener<K, V> listener) {
+
+	}
+
+	/**
+	 * Add a listener at a specific index.
+	 * @param index the index (list position).
+	 * @param listener the listener.
+	 * @since 2.5.3
+	 */
+	default void addListener(int index, Listener<K, V> listener) {
+
+	}
+
+	/**
+	 * Remove a listener.
+	 * @param listener the listener.
+	 * @return true if removed.
+	 * @since 2.5.3
+	 */
+	default boolean removeListener(Listener<K, V> listener) {
+		return false;
+	}
+
+	/**
+	 * Get the current list of listeners.
+	 * @return the listeners.
+	 * @since 2.5.3
+	 */
+	default List<Listener<K, V>> getListeners() {
+		return Collections.emptyList();
+	}
+
+	/**
+	 * Add a post processor.
+	 * @param postProcessor the post processor.
+	 * @since 2.5.3
+	 */
+	default void addPostProcessor(ProducerPostProcessor<K, V> postProcessor) {
+
+	}
+
+	/**
+	 * Remove a post processor.
+	 * @param postProcessor the post processor.
+	 * @return true if removed.
+	 * @since 2.5.3
+	 */
+	default boolean removePostProcessor(ProducerPostProcessor<K, V> postProcessor) {
+		return false;
+	}
+
+	/**
+	 * Get the current list of post processors.
+	 * @return the post processors.
+	 * @since 2.5.3
+	 */
+	default List<ProducerPostProcessor<K, V>> getPostProcessors() {
+		return Collections.emptyList();
+	}
+
 	/**
 	 * Called whenever a producer is added or removed.
 	 *

spring-kafka/src/main/java/org/springframework/kafka/core/ProducerPostProcessor.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 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.kafka.core;
+
+import java.util.function.Function;
+
+import org.apache.kafka.clients.producer.Producer;
+
+/**
+ * Called by producer factories to perform post processing on newly created producers.
+ *
+ * @param <K> the key type.
+ * @param <V> the value type
+ *
+ * @author Gary Russell
+ * @since 2.5.3
+ *
+ */
+public interface ProducerPostProcessor<K, V> extends Function<Producer<K, V>, Producer<K, V>> {
+
+}