Kafka 2257

clients/src/main/java/org/apache/kafka/clients/tools/VerifiableProducer.java

@@ -27,11 +27,15 @@
 import org.json.simple.JSONObject;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Properties;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import static net.sourceforge.argparse4j.impl.Arguments.store;
 
 import net.sourceforge.argparse4j.ArgumentParsers;
+import net.sourceforge.argparse4j.impl.Arguments;
 import net.sourceforge.argparse4j.inf.ArgumentParser;
 import net.sourceforge.argparse4j.inf.ArgumentParserException;
 import net.sourceforge.argparse4j.inf.Namespace;
@@ -66,14 +70,18 @@ public class VerifiableProducer {
     private long throughput;
 
     // Hook to trigger producing thread to stop sending messages
-    private boolean stopProducing = false;
+    private AtomicBoolean stopProducing = new AtomicBoolean(false);
+
+    // Timeout on producer.close() call
+    private long closeTimeoutSeconds;
 
     public VerifiableProducer(
-            Properties producerProps, String topic, int throughput, int maxMessages) {
+            Properties producerProps, String topic, int throughput, int maxMessages, long closeTimeoutSeconds) {
 
         this.topic = topic;
         this.throughput = throughput;
         this.maxMessages = maxMessages;
+        this.closeTimeoutSeconds = closeTimeoutSeconds;
         this.producer = new KafkaProducer<String, String>(producerProps);
     }
 
@@ -124,6 +132,23 @@ private static ArgumentParser argParser() {
                 .choices(0, 1, -1)
                 .metavar("ACKS")
                 .help("Acks required on each produced message. See Kafka docs on request.required.acks for details.");
+
+        parser.addArgument("--producer-prop")
+                .action(Arguments.append())
+                .required(false)
+                .type(String.class)
+                .dest("userSpecifiedProducerProps")
+                .metavar("PRODUCER-PROP")
+                .help("Any custom producer properties of the form key=val.");
+
+        parser.addArgument("--close-timeout")
+                .action(store())
+                .required(false)
+                .setDefault(10L)
+                .type(Long.class)
+                .metavar("CLOSE-TIMEOUT")
+                .dest("closeTimeoutSeconds")
+                .help("When SIGTERM is caught, wait at most this many seconds for unsent messages to flush before stopping the VerifiableProducer process.");
 
         return parser;
     }
@@ -134,13 +159,12 @@ public static VerifiableProducer createFromArgs(String[] args) {
         VerifiableProducer producer = null;
 
         try {
-            Namespace res;
-            res = parser.parseArgs(args);
-
+            Namespace res = parser.parseArgs(args);
             int maxMessages = res.getInt("maxMessages");
             String topic = res.getString("topic");
             int throughput = res.getInt("throughput");
-
+            long closeTimeoutSeconds = res.getLong("closeTimeoutSeconds");
+
             Properties producerProps = new Properties();
             producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, res.getString("brokerList"));
             producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG,
@@ -150,8 +174,24 @@ public static VerifiableProducer createFromArgs(String[] args) {
             producerProps.put(ProducerConfig.ACKS_CONFIG, Integer.toString(res.getInt("acks")));
             // No producer retries
             producerProps.put("retries", "0");
+
+            // Properties specified using the --producer-prop key=val format override
+            // other properties
+            List<String> userSpecifiedProducerProps = res.getList("userSpecifiedProducerProps");
+            for (String prop: userSpecifiedProducerProps) {
+                int index = prop.indexOf('=');
+                if (index < 0 || index == prop.length() - 1) {
+                    throw new IllegalArgumentException(
+                            "Invalid format for producer prop: " + prop);
+                }
+
+                String key = prop.substring(0, index);
+                String val = prop.substring(index + 1, prop.length());
+                producerProps.put(key, val);
+            }
 
-            producer = new VerifiableProducer(producerProps, topic, throughput, maxMessages);
+            producer = new VerifiableProducer(
+                    producerProps, topic, throughput, maxMessages, closeTimeoutSeconds);
         } catch (ArgumentParserException e) {
             if (args.length == 0) {
                 parser.printHelp();
@@ -181,7 +221,7 @@ public void send(String key, String value) {
 
     /** Close the producer to flush any remaining messages. */
     public void close() {
-        producer.close();
+        producer.close(this.closeTimeoutSeconds, TimeUnit.SECONDS);
     }
 
     /**
@@ -193,7 +233,7 @@ String errorString(Exception e, String key, String value, Long nowMs) {
 
         JSONObject obj = new JSONObject();
         obj.put("class", this.getClass().toString());
-        obj.put("name", "producer_send_error");
+        obj.put("name", "f_error");
 
         obj.put("time_ms", nowMs);
         obj.put("exception", e.getClass().toString());
@@ -233,6 +273,11 @@ private class PrintInfoCallback implements Callback {
 
         public void onCompletion(RecordMetadata recordMetadata, Exception e) {
             synchronized (System.out) {
+                if (VerifiableProducer.this.stopProducing.get()) {
+                    // Stop writing to stdout
+                    return;
+                }
+
                 if (e == null) {
                     VerifiableProducer.this.numAcked++;
                     System.out.println(successString(recordMetadata, this.key, this.value, System.currentTimeMillis()));
@@ -253,7 +298,7 @@ public static void main(String[] args) throws IOException {
             @Override
             public void run() {
                 // Trigger main thread to stop producing messages
-                producer.stopProducing = true;
+                producer.stopProducing.set(true);
 
                 // Flush any remaining messages
                 producer.close();
@@ -274,7 +319,7 @@ public void run() {
 
         ThroughputThrottler throttler = new ThroughputThrottler(producer.throughput, startMs);
         for (int i = 0; i < producer.maxMessages || infinite; i++) {
-            if (producer.stopProducing) {
+            if (producer.stopProducing.get()) {
                 break;
             }
             long sendStartMs = System.currentTimeMillis();

tests/kafkatest/services/console_consumer.py

@@ -91,6 +91,7 @@ def __init__(self, context, num_nodes, kafka, topic, message_validator=is_int, f
         """
         super(ConsoleConsumer, self).__init__(context, num_nodes)
         self.kafka = kafka
+        self.topic = topic
         self.args = {
             'topic': topic,
         }
@@ -132,7 +133,6 @@ def _worker(self, idx, node):
             msg = line.strip()
             msg = self.message_validator(msg)
             if msg is not None:
-                self.logger.debug("consumed a message: " + str(msg))
                 self.messages_consumed[idx].append(msg)
 
     def start_node(self, node):

tests/kafkatest/services/kafka.py

@@ -54,6 +54,14 @@ def start(self):
                 topic_cfg["topic"] = topic
                 self.create_topic(topic_cfg)
 
+        # Sanity check - are the broker processes actually running?
+        for node in self.nodes:
+            for pid in self.pids(node):
+                if not node.account.alive(pid):
+                    raise Exception(
+                        "Expected Kafka process %s to be running on node %s, but found that the process is not alive." %
+                        (str(pid), str(node.account)))
+
     def start_node(self, node):
         props_file = self.render('kafka.properties', node=node, broker_id=self.idx(node))
         self.logger.info("kafka.properties:")
@@ -96,7 +104,7 @@ def clean_node(self, node):
         node.account.ssh("rm -rf /mnt/kafka-logs /mnt/kafka.properties /mnt/kafka.log /mnt/kafka.pid", allow_fail=False)
 
     def create_topic(self, topic_cfg):
-        node = self.nodes[0] # any node is fine here
+        node = self.nodes[0]  # any node is fine here
         self.logger.info("Creating topic %s with settings %s", topic_cfg["topic"], topic_cfg)
 
         cmd = "/opt/kafka/bin/kafka-topics.sh --zookeeper %(zk_connect)s --create "\
@@ -116,6 +124,7 @@ def create_topic(self, topic_cfg):
 
         time.sleep(1)
         self.logger.info("Checking to see if topic was properly created...\n%s" % cmd)
+
         for line in self.describe_topic(topic_cfg["topic"]).split("\n"):
             self.logger.info(line)
 
@@ -199,29 +208,53 @@ def restart_node(self, node, wait_sec=0, clean_shutdown=True):
     def leader(self, topic, partition=0):
         """ Get the leader replica for the given topic and partition.
         """
-        cmd = "/opt/kafka/bin/kafka-run-class.sh kafka.tools.ZooKeeperMainWrapper -server %s " \
-              % self.zk.connect_setting()
-        cmd += "get /brokers/topics/%s/partitions/%d/state" % (topic, partition)
-        self.logger.debug(cmd)
+        topic_partition_data = self.zk_get_data("/brokers/topics/%s/partitions/%d/state" % (topic, partition))
+
+        if topic_partition_data is None:
+            raise Exception("Error finding data for topic %s and partition %d." % (topic, partition))
+        self.logger.info(topic_partition_data)
+
+        leader_idx = int(topic_partition_data["leader"])
+        self.logger.info("Leader for topic %s and partition %d is now: %d" % (topic, partition, leader_idx))
+        return self.get_node(leader_idx)
+
+    def controller(self):
+        """
+        Get the current controller
+        """
+        controller_data = self.zk_get_data("/controller")
+        if controller_data is None:
+            raise Exception("Error finding controller.")
+        self.logger.info(controller_data)
 
+        controller_idx = int(controller_data["brokerid"])
+        self.logger.info("Controller is: %d" % controller_idx)
+        controller_node = self.get_node(controller_idx)
+
+        if controller_node is None:
+            raise Exception("Found no node corresponding to the id: %d" % controller_idx)
+        return controller_node
+
+    def zk_get_data(self, path):
+        """
+        Get data from zookeeper on the given path. This method assumes the data is in JSON format.
+
+        :param path Zookeeper path to query for data
+        :param node optional node on which query will be run
+        """
+        cmd = "/opt/kafka/bin/kafka-run-class.sh kafka.tools.ZooKeeperMainWrapper -server %s get %s" % \
+              (self.zk.connect_setting(), path)
+        data = None
         node = self.nodes[0]
-        self.logger.debug("Querying zookeeper to find leader replica for topic %s: \n%s" % (cmd, topic))
-        partition_state = None
         for line in node.account.ssh_capture(cmd):
-            match = re.match("^({.+})$", line)
-            if match is not None:
-                partition_state = match.groups()[0]
+            try:
+                data = json.loads(line)
                 break
+            except ValueError:
+                pass
 
-        if partition_state is None:
-            raise Exception("Error finding partition state for topic %s and partition %d." % (topic, partition))
-
-        partition_state = json.loads(partition_state)
-        self.logger.info(partition_state)
-
-        leader_idx = int(partition_state["leader"])
-        self.logger.info("Leader for topic %s and partition %d is now: %d" % (topic, partition, leader_idx))
-        return self.get_node(leader_idx)
+        return data
 
     def bootstrap_servers(self):
-        return ','.join([node.account.hostname + ":9092" for node in self.nodes])
+        return ','.join([node.account.hostname + ":9092" for node in self.nodes])
+

tests/kafkatest/services/verifiable_producer.py

@@ -23,16 +23,18 @@ class VerifiableProducer(BackgroundThreadService):
     logs = {
         "producer_log": {
             "path": "/mnt/producer.log",
-            "collect_default": False}
+            "collect_default": True}
     }
 
-    def __init__(self, context, num_nodes, kafka, topic, max_messages=-1, throughput=100000):
+    def __init__(self, context, num_nodes, kafka, topic, max_messages=-1, throughput=100000, close_timeout_sec=60, producer_props=None):
         super(VerifiableProducer, self).__init__(context, num_nodes)
 
         self.kafka = kafka
         self.topic = topic
         self.max_messages = max_messages
         self.throughput = throughput
+        self.close_timeout_sec = close_timeout_sec
+        self.producer_props = producer_props  # Overrides for Kafka producer settings
 
         self.acked_values = []
         self.not_acked_values = []
@@ -63,6 +65,11 @@ def start_cmd(self):
             cmd += " --max-messages %s" % str(self.max_messages)
         if self.throughput > 0:
             cmd += " --throughput %s" % str(self.throughput)
+        cmd += " --close-timeout %s" % str(self.close_timeout_sec)
+
+        if self.producer_props is not None:
+            for key in self.producer_props:
+                cmd += " --producer-prop %s=%s" % (key, str(self.producer_props[key]))
 
         cmd += " 2>> /mnt/producer.log | tee -a /mnt/producer.log &"
         return cmd