GH-61: Add (De)Serializer options for Factories

artembilan · artembilan · commit 48845e1de043 · 2016-05-11T17:13:32.000-04:00
Fixes GH-61 (#61) Added setter injections for both (De)Serializers in Consumer and Producer
diff --git a/spring-kafka/src/main/java/org/springframework/kafka/core/DefaultKafkaConsumerFactory.java b/spring-kafka/src/main/java/org/springframework/kafka/core/DefaultKafkaConsumerFactory.java
@@ -22,28 +22,55 @@
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.common.serialization.Deserializer;
 
 /**
  * The {@link ConsumerFactory} implementation to produce a new {@link Consumer} instance
- * for provided {@link Map} {@code configs} on each {@link #createConsumer()}
+ * for provided {@link Map} {@code configs} and optional {@link Deserializer} {@code keyDeserializer},
+ * {@code valueDeserializer} implementations on each {@link #createConsumer()}
  * invocation.
  *
  * @param <K> the key type.
  * @param <V> the value type.
  *
  * @author Gary Russell
+ * @author Murali Reddy
  */
 public class DefaultKafkaConsumerFactory<K, V> implements ConsumerFactory<K, V> {
 
 	private final Map<String, Object> configs;
 
+	private Deserializer<K> keyDeserializer;
+
+	private Deserializer<V> valueDeserializer;
+
 	public DefaultKafkaConsumerFactory(Map<String, Object> configs) {
+		this(configs, null, null);
+	}
+
+	public DefaultKafkaConsumerFactory(Map<String, Object> configs,
+			Deserializer<K> keyDeserializer,
+			Deserializer<V> valueDeserializer) {
 		this.configs = new HashMap<>(configs);
+		this.keyDeserializer = keyDeserializer;
+		this.valueDeserializer = valueDeserializer;
+	}
+
+	public void setKeyDeserializer(Deserializer<K> keyDeserializer) {
+		this.keyDeserializer = keyDeserializer;
+	}
+
+	public void setValueDeserializer(Deserializer<V> valueDeserializer) {
+		this.valueDeserializer = valueDeserializer;
 	}
 
 	@Override
 	public Consumer<K, V> createConsumer() {
-		return new KafkaConsumer<>(this.configs);
+		return createKafkaConsumer();
+	}
+
+	protected KafkaConsumer<K, V> createKafkaConsumer() {
+		return new KafkaConsumer<K, V>(this.configs, this.keyDeserializer, this.valueDeserializer);
 	}
 
 	@Override
diff --git a/spring-kafka/src/main/java/org/springframework/kafka/core/DefaultKafkaProducerFactory.java b/spring-kafka/src/main/java/org/springframework/kafka/core/DefaultKafkaProducerFactory.java
@@ -24,6 +24,7 @@
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.Producer;
@@ -32,6 +33,7 @@
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.serialization.Serializer;
 
 import org.springframework.beans.factory.DisposableBean;
 import org.springframework.context.Lifecycle;
@@ -40,6 +42,11 @@
  * The {@link ProducerFactory} implementation for the {@code singleton} shared {@link Producer}
  * instance.
  * <p>
+ * This implementation will produce a new {@link Producer} instance
+ * for provided {@link Map} {@code configs} and optional {@link Serializer} {@code keySerializer},
+ * {@code valueSerializer} implementations on each {@link #createProducer()}
+ * invocation.
+ * <p>
  * The {@link Producer} instance is freed from the external {@link Producer#close()} invocation
  * with the internal wrapper. The real {@link Producer#close()} is called on the target
  * {@link Producer} during the {@link Lifecycle#stop()} or {@link DisposableBean#destroy()}.
@@ -48,6 +55,7 @@
  * @param <V> the value type.
  *
  * @author Gary Russell
+ * @author Murali Reddy
  */
 public class DefaultKafkaProducerFactory<K, V> implements ProducerFactory<K, V>, Lifecycle, DisposableBean {
 
@@ -57,10 +65,29 @@ public class DefaultKafkaProducerFactory<K, V> implements ProducerFactory<K, V>,
 
 	private volatile CloseSafeProducer<K, V> producer;
 
+	private Serializer<K> keySerializer;
+
+	private Serializer<V> valueSerializer;
+
 	private volatile boolean running;
 
 	public DefaultKafkaProducerFactory(Map<String, Object> configs) {
+		this(configs, null, null);
+	}
+
+	public DefaultKafkaProducerFactory(Map<String, Object> configs, Serializer<K> keySerializer,
+			Serializer<V> valueSerializer) {
 		this.configs = new HashMap<>(configs);
+		this.keySerializer = keySerializer;
+		this.valueSerializer = valueSerializer;
+	}
+
+	public void setKeySerializer(Serializer<K> keySerializer) {
+		this.keySerializer = keySerializer;
+	}
+
+	public void setValueSerializer(Serializer<V> valueSerializer) {
+		this.valueSerializer = valueSerializer;
 	}
 
 	@Override
@@ -100,13 +127,17 @@ public Producer<K, V> createProducer() {
 		if (this.producer == null) {
 			synchronized (this) {
 				if (this.producer == null) {
-					this.producer = new CloseSafeProducer<K, V>(new KafkaProducer<K, V>(this.configs));
+					this.producer = new CloseSafeProducer<K, V>(createKafkaProducer());
 				}
 			}
 		}
 		return this.producer;
 	}
 
+	protected KafkaProducer<K, V> createKafkaProducer() {
+		return new KafkaProducer<K, V>(this.configs, this.keySerializer, this.valueSerializer);
+	}
+
 	private static class CloseSafeProducer<K, V> implements Producer<K, V> {
 
 		private final Producer<K, V> delegate;
diff --git a/spring-kafka/src/test/java/org/springframework/kafka/annotation/EnableKafkaIntegrationTests.java b/spring-kafka/src/test/java/org/springframework/kafka/annotation/EnableKafkaIntegrationTests.java
@@ -187,15 +187,15 @@ public PlatformTransactionManager transactionManager() {
 
 		@Bean
 		public KafkaListenerContainerFactory<ConcurrentMessageListenerContainer<Integer, String>>
-				kafkaListenerContainerFactory() {
+		kafkaListenerContainerFactory() {
 			SimpleKafkaListenerContainerFactory<Integer, String> factory = new SimpleKafkaListenerContainerFactory<>();
 			factory.setConsumerFactory(consumerFactory());
 			return factory;
 		}
 
 		@Bean
 		public KafkaListenerContainerFactory<ConcurrentMessageListenerContainer<Integer, String>>
-				kafkaJsonListenerContainerFactory() {
+		kafkaJsonListenerContainerFactory() {
 			SimpleKafkaListenerContainerFactory<Integer, String> factory = new SimpleKafkaListenerContainerFactory<>();
 			factory.setConsumerFactory(consumerFactory());
 			factory.setMessageConverter(new StringJsonMessageConverter());
@@ -204,7 +204,7 @@ public PlatformTransactionManager transactionManager() {
 
 		@Bean
 		public KafkaListenerContainerFactory<ConcurrentMessageListenerContainer<Integer, String>>
-				kafkaManualAckListenerContainerFactory() {
+		kafkaManualAckListenerContainerFactory() {
 			SimpleKafkaListenerContainerFactory<Integer, String> factory = new SimpleKafkaListenerContainerFactory<>();
 			factory.setConsumerFactory(manualConsumerFactory());
 			factory.setAckMode(AckMode.MANUAL_IMMEDIATE);
@@ -213,7 +213,7 @@ public PlatformTransactionManager transactionManager() {
 
 		@Bean
 		public KafkaListenerContainerFactory<ConcurrentMessageListenerContainer<Integer, String>>
-				kafkaAutoStartFalseListenerContainerFactory() {
+		kafkaAutoStartFalseListenerContainerFactory() {
 			SimpleKafkaListenerContainerFactory<Integer, String> factory = new SimpleKafkaListenerContainerFactory<>();
 			factory.setConsumerFactory(consumerFactory());
 			factory.setAutoStartup(false);
@@ -222,7 +222,7 @@ public PlatformTransactionManager transactionManager() {
 
 		@Bean
 		public KafkaListenerContainerFactory<ConcurrentMessageListenerContainer<Integer, String>>
-				kafkaRebalanceListenerContainerFactory() {
+		kafkaRebalanceListenerContainerFactory() {
 			SimpleKafkaListenerContainerFactory<Integer, String> factory = new SimpleKafkaListenerContainerFactory<>();
 			factory.setConsumerFactory(consumerFactory());
 			factory.setConsumerRebalanceListener(consumerRebalanceListener());
@@ -366,9 +366,9 @@ public void listen4(@Payload String foo, Acknowledgment ack) {
 		}
 
 		@KafkaListener(id = "fiz", topicPartitions = {
-				@TopicPartition(topic = "annotated5", partitions = {"0", "1"}),
-				@TopicPartition(topic = "annotated6", partitions = {"0", "1"})
-			})
+				@TopicPartition(topic = "annotated5", partitions = { "0", "1" }),
+				@TopicPartition(topic = "annotated6", partitions = { "0", "1" })
+		})
 		public void listen5(ConsumerRecord<?, ?> record) {
 			this.record = record;
 			this.latch5.countDown();
diff --git a/spring-kafka/src/test/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainerTests.java b/spring-kafka/src/test/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainerTests.java
@@ -172,6 +172,60 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
 		logger.info("Stop auto");
 	}
 
+	@Test
+	public void testAutoCommitWithRebalanceListener() throws Exception {
+		logger.info("Start auto");
+		Map<String, Object> props = KafkaTestUtils.consumerProps("test10", "true", embeddedKafka);
+		DefaultKafkaConsumerFactory<Integer, String> cf = new DefaultKafkaConsumerFactory<Integer, String>(props);
+		ConcurrentMessageListenerContainer<Integer, String> container =
+				new ConcurrentMessageListenerContainer<>(cf, topic1);
+		final CountDownLatch latch = new CountDownLatch(4);
+		container.setMessageListener(new MessageListener<Integer, String>() {
+
+			@Override
+			public void onMessage(ConsumerRecord<Integer, String> message) {
+				logger.info("auto: " + message);
+				latch.countDown();
+			}
+		});
+		final CountDownLatch rebalancePartitionsAssignedLatch = new CountDownLatch(2);
+		final CountDownLatch rebalancePartitionsRevokedLatch = new CountDownLatch(2);
+		container.setConsumerRebalanceListener(new ConsumerRebalanceListener() {
+
+			@Override
+			public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+				logger.info("In test, partitions revoked:" + partitions);
+				rebalancePartitionsRevokedLatch.countDown();
+			}
+
+			@Override
+			public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
+				logger.info("In test, partitions assigned:" + partitions);
+				rebalancePartitionsAssignedLatch.countDown();
+			}
+
+		});
+
+		container.setConcurrency(2);
+		container.setBeanName("testAuto");
+		container.start();
+		ContainerTestUtils.waitForAssignment(container, embeddedKafka.getPartitionsPerTopic());
+		Map<String, Object> senderProps = KafkaTestUtils.producerProps(embeddedKafka);
+		ProducerFactory<Integer, String> pf = new DefaultKafkaProducerFactory<Integer, String>(senderProps);
+		KafkaTemplate<Integer, String> template = new KafkaTemplate<>(pf);
+		template.setDefaultTopic(topic1);
+		template.send(0, "foo");
+		template.send(2, "bar");
+		template.send(0, "baz");
+		template.send(2, "qux");
+		template.flush();
+		assertThat(latch.await(60, TimeUnit.SECONDS)).isTrue();
+		assertThat(rebalancePartitionsAssignedLatch.await(60, TimeUnit.SECONDS)).isTrue();
+		assertThat(rebalancePartitionsRevokedLatch.await(60, TimeUnit.SECONDS)).isTrue();
+		container.stop();
+		logger.info("Stop auto");
+	}
+
 	@Test
 	public void testAfterListenCommit() throws Exception {
 		logger.info("Start manual");
