Kafka Consumer Quick Rebalance

/*
 * 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 java.io.File;
import java.io.IOException;
import java.util.Collections;
import java.util.Properties;

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 kafka.utils.ZkUtils;
import org.I0Itec.zkclient.ZkClient;
import org.apache.curator.test.InstanceSpec;
import org.apache.curator.test.TestingServer;


/**
 * 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);
		}
		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 String getZkConnectString() {
		return zkServer.getConnectString();
	}

	public void stop() throws IOException {
		kafkaServer.shutdown();
		zkServer.stop();
	}

    public static void main(String[] args) {
        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");
        ConsumerConfig config = new ConsumerConfig(props);
        // This works fine:
        // connector.createMessageStreams(Collections.singletonMap("mytopic", 3));

        // This does not:
        for (int i = 0 ; i < 3 ; i++) {
            System.out.format("%nCreating consumer #%d%n", i+1);
            ConsumerConnector connector = Consumer.createJavaConsumerConnector(config);
            connector.createMessageStreams(Collections.singletonMap("mytopic", 1));
        }

    }
}