/*
* Copyright 2014 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
*
* 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.springframework.xd.dirt.integration.kafka;
import kafka.admin.AdminUtils;
import kafka.consumer.Consumer;
import kafka.consumer.ConsumerConfig;
import kafka.javaapi.consumer.ConsumerConnector;
import kafka.server.KafkaConfig;
import kafka.server.KafkaServerStartable;
import kafka.utils.TestUtils;
import org.I0Itec.zkclient.ZkClient;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheEvent;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.apache.curator.retry.RetryUntilElapsed;
import org.apache.curator.test.InstanceSpec;
import org.apache.curator.test.TestingServer;
import org.springframework.util.Assert;
import java.io.IOException;
import java.util.Collections;
import java.util.Properties;
/**
* A test Kafka + ZooKeeper pair for testing purposes.
*
* @author Eric Bottard
*/
public class TestKafkaCluster {
private KafkaServerStartable kafkaServer;
private TestingServer zkServer;
public TestKafkaCluster() {
try {
// Use very short ticktime so that disconnects are picked up quickly
// and rebalances go through
// int tickTime = 100;
// InstanceSpec spec = new InstanceSpec(null, -1, -1, -1, true, -1, tickTime, -1);
// zkServer = new TestingServer(spec, true);
zkServer = new TestingServer();
} catch (Exception e) {
throw new IllegalStateException(e);
}
KafkaConfig config = getKafkaConfig(zkServer.getConnectString());
kafkaServer = new KafkaServerStartable(config);
kafkaServer.startup();
}
private static KafkaConfig getKafkaConfig(final String zkConnectString) {
scala.collection.Iterator<Properties> propsI =
TestUtils.createBrokerConfigs(1).iterator();
assert propsI.hasNext();
Properties props = propsI.next();
assert props.containsKey("zookeeper.connect");
props.put("zookeeper.connect", zkConnectString);
return new KafkaConfig(props);
}
public String getKafkaBrokerString() {
return String.format("localhost:%d",
kafkaServer.serverConfig().port());
}
public void stop() throws IOException {
kafkaServer.shutdown();
zkServer.stop();
}
/**
* See XD-2293. This is used to reproduce Kafka rebalance issues.
*/
public static void main(String[] args) throws Exception {
TestKafkaCluster cluster = new TestKafkaCluster();
ZkClient client = new ZkClient(cluster.getZkConnectString(), 10000, 10000, KafkaMessageBus.utf8Serializer);
int partitions = 5;
int replication = 1;
AdminUtils.createTopic(client, "mytopic", partitions, replication, new Properties());
Properties props = new Properties();
props.put("zookeeper.connect", cluster.getZkConnectString());
props.put("group.id", "foo");
props.put("rebalance.backoff.ms", "2000");
props.put("rebalance.max.retries", "2000");
ConsumerConfig config = new ConsumerConfig(props);
CuratorFramework curator = CuratorFrameworkFactory.newClient(cluster.getZkConnectString(), new RetryUntilElapsed(1000, 100));
curator.start();
RebalanceListener listener = null;
for (int i = 0; i < 5; i++) {
System.out.format("%nCreating consumer #%d%n", i + 1);
ConsumerConnector connector = Consumer.createJavaConsumerConnector(config);
connector.createMessageStreams(Collections.singletonMap("mytopic", 1));
if (i == 0) {
PathChildrenCache cache = new PathChildrenCache(curator, "/consumers/foo/owners/mytopic", true);
listener = new RebalanceListener(5);
cache.getListenable().addListener(listener);
cache.start(PathChildrenCache.StartMode.POST_INITIALIZED_EVENT);
}
synchronized (listener) {
System.out.println("******** Waiting for rebalance...");
listener.wait();
}
}
System.out.println();
}
public String getZkConnectString() {
return zkServer.getConnectString();
}
private static class RebalanceListener implements PathChildrenCacheListener {
private int expected;
private int actual;
private boolean ready;
public RebalanceListener(int expected) {
this.expected = expected;
}
@Override
public synchronized void childEvent(CuratorFramework client, PathChildrenCacheEvent event) throws Exception {
System.out.println(event);
System.out.println(event.getData() != null ? new String(event.getData().getData()) : "no data");
switch (event.getType()) {
case CHILD_ADDED:
actual++;
if (ready && actual == expected) {
System.out.println("*** Moving on... ");
this.notify();
}
break;
case CHILD_REMOVED:
actual--;
break;
case INITIALIZED:
Assert.isTrue(actual == expected);
ready = true;
this.notify();
break;
}
// System.out.println("" + actual + "/" + expected);
}
}
}