/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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
*
* http://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.apache.beam.runners.spark.translation.streaming.utils;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Properties;
import java.util.Random;
import kafka.server.KafkaConfig;
import kafka.server.KafkaServerStartable;
import org.apache.zookeeper.server.NIOServerCnxnFactory;
import org.apache.zookeeper.server.ServerCnxnFactory;
import org.apache.zookeeper.server.ZooKeeperServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Embedded Kafka cluster.
* https://gist.github.com/fjavieralba/7930018
*/
public class EmbeddedKafkaCluster {
private static final Logger LOG = LoggerFactory.getLogger(EmbeddedKafkaCluster.class);
private final List<Integer> ports;
private final String zkConnection;
private final Properties baseProperties;
private final String brokerList;
private final List<KafkaServerStartable> brokers;
private final List<File> logDirs;
public EmbeddedKafkaCluster(String zkConnection) {
this(zkConnection, new Properties());
}
public EmbeddedKafkaCluster(String zkConnection, Properties baseProperties) {
this(zkConnection, baseProperties, Collections.singletonList(-1));
}
public EmbeddedKafkaCluster(String zkConnection, Properties baseProperties, List<Integer> ports) {
this.zkConnection = zkConnection;
this.ports = resolvePorts(ports);
this.baseProperties = baseProperties;
this.brokers = new ArrayList<>();
this.logDirs = new ArrayList<>();
this.brokerList = constructBrokerList(this.ports);
}
private static List<Integer> resolvePorts(List<Integer> ports) {
List<Integer> resolvedPorts = new ArrayList<>();
for (Integer port : ports) {
resolvedPorts.add(resolvePort(port));
}
return resolvedPorts;
}
private static int resolvePort(int port) {
if (port == -1) {
return TestUtils.getAvailablePort();
}
return port;
}
private static String constructBrokerList(List<Integer> ports) {
StringBuilder sb = new StringBuilder();
for (Integer port : ports) {
if (sb.length() > 0) {
sb.append(",");
}
sb.append("127.0.0.1:").append(port);
}
return sb.toString();
}
public void startup() {
for (int i = 0; i < ports.size(); i++) {
Integer port = ports.get(i);
File logDir = TestUtils.constructTempDir("kafka-local");
Properties properties = new Properties();
properties.putAll(baseProperties);
properties.setProperty("zookeeper.connect", zkConnection);
properties.setProperty("broker.id", String.valueOf(i + 1));
properties.setProperty("advertised.host.name", "127.0.0.1");
properties.setProperty("host.name", "127.0.0.1");
properties.setProperty("advertised.port", Integer.toString(port));
properties.setProperty("port", Integer.toString(port));
properties.setProperty("log.dirs", logDir.getAbsolutePath());
properties.setProperty("offsets.topic.num.partitions", "1");
properties.setProperty("offsets.topic.replication.factor", "1");
properties.setProperty("log.flush.interval.messages", String.valueOf(1));
KafkaServerStartable broker = startBroker(properties);
brokers.add(broker);
logDirs.add(logDir);
}
}
private static KafkaServerStartable startBroker(Properties props) {
KafkaServerStartable server = new KafkaServerStartable(new KafkaConfig(props));
server.startup();
return server;
}
public Properties getProps() {
Properties props = new Properties();
props.putAll(baseProperties);
props.put("bootstrap.servers", brokerList);
return props;
}
public String getBrokerList() {
return brokerList;
}
public List<Integer> getPorts() {
return ports;
}
public String getZkConnection() {
return zkConnection;
}
public void shutdown() {
for (KafkaServerStartable broker : brokers) {
try {
broker.shutdown();
} catch (Exception e) {
LOG.warn("{}", e.getMessage(), e);
}
}
for (File logDir : logDirs) {
try {
TestUtils.deleteFile(logDir);
} catch (FileNotFoundException e) {
LOG.warn("{}", e.getMessage(), e);
}
}
}
@Override
public String toString() {
return "EmbeddedKafkaCluster{" + "brokerList='" + brokerList + "'}";
}
/**
* Embedded Zookeeper.
*/
public static class EmbeddedZookeeper {
private int port = -1;
private int tickTime = 500;
private ServerCnxnFactory factory;
private File snapshotDir;
private File logDir;
public EmbeddedZookeeper() {
this(-1);
}
public EmbeddedZookeeper(int port) {
this(port, 500);
}
public EmbeddedZookeeper(int port, int tickTime) {
this.port = resolvePort(port);
this.tickTime = tickTime;
}
private static int resolvePort(int port) {
if (port == -1) {
return TestUtils.getAvailablePort();
}
return port;
}
public void startup() throws IOException {
if (this.port == -1) {
this.port = TestUtils.getAvailablePort();
}
this.factory = NIOServerCnxnFactory.createFactory(new InetSocketAddress("127.0.0.1", port),
1024);
this.snapshotDir = TestUtils.constructTempDir("embedded-zk/snapshot");
this.logDir = TestUtils.constructTempDir("embedded-zk/log");
try {
factory.startup(new ZooKeeperServer(snapshotDir, logDir, tickTime));
} catch (InterruptedException e) {
throw new IOException(e);
}
}
public void shutdown() {
if (factory != null) {
factory.shutdown();
}
try {
TestUtils.deleteFile(snapshotDir);
} catch (FileNotFoundException e) {
// ignore
}
try {
TestUtils.deleteFile(logDir);
} catch (FileNotFoundException e) {
// ignore
}
}
public String getConnection() {
return "127.0.0.1:" + port;
}
public void setPort(int port) {
this.port = port;
}
public void setTickTime(int tickTime) {
this.tickTime = tickTime;
}
public int getPort() {
return port;
}
public int getTickTime() {
return tickTime;
}
@Override
public String toString() {
return "EmbeddedZookeeper{" + "connection=" + getConnection() + "}";
}
}
static final class TestUtils {
private static final Random RANDOM = new Random();
private TestUtils() {
}
static File constructTempDir(String dirPrefix) {
File file = new File(System.getProperty("java.io.tmpdir"), dirPrefix + RANDOM.nextInt
(10000000));
if (!file.mkdirs()) {
throw new RuntimeException("could not create temp directory: " + file.getAbsolutePath());
}
file.deleteOnExit();
return file;
}
static int getAvailablePort() {
try {
try (ServerSocket socket = new ServerSocket(0)) {
return socket.getLocalPort();
}
} catch (IOException e) {
throw new IllegalStateException("Cannot find available port: " + e.getMessage(), e);
}
}
static boolean deleteFile(File path) throws FileNotFoundException {
if (!path.exists()) {
throw new FileNotFoundException(path.getAbsolutePath());
}
boolean ret = true;
if (path.isDirectory()) {
for (File f : path.listFiles()) {
ret = ret && deleteFile(f);
}
}
return ret && path.delete();
}
}
}