Add timeout when flushing items to Kafka in KafkaItemWriter

Issue #3773

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

@@ -25,6 +25,7 @@ import org.springframework.util.concurrent.ListenableFuture;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 /**
  * <p>
@@ -33,6 +34,7 @@ import java.util.List;
  * </p>
  *
  * @author Mathieu Ouellet
+ * @author Mahmoud Ben Hassine
  * @since 4.2
  *
  */
@@ -40,6 +42,7 @@ 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 @@ public class KafkaItemWriter<K, T> extends KeyValueItemWriter<K, T> {
 	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 @@ public class KafkaItemWriter<K, T> extends KeyValueItemWriter<K, T> {
 	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 @@ import org.springframework.util.Assert;
  * 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 class KafkaItemWriterBuilder<K, V> {
 		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 class KafkaItemWriterBuilder<K, V> {
 		writer.setKafkaTemplate(this.kafkaTemplate);
 		writer.setItemKeyMapper(this.itemKeyMapper);
 		writer.setDelete(this.delete);
+		writer.setTimeout(this.timeout);
 		return writer;
 	}
 }