Add timeout when flushing items to Kafka in KafkaItemWriter

Issue #3773

Author: fmbenhassine

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/KafkaItemWriter.java

@@ -25,6 +25,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 /**
  * <p>
@@ -33,13 +34,15 @@
  * </p>
  *
  * @author Mathieu Ouellet
+ * @author Mahmoud Ben Hassine
  * @since 4.2
  *
  */
 public class KafkaItemWriter<K, T> extends KeyValueItemWriter<K, T> {
 
 	protected KafkaTemplate<K, T> kafkaTemplate;
 	private final List<ListenableFuture<SendResult<K, T>>> listenableFutures = new ArrayList<>();
+	private long timeout = -1;
 
 	@Override
 	protected void writeKeyValue(K key, T value) {
@@ -55,7 +58,12 @@ protected void writeKeyValue(K key, T value) {
 	protected void flush() throws Exception{
 		this.kafkaTemplate.flush();
 		for(ListenableFuture<SendResult<K,T>> future: this.listenableFutures){
-			future.get();
+			if (this.timeout >= 0) {
+				future.get(this.timeout, TimeUnit.MILLISECONDS);
+			}
+			else {
+				future.get();
+			}
 		}
 		this.listenableFutures.clear();
 	}
@@ -73,4 +81,15 @@ protected void init() {
 	public void setKafkaTemplate(KafkaTemplate<K, T> kafkaTemplate) {
 		this.kafkaTemplate = kafkaTemplate;
 	}
+
+	/**
+	 * The time limit to wait when flushing items to Kafka.
+	 *
+	 * @param timeout milliseconds to wait, defaults to -1 (no timeout).
+	 * @since 4.3.2
+	 */
+	public void setTimeout(long timeout) {
+		this.timeout = timeout;
+	}
+
 }

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/builder/KafkaItemWriterBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 the original author or authors.
+ * Copyright 2019-2021 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.
@@ -25,6 +25,7 @@
  * A builder implementation for the {@link KafkaItemWriter}
  *
  * @author Mathieu Ouellet
+ * @author Mahmoud Ben Hassine
  * @since 4.2
  */
 public class KafkaItemWriterBuilder<K, V> {
@@ -35,6 +36,8 @@ public class KafkaItemWriterBuilder<K, V> {
 
 	private boolean delete;
 
+	private long timeout = -1;
+
 	/**
 	 * Establish the KafkaTemplate to be used by the KafkaItemWriter.
 	 * @param kafkaTemplate the template to be used
@@ -71,6 +74,19 @@ public KafkaItemWriterBuilder<K, V> delete(boolean delete) {
 		return this;
 	}
 
+	/**
+	 * The time limit to wait when flushing items to Kafka.
+	 *
+	 * @param timeout milliseconds to wait, defaults to -1 (no timeout).
+	 * @return The current instance of the builder.
+	 * @see KafkaItemWriter#setTimeout(long)
+	 * @since 4.3.2
+	 */
+	public KafkaItemWriterBuilder<K, V> timeout(long timeout) {
+		this.timeout = timeout;
+		return this;
+	}
+
 	/**
 	 * Validates and builds a {@link KafkaItemWriter}.
 	 * @return a {@link KafkaItemWriter}
@@ -83,6 +99,7 @@ public KafkaItemWriter<K, V> build() {
 		writer.setKafkaTemplate(this.kafkaTemplate);
 		writer.setItemKeyMapper(this.itemKeyMapper);
 		writer.setDelete(this.delete);
+		writer.setTimeout(this.timeout);
 		return writer;
 	}
 }

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/KafkaItemWriterTests.java

@@ -17,6 +17,7 @@
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -57,6 +58,7 @@ public void setUp() throws Exception {
 		this.writer.setKafkaTemplate(this.kafkaTemplate);
 		this.writer.setItemKeyMapper(this.itemKeyMapper);
 		this.writer.setDelete(false);
+		this.writer.setTimeout(10L);
 		this.writer.afterPropertiesSet();
 	}
 
@@ -99,7 +101,7 @@ public void testBasicWrite() throws Exception {
 		verify(this.kafkaTemplate).sendDefault(items.get(0), items.get(0));
 		verify(this.kafkaTemplate).sendDefault(items.get(1), items.get(1));
 		verify(this.kafkaTemplate).flush();
-		verify(this.future, times(2)).get();
+		verify(this.future, times(2)).get(10L, TimeUnit.MILLISECONDS);
 	}
 
 	@Test
@@ -112,7 +114,7 @@ public void testBasicDelete() throws Exception {
 		verify(this.kafkaTemplate).sendDefault(items.get(0), null);
 		verify(this.kafkaTemplate).sendDefault(items.get(1), null);
 		verify(this.kafkaTemplate).flush();
-		verify(this.future, times(2)).get();
+		verify(this.future, times(2)).get(10L, TimeUnit.MILLISECONDS);
 	}
 
 	@Test

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/builder/KafkaItemWriterBuilderTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 the original author or authors.
+ * Copyright 2019-2021 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.
@@ -33,6 +33,7 @@
 
 /**
  * @author Mathieu Ouellet
+ * @author Mahmoud Ben Hassine
  */
 public class KafkaItemWriterBuilderTests {
 
@@ -70,13 +71,19 @@ public void testNullItemKeyMapper() {
 	public void testKafkaItemWriterBuild() {
 		// given
 		boolean delete = true;
+		long timeout = 10L;
 
 		// when
 		KafkaItemWriter<String, String> writer = new KafkaItemWriterBuilder<String, String>()
-				.kafkaTemplate(this.kafkaTemplate).itemKeyMapper(this.itemKeyMapper).delete(delete).build();
+				.kafkaTemplate(this.kafkaTemplate)
+				.itemKeyMapper(this.itemKeyMapper)
+				.delete(delete)
+				.timeout(timeout)
+				.build();
 
 		// then
 		assertTrue((Boolean) ReflectionTestUtils.getField(writer, "delete"));
+		assertEquals(timeout, ReflectionTestUtils.getField(writer, "timeout"));
 		assertEquals(this.itemKeyMapper, ReflectionTestUtils.getField(writer, "itemKeyMapper"));
 		assertEquals(this.kafkaTemplate, ReflectionTestUtils.getField(writer, "kafkaTemplate"));
 	}