Make ZK client timeouts configurable (#1380)

* Make ZK client timeouts configurable

* Wire ZK clients timeouts to @embeddedkafka properties

* Splitting the ZK client timeout properties, @author

* @SInCE on new API

Author: pawellozinski

spring-kafka-test/src/main/java/org/springframework/kafka/test/EmbeddedKafkaBroker.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2019 the original author or authors.
+ * Copyright 2018-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.
@@ -77,6 +77,7 @@
  * @author Kamill Sokol
  * @author Elliot Kennedy
  * @author Nakul Mishra
+ * @author Pawel Lozinski
  *
  * @since 2.2
  */
@@ -100,9 +101,9 @@ public class EmbeddedKafkaBroker implements InitializingBean, DisposableBean {
 
 	private static final Duration DEFAULT_ADMIN_TIMEOUT = Duration.ofSeconds(10);
 
-	private static final int ZK_CONNECTION_TIMEOUT = 6000;
+	public static final int DEFAULT_ZK_CONNECTION_TIMEOUT = 6000;
 
-	private static final int ZK_SESSION_TIMEOUT = 6000;
+	public static final int DEFAULT_ZK_SESSION_TIMEOUT = 6000;
 
 	private final int count;
 
@@ -126,6 +127,10 @@ public class EmbeddedKafkaBroker implements InitializingBean, DisposableBean {
 
 	private Duration adminTimeout = DEFAULT_ADMIN_TIMEOUT;
 
+	private int zkConnectionTimeout = DEFAULT_ZK_CONNECTION_TIMEOUT;
+
+	private int zkSessionTimeout = DEFAULT_ZK_SESSION_TIMEOUT;
+
 	private String brokerListProperty;
 
 	private volatile ZooKeeperClient zooKeeperClient;
@@ -249,6 +254,28 @@ public void setZkPort(int zkPort) {
 		this.zkPort = zkPort;
 	}
 
+	/**
+	 * Set connection timeout for the client to the embedded Zookeeper.
+	 * @param zkConnectionTimeout the connection timeout,
+	 * @return the {@link EmbeddedKafkaBroker}.
+	 * @since 2.4
+	 */
+	public EmbeddedKafkaBroker zkConnectionTimeout(int zkConnectionTimeout) {
+		this.zkConnectionTimeout = zkConnectionTimeout;
+		return this;
+	}
+
+	/**
+	 * Set session timeout for the client to the embedded Zookeeper.
+	 * @param zkSessionTimeout the session timeout.
+	 * @return the {@link EmbeddedKafkaBroker}.
+	 * @since 2.4
+	 */
+	public EmbeddedKafkaBroker zkSessionTimeout(int zkSessionTimeout) {
+		this.zkSessionTimeout = zkSessionTimeout;
+		return this;
+	}
+
 	@Override
 	public void afterPropertiesSet() {
 		try {
@@ -336,8 +363,8 @@ private void createKafkaTopics(Set<String> topicsToCreate) {
 		doWithAdmin(admin -> {
 			createTopics(admin,
 					topicsToCreate.stream()
-							.map(t -> new NewTopic(t, this.partitionsPerTopic, (short) this.count))
-							.collect(Collectors.toList()));
+						.map(t -> new NewTopic(t, this.partitionsPerTopic, (short) this.count))
+						.collect(Collectors.toList()));
 		});
 	}
 
@@ -428,7 +455,7 @@ public EmbeddedZookeeper getZookeeper() {
 	 */
 	public synchronized ZooKeeperClient getZooKeeperClient() {
 		if (this.zooKeeperClient == null) {
-			this.zooKeeperClient = new ZooKeeperClient(this.zkConnect, ZK_SESSION_TIMEOUT, ZK_CONNECTION_TIMEOUT,
+			this.zooKeeperClient = new ZooKeeperClient(this.zkConnect, zkSessionTimeout, zkConnectionTimeout,
 					1, Time.SYSTEM, "embeddedKafkaZK", "embeddedKafkaZK");
 		}
 		return this.zooKeeperClient;

spring-kafka-test/src/main/java/org/springframework/kafka/test/condition/EmbeddedKafkaCondition.java

@@ -50,6 +50,7 @@
  *
  * @author Gary Russell
  * @author Artem Bilan
+ * @author Pawel Lozinski
  *
  * @since 2.3
  *
@@ -118,7 +119,9 @@ private EmbeddedKafkaBroker createBroker(EmbeddedKafka embedded) {
 		int[] ports = setupPorts(embedded);
 		broker = new EmbeddedKafkaBroker(embedded.count(), embedded.controlledShutdown(), embedded.topics())
 				.zkPort(embedded.zookeeperPort())
-				.kafkaPorts(ports);
+				.kafkaPorts(ports)
+				.zkConnectionTimeout(embedded.zkConnectionTimeout())
+				.zkSessionTimeout(embedded.zkSessionTimeout());
 		Properties properties = new Properties();
 
 		for (String pair : embedded.brokerProperties()) {

spring-kafka-test/src/main/java/org/springframework/kafka/test/context/EmbeddedKafka.java

@@ -26,6 +26,7 @@
 import org.junit.jupiter.api.extension.ExtendWith;
 
 import org.springframework.core.annotation.AliasFor;
+import org.springframework.kafka.test.EmbeddedKafkaBroker;
 import org.springframework.kafka.test.condition.EmbeddedKafkaCondition;
 
 /**
@@ -57,6 +58,7 @@
  * @author Zach Olauson
  * @author Gary Russell
  * @author Sergio Lourenco
+ * @author Pawel Lozinski
  *
  * @since 1.3
  *
@@ -155,5 +157,19 @@
 	 */
 	String bootstrapServersProperty() default "";
 
+	/**
+	 * Timeout for internal ZK client connection.
+	 * @return default {@link EmbeddedKafkaBroker#DEFAULT_ZK_CONNECTION_TIMEOUT}.
+	 * @since 2.4
+	 */
+	int zkConnectionTimeout() default EmbeddedKafkaBroker.DEFAULT_ZK_CONNECTION_TIMEOUT;
+
+	/**
+	 * Timeout for internal ZK client session.
+	 * @return default {@link EmbeddedKafkaBroker#DEFAULT_ZK_SESSION_TIMEOUT}.
+	 * @since 2.4
+	 */
+	int zkSessionTimeout() default EmbeddedKafkaBroker.DEFAULT_ZK_SESSION_TIMEOUT;
+
 }
 

spring-kafka-test/src/main/java/org/springframework/kafka/test/context/EmbeddedKafkaContextCustomizer.java

@@ -42,6 +42,7 @@
  * @author Zach Olauson
  * @author Oleg Artyomov
  * @author Sergio Lourenco
+ * @author Pawel Lozinski
  *
  * @since 1.3
  */
@@ -72,7 +73,9 @@ public void customizeContext(ConfigurableApplicationContext context, MergedConte
 					this.embeddedKafka.partitions(),
 					topics)
 				.kafkaPorts(ports)
-				.zkPort(this.embeddedKafka.zookeeperPort());
+				.zkPort(this.embeddedKafka.zookeeperPort())
+				.zkConnectionTimeout(this.embeddedKafka.zkConnectionTimeout())
+				.zkSessionTimeout(this.embeddedKafka.zkSessionTimeout());
 
 		Properties properties = new Properties();