/*
* Author: cbedford
* Date: 10/20/13
* Time: 8:54 PM
*/
import com.google.common.collect.ImmutableMap;
import com.google.common.io.Files;
import com.netflix.curator.test.TestingServer;
import kafka.consumer.ConsumerConfig;
import kafka.consumer.ConsumerIterator;
import kafka.consumer.KafkaStream;
import kafka.javaapi.consumer.ConsumerConnector;
import kafka.javaapi.producer.Producer;
import kafka.producer.KeyedMessage;
import kafka.producer.ProducerConfig;
import kafka.serializer.StringDecoder;
import kafka.server.KafkaConfig;
import kafka.server.KafkaServer;
import kafka.utils.MockTime;
import kafka.utils.VerifiableProperties;
import org.apache.commons.lang.StringUtils;
import java.io.File;
import java.io.IOException;
import java.util.*;
class TestKafkaProducer {
private String topic = "";
private String zkConnectString = "";
private List<String> messages = null;
private List<String> messagesReceived = new ArrayList<String>();
private Producer<String, String> producer;
private KafkaServer kafkaServer;
private Thread kafkaMessageReceiverThread;
private static final String RANDOM_GROUP_ID = "RANDOM-GROUP-ID";
public static void main(String[] args) {
TestKafkaProducer tkp = null;
boolean success = false;
try (TestingServer zookeeperTestServer = new TestingServer()) {
final String theTopic = "someTopic-" + new Random().nextInt();
tkp = new TestKafkaProducer(
theTopic,
"localhost:" + zookeeperTestServer.getPort(),
4400);
tkp.sendMessages();
tkp.consumeMessages();
try { // Give consumer some time...
tkp.shutdownConsumers();
Thread.sleep(1000);
tkp.kafkaMessageReceiverThread.join();
tkp.shutdown();
} catch (Exception e) {
System.out.println("Error in shut down. we will ignore it as long as our messages came through");
e.printStackTrace();
}
String got = StringUtils.join(tkp.messagesReceived, "+");
String expected = StringUtils.join(tkp.messages, "+");
if (got.equals(expected)) {
success = true;
}
} catch (Exception e) {
e.printStackTrace();
}
if (! success) {
throw new RuntimeException("oh rats... we failed");
}
System.out.println("SUCCESS -- WE ARE HAPPY !...");
}
private void consumeMessages() {
final ConsumerConnector consumer =
kafka.consumer.Consumer.createJavaConsumerConnector(createConsumerConfig());
final Map<String, Integer> topicCountMap =
ImmutableMap.of(topic, 1);
final StringDecoder decoder =
new StringDecoder(new VerifiableProperties());
final Map<String, List<KafkaStream<String,String>>> consumerMap =
consumer.createMessageStreams(topicCountMap, decoder, decoder);
final KafkaStream<String,String> stream =
consumerMap.get(topic).get(0);
final ConsumerIterator<String,String> iterator = stream.iterator();
kafkaMessageReceiverThread = new Thread(
new Runnable() {
@Override
public void run() {
while (iterator.hasNext()) {
String msg = iterator.next().message();
msg = msg == null ? "<null>" : msg;
System.out.println("got message" + msg);
if (msg.equals("SHUTDOWN")) {
consumer.shutdown();
return;
}
messagesReceived.add(msg);
}
}
},
"kafkaMessageReceiverThread"
);
kafkaMessageReceiverThread.start();
}
private ConsumerConfig createConsumerConfig() {
Properties props = new Properties();
props.put("zookeeper.connect", this.zkConnectString);
props.put("group.id", RANDOM_GROUP_ID);
props.put("zk.sessiontimeout.ms", "400");
props.put("zk.synctime.ms", "200");
props.put("autocommit.interval.ms", "1000");
props.put("serializer.class", "kafka.serializer.StringEncoder");
return new ConsumerConfig(props);
}
public void shutdownConsumers() {
sendMessage("SHUTDOWN");
}
public void shutdown() {
producer.close();
try { // Give producer some time...
Thread.sleep(1000);
} catch (InterruptedException e) {
e.printStackTrace();
}
kafkaServer.shutdown();
kafkaServer.awaitShutdown();
}
TestKafkaProducer(String topic, String zkConnectString, int numRandomMessages) throws IOException {
final Random generator = new Random();
if (numRandomMessages <= 0) {
throw new RuntimeException("no messages defined for test");
}
messages = new ArrayList<String>();
for (int i = 0; i < numRandomMessages; i++) {
int num1 = Math.abs(generator.nextInt());
int num2 = Math.abs(generator.nextInt());
String messageToSend = num1 + ":-(a)-" + num2;
messages.add(messageToSend);
}
this.topic = topic;
this.zkConnectString = zkConnectString;
initProducer(zkConnectString);
}
public void sendMessages() throws IOException {
for (String msg : messages) {
sendMessage(msg);
}
}
private void sendMessage(String msg) {
KeyedMessage<String, String> data = new KeyedMessage<String, String>(topic, msg);
producer.send(data);
}
private void initProducer(String zkConnectString) throws IOException {
kafkaServer = startKafkaServer();
Properties props = new Properties();
props.put("metadata.broker.list", "localhost:9092");
props.put("serializer.class", "kafka.serializer.StringEncoder");
props.put("producer.type", "async");
props.put("batch.size", "1");
ProducerConfig config = new ProducerConfig(props);
producer = new Producer<String, String>(config);
}
private KafkaServer startKafkaServer() {
File tmpDir = Files.createTempDir();
Properties props = createProperties(tmpDir.getAbsolutePath(), 9092, 1);
KafkaConfig kafkaConfig = new KafkaConfig(props);
kafkaServer = new KafkaServer(kafkaConfig, new MockTime());
kafkaServer.startup();
return kafkaServer;
}
private Properties createProperties(String logDir, int port, int brokerId) {
Properties properties = new Properties();
properties.put("port", port + "");
properties.put("broker.id", brokerId + "");
properties.put("log.dir", logDir);
properties.put("zookeeper.connect", this.zkConnectString);
return properties;
}
}