GH-1254: Configurable Embedded Zookeeper Port

Resolves #1254

- port `EmbeddedKafkaBroker` from scala and add ZK port configuration
- add `zookeeperPort` to `@EmbeddedKafka`

Author: garyrussell

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

@@ -18,6 +18,9 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import java.io.File;
+import java.io.IOException;
+import java.net.InetSocketAddress;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -46,6 +49,9 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.zookeeper.server.NIOServerCnxnFactory;
+import org.apache.zookeeper.server.ZooKeeperServer;
 
 import org.springframework.beans.factory.DisposableBean;
 import org.springframework.beans.factory.InitializingBean;
@@ -63,7 +69,7 @@
 import kafka.utils.CoreUtils;
 import kafka.utils.TestUtils;
 import kafka.utils.ZKStringSerializer$;
-import kafka.zk.EmbeddedZookeeper;
+import kafka.zk.ZkFourLetterWords;
 
 /**
  * An embedded Kafka Broker(s) and Zookeeper manager.
@@ -114,6 +120,8 @@ public class EmbeddedKafkaBroker implements InitializingBean, DisposableBean {
 
 	private String zkConnect;
 
+	private int zkPort;
+
 	private int[] kafkaPorts;
 
 	private Duration adminTimeout = DEFAULT_ADMIN_TIMEOUT;
@@ -190,6 +198,16 @@ public EmbeddedKafkaBroker kafkaPorts(int... ports) {
 		return this;
 	}
 
+	/**
+	 * Set an explicit port for the embedded Zookeeper.
+	 * @param port the port.
+	 * @return the {@link EmbeddedKafkaBroker}.
+	 * @since 2.3
+	 */
+	public EmbeddedKafkaBroker zkPort(int port) {
+		this.zkPort = port;
+		return this;
+	}
 	/**
 	 * Set the timeout in seconds for admin operations (e.g. topic creation, close).
 	 * Default 30 seconds.
@@ -211,13 +229,36 @@ public EmbeddedKafkaBroker brokerListProperty(String brokerListProperty) {
 		return this;
 	}
 
+	/**
+	 * Get the port that the embedded Zookeeper is running on or will run on.
+	 * @param zkPort the port.
+	 * @since 2.3
+	 */
+	public int getZkPort() {
+		return this.zookeeper != null ? this.zookeeper.getPort() : this.zkPort;
+	}
+
+	/**
+	 * Set the port to run the embedded Zookeeper on (default random).
+	 * @param zkPort the port.
+	 * @since 2.3
+	 */
+	public void setZkPort(int zkPort) {
+		this.zkPort = zkPort;
+	}
+
 	@Override
 	public void afterPropertiesSet() {
-		this.zookeeper = new EmbeddedZookeeper();
+		try {
+			this.zookeeper = new EmbeddedZookeeper(this.zkPort);
+		}
+		catch (IOException | InterruptedException e) {
+			throw new IllegalStateException("Failed to create embedded Zookeeper", e);
+		}
 		int zkConnectionTimeout = 6000; // NOSONAR magic #
 		int zkSessionTimeout = 6000; // NOSONAR magic #
 
-		this.zkConnect = "127.0.0.1:" + this.zookeeper.port();
+		this.zkConnect = "127.0.0.1:" + this.zookeeper.getPort();
 		this.zookeeperClient = new ZkClient(this.zkConnect, zkSessionTimeout, zkConnectionTimeout,
 				ZKStringSerializer$.MODULE$);
 		this.kafkaServers.clear();
@@ -360,6 +401,7 @@ public void destroy() {
 		}
 		try {
 			this.zookeeper.shutdown();
+			this.zkConnect = null;
 		}
 		catch (Exception e) {
 			// do nothing
@@ -510,4 +552,86 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
 		logger.debug("Subscription Initiated");
 	}
 
+	/**
+	 * Ported from scala to allow setting the port.
+	 *
+	 * @author Gary Russell
+	 * @since 2.3
+	 */
+	public static final class EmbeddedZookeeper {
+
+		private final NIOServerCnxnFactory factory;
+
+		private final ZooKeeperServer zookeeper;
+
+		private final int port;
+
+		private final File snapshotDir;
+
+		private final File logDir;
+
+		public EmbeddedZookeeper(int zkPort) throws IOException, InterruptedException {
+			this.snapshotDir = TestUtils.tempDir();
+			this.logDir = TestUtils.tempDir();
+			int tickTime = 800; // allow a maxSessionTimeout of 20 * 800ms = 16 secs
+
+			System.setProperty("zookeeper.forceSync", "no"); // disable fsync to ZK txn
+																// log in tests to avoid
+																// timeout
+			this.zookeeper = new ZooKeeperServer(this.snapshotDir, this.logDir, tickTime);
+			this.factory = new NIOServerCnxnFactory();
+			InetSocketAddress addr = new InetSocketAddress("127.0.0.1", zkPort == 0 ? TestUtils.RandomPort() : zkPort);
+			this.factory.configure(addr, 0);
+			this.factory.startup(zookeeper);
+			this.port = zookeeper.getClientPort();
+		}
+
+		public int getPort() {
+			return this.port;
+		}
+
+		public File getSnapshotDir() {
+			return this.snapshotDir;
+		}
+
+		public File getLogDir() {
+			return this.logDir;
+		}
+
+		public void shutdown() throws IOException {
+			// Also shuts down ZooKeeperServer
+			try {
+				this.factory.shutdown();
+			}
+			catch (Exception e) {
+				logger.error(e, "ZK shutdown failed");
+			}
+
+			int n = 0;
+			while (n++ < 100) {
+				try {
+					ZkFourLetterWords.sendStat("127.0.0.1", port, 3000);
+					Thread.sleep(100);
+				}
+				catch (@SuppressWarnings("unused") Exception e) {
+					break;
+				}
+			}
+			if (n == 100) {
+				logger.debug("Zookeeper failed to stop");
+			}
+
+			try {
+				this.zookeeper.getZKDatabase().close();
+			}
+			catch (Exception e) {
+				logger.error(e, "ZK db close failed");
+			}
+
+			Utils.delete(this.logDir);
+			Utils.delete(this.snapshotDir);
+		}
+
+	}
+
 }

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

@@ -115,9 +115,9 @@ private boolean springTestContext(AnnotatedElement annotatedElement) {
 	@SuppressWarnings("unchecked")
 	private EmbeddedKafkaBroker createBroker(EmbeddedKafka embedded) {
 		EmbeddedKafkaBroker broker;
-		broker = new EmbeddedKafkaBroker(embedded.count(),
-				embedded.controlledShutdown(), embedded.topics());
-		broker.kafkaPorts(embedded.ports());
+		broker = new EmbeddedKafkaBroker(embedded.count(), embedded.controlledShutdown(), embedded.topics())
+				.zkPort(embedded.zookeeperPort())
+				.kafkaPorts(embedded.ports());
 		Properties properties = new Properties();
 
 		for (String pair : embedded.brokerProperties()) {

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

@@ -95,6 +95,13 @@
 	 */
 	int[] ports() default {0};
 
+	/**
+	 * Set the port on which the embedded Zookeeper should listen;
+	 * @return the port.
+	 * @since 2.3
+	 */
+	int zookeeperPort() default 0;
+
 	/**
 	 * @return partitions per topic
 	 */

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

@@ -67,11 +67,11 @@ public void customizeContext(ConfigurableApplicationContext context, MergedConte
 						.toArray(String[]::new);
 
 		EmbeddedKafkaBroker embeddedKafkaBroker = new EmbeddedKafkaBroker(this.embeddedKafka.count(),
-				this.embeddedKafka.controlledShutdown(),
-				this.embeddedKafka.partitions(),
-				topics);
-
-		embeddedKafkaBroker.kafkaPorts(this.embeddedKafka.ports());
+					this.embeddedKafka.controlledShutdown(),
+					this.embeddedKafka.partitions(),
+					topics)
+				.kafkaPorts(this.embeddedKafka.ports())
+				.zkPort(this.embeddedKafka.zookeeperPort());
 
 		Properties properties = new Properties();
 

spring-kafka-test/src/main/java/org/springframework/kafka/test/rule/EmbeddedKafkaRule.java

@@ -96,6 +96,11 @@ public EmbeddedKafkaRule kafkaPorts(int... kafkaPorts) {
 		return this;
 	}
 
+	public EmbeddedKafkaRule zkPort(int port) {
+		this.embeddedKafka.setZkPort(port);
+		return this;
+	}
+
 	/**
 	 * Return an underlying delegator {@link EmbeddedKafkaBroker} instance.
 	 * @return the {@link EmbeddedKafkaBroker} instance.

spring-kafka-test/src/test/java/org/springframework/kafka/test/EmbeddedKafkaBrokerTests.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2019 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.test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import org.junit.jupiter.api.Test;
+
+/**
+ * @author Gary Russell
+ * @since 2.3
+ *
+ */
+public class EmbeddedKafkaBrokerTests {
+
+	@Test
+	void testUpDown() {
+		EmbeddedKafkaBroker kafka = new EmbeddedKafkaBroker(1);
+		kafka.afterPropertiesSet();
+		assertThat(kafka.getZookeeperConnectionString()).startsWith("127");
+		kafka.destroy();
+		assertThat(kafka.getZookeeperConnectionString()).isNull();
+	}
+
+}

spring-kafka-test/src/test/java/org/springframework/kafka/test/rule/AddressableEmbeddedBrokerTests.java

@@ -61,15 +61,16 @@ public class AddressableEmbeddedBrokerTests {
 
 	@Test
 	public void testKafkaEmbedded() {
-		assertThat(broker.getBrokersAsString()).isEqualTo("127.0.0.1:" + this.config.port);
+		assertThat(broker.getBrokersAsString()).isEqualTo("127.0.0.1:" + this.config.kafkaPort);
+		assertThat(broker.getZkPort()).isEqualTo(this.config.zkPort);
 		assertThat(broker.getBrokersAsString())
 				.isEqualTo(System.getProperty(EmbeddedKafkaBroker.SPRING_EMBEDDED_KAFKA_BROKERS));
 		assertThat(broker.getZookeeperConnectionString())
 				.isEqualTo(System.getProperty(EmbeddedKafkaBroker.SPRING_EMBEDDED_ZOOKEEPER_CONNECT));
 	}
 
 	@Test
-	public void testLateStartedConsumer() throws Exception {
+	public void testLateStartedConsumer() {
 		Map<String, Object> consumerProps = KafkaTestUtils.consumerProps(TEST_EMBEDDED, "false", this.broker);
 		consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
 		Consumer<Integer, String> consumer = new KafkaConsumer<>(consumerProps);
@@ -93,16 +94,22 @@ public void testLateStartedConsumer() throws Exception {
 	@Configuration
 	public static class Config {
 
-		private int port;
+		private int kafkaPort;
+
+		private int zkPort;
 
 		@Bean
 		public EmbeddedKafkaBroker broker() throws IOException {
 			ServerSocket ss = ServerSocketFactory.getDefault().createServerSocket(0);
-			this.port = ss.getLocalPort();
+			this.kafkaPort = ss.getLocalPort();
+			ss.close();
+			ss = ServerSocketFactory.getDefault().createServerSocket(0);
+			this.zkPort = ss.getLocalPort();
 			ss.close();
 
 			return new EmbeddedKafkaBroker(1, true, TEST_EMBEDDED)
-					.kafkaPorts(this.port);
+					.zkPort(this.zkPort)
+					.kafkaPorts(this.kafkaPort);
 		}
 
 	}