KAFKA-2527; System Test for Quotas in Ducktape

tests/kafkatest/services/console_consumer.py

@@ -15,6 +15,7 @@
 
 from ducktape.services.background_thread import BackgroundThreadService
 from ducktape.utils.util import wait_until
+from kafkatest.services.performance.jmx_mixin import JmxMixin
 
 import os
 import subprocess
@@ -71,7 +72,7 @@ def is_int(msg):
 """
 
 
-class ConsoleConsumer(BackgroundThreadService):
+class ConsoleConsumer(JmxMixin):
     # Root directory for persistent output
     PERSISTENT_ROOT = "/mnt/console_consumer"
     STDOUT_CAPTURE = os.path.join(PERSISTENT_ROOT, "console_consumer.stdout")
@@ -93,7 +94,8 @@ class ConsoleConsumer(BackgroundThreadService):
             "collect_default": True}
         }
 
-    def __init__(self, context, num_nodes, kafka, topic, message_validator=None, from_beginning=True, consumer_timeout_ms=None):
+    def __init__(self, context, num_nodes, kafka, topic, message_validator=None, from_beginning=True,
+                 client_id="console-consumer", consumer_timeout_ms=None, jmx_object_name=None, jmx_attributes=None):
         """
         Args:
             context:                    standard context
@@ -107,7 +109,7 @@ def __init__(self, context, num_nodes, kafka, topic, message_validator=None, fro
                                         waiting for the consumer to stop is a pretty good way to consume all messages
                                         in a topic.
         """
-        super(ConsoleConsumer, self).__init__(context, num_nodes)
+        super(ConsoleConsumer, self).__init__(context, num_nodes, jmx_object_name, jmx_attributes)
         self.kafka = kafka
         self.args = {
             'topic': topic,
@@ -118,6 +120,7 @@ def __init__(self, context, num_nodes, kafka, topic, message_validator=None, fro
         self.from_beginning = from_beginning
         self.message_validator = message_validator
         self.messages_consumed = {idx: [] for idx in range(1, num_nodes + 1)}
+        self.client_id = client_id
 
     @property
     def start_cmd(self):
@@ -126,10 +129,11 @@ def start_cmd(self):
         args['stdout'] = ConsoleConsumer.STDOUT_CAPTURE
         args['stderr'] = ConsoleConsumer.STDERR_CAPTURE
         args['config_file'] = ConsoleConsumer.CONFIG_FILE
+        args['jmx_port'] = self.jmx_port
 
         cmd = "export LOG_DIR=%s;" % ConsoleConsumer.LOG_DIR
         cmd += " export KAFKA_LOG4J_OPTS=\"-Dlog4j.configuration=file:%s\";" % ConsoleConsumer.LOG4J_CONFIG
-        cmd += " /opt/kafka/bin/kafka-console-consumer.sh --topic %(topic)s --zookeeper %(zk_connect)s" \
+        cmd += " JMX_PORT=%(jmx_port)d /opt/kafka/bin/kafka-console-consumer.sh --topic %(topic)s --zookeeper %(zk_connect)s" \
             " --consumer.config %(config_file)s" % args
 
         if self.from_beginning:
@@ -153,10 +157,7 @@ def _worker(self, idx, node):
         node.account.ssh("mkdir -p %s" % ConsoleConsumer.PERSISTENT_ROOT, allow_fail=False)
 
         # Create and upload config file
-        if self.consumer_timeout_ms is not None:
-            prop_file = self.render('console_consumer.properties', consumer_timeout_ms=self.consumer_timeout_ms)
-        else:
-            prop_file = self.render('console_consumer.properties')
+        prop_file = self.render('console_consumer.properties', consumer_timeout_ms=self.consumer_timeout_ms, client_id=self.client_id)
 
         self.logger.info("console_consumer.properties:")
         self.logger.info(prop_file)
@@ -170,24 +171,26 @@ def _worker(self, idx, node):
         cmd = self.start_cmd
         self.logger.debug("Console consumer %d command: %s", idx, cmd)
         for line in node.account.ssh_capture(cmd, allow_fail=False):
+            self.maybe_start_jmx_tool(idx, node)
             msg = line.strip()
             if self.message_validator is not None:
                 msg = self.message_validator(msg)
             if msg is not None:
                 self.messages_consumed[idx].append(msg)
 
+        self.read_jmx_output(idx, node)
+
     def start_node(self, node):
         super(ConsoleConsumer, self).start_node(node)
 
     def stop_node(self, node):
-        node.account.kill_process("java", allow_fail=True)
+        node.account.kill_process("console-consumer", allow_fail=True)
         wait_until(lambda: not self.alive(node), timeout_sec=10, backoff_sec=.2,
                    err_msg="Timed out waiting for consumer to stop.")
 
     def clean_node(self, node):
         if self.alive(node):
             self.logger.warn("%s %s was still alive at cleanup time. Killing forcefully..." %
                              (self.__class__.__name__, node.account))
-        node.account.kill_process("java", clean_shutdown=False, allow_fail=True)
+        super(JmxMixin, self).clean_node(node)
         node.account.ssh("rm -rf %s" % ConsoleConsumer.PERSISTENT_ROOT, allow_fail=False)
-

tests/kafkatest/services/kafka.py

@@ -33,7 +33,7 @@ class KafkaService(Service):
             "collect_default": False}
     }
 
-    def __init__(self, context, num_nodes, zk, topics=None):
+    def __init__(self, context, num_nodes, zk, topics=None, quota_config=None):
         """
         :type context
         :type zk: ZookeeperService
@@ -42,6 +42,7 @@ def __init__(self, context, num_nodes, zk, topics=None):
         super(KafkaService, self).__init__(context, num_nodes)
         self.zk = zk
         self.topics = topics
+        self.quota_config = quota_config
 
     def start(self):
         super(KafkaService, self).start()
@@ -56,7 +57,8 @@ def start(self):
                 self.create_topic(topic_cfg)
 
     def start_node(self, node):
-        props_file = self.render('kafka.properties', node=node, broker_id=self.idx(node))
+        props_file = self.render('kafka.properties', node=node, broker_id=self.idx(node), quota_config=self.quota_config)
+
         self.logger.info("kafka.properties:")
         self.logger.info(props_file)
         node.account.create_file("/mnt/kafka.properties", props_file)

tests/kafkatest/services/performance/jmx_mixin.py

@@ -0,0 +1,78 @@
+# 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.
+
+from kafkatest.services.performance import PerformanceService
+
+
+class JmxMixin(PerformanceService):
+
+    def __init__(self, context, num_nodes, jmx_object_name=None, jmx_attributes=None):
+        super(JmxMixin, self).__init__(context, num_nodes)
+        self.jmx_object_name = jmx_object_name
+        self.jmx_attributes = jmx_attributes
+        self.jmx_port = 9192
+
+        self.started = [False] * self.num_nodes
+        self.jmx_stats = [{} for x in range(self.num_nodes)]
+        self.maximum_jmx_value = None
+        self.average_jmx_value = None
+
+    def clean_node(self, node):
+        super(JmxMixin, self).clean_node(node)
+        node.account.kill_process("jmx", clean_shutdown=False, allow_fail=True)
+        node.account.ssh("rm -rf /mnt/jmx_tool.log", allow_fail=False)
+
+    def maybe_start_jmx_tool(self, idx, node):
+        if self.started[idx-1] == True or self.jmx_object_name == None:
+            return
+        self.started[idx-1] = True
+
+        cmd = "/opt/kafka/bin/kafka-run-class.sh kafka.tools.JmxTool " \
+              "--reporting-interval 1000 --jmx-url service:jmx:rmi:///jndi/rmi://127.0.0.1:%d/jmxrmi " \
+              "--object-name %s " % (self.jmx_port, self.jmx_object_name)
+        if self.jmx_attributes != None:
+            cmd += "--attributes %s " % self.jmx_attributes
+        cmd += "> /mnt/jmx_tool.log &"
+
+        self.logger.debug("Start JmxTool %d command: %s", idx, cmd)
+        for l in node.account.ssh_capture(cmd, allow_fail=False):
+            break
+
+    def read_jmx_output(self, idx, node):
+        if self.started[idx-1] == False:
+            return
+
+        cmd = "grep -v time /mnt/jmx_tool.log"
+        self.logger.debug("Read jmx output %d command: %s", idx, cmd)
+
+        for line in node.account.ssh_capture(cmd, allow_fail=False):
+            time_sec = int(line.split(',')[0])/1000
+            value = float(line.split(',')[1])
+            self.jmx_stats[idx-1][time_sec] = value
+
+        if any(len(dict)==0 for dict in self.jmx_stats):
+            return
+
+        # since we have read jmx stats from all nodes, calculate average and maximum of jmx stats
+        start_time_sec = min([min(dict.keys()) for dict in self.jmx_stats])
+        end_time_sec = max([max(dict.keys()) for dict in self.jmx_stats])
+        values = []
+
+        for time_sec in xrange(start_time_sec, end_time_sec+1):
+            collection = [dict.get(time_sec, 0) for dict in self.jmx_stats]
+            values.append(sum(collection))
+
+        self.average_jmx_value = sum(values)/len(values)
+        self.maximum_jmx_value = max(values)

tests/kafkatest/services/performance/producer_performance.py

@@ -13,19 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from kafkatest.services.performance import PerformanceService
+from kafkatest.services.performance.jmx_mixin import JmxMixin
 
-
-class ProducerPerformanceService(PerformanceService):
+class ProducerPerformanceService(JmxMixin):
 
     logs = {
         "producer_performance_log": {
             "path": "/mnt/producer-performance.log",
             "collect_default": True},
     }
 
-    def __init__(self, context, num_nodes, kafka, topic, num_records, record_size, throughput, settings={}, intermediate_stats=False):
-        super(ProducerPerformanceService, self).__init__(context, num_nodes)
+    def __init__(self, context, num_nodes, kafka, topic, num_records, record_size, throughput, settings={},
+                 intermediate_stats=False, client_id="producer-performance", jmx_object_name=None, jmx_attributes=None):
+        super(ProducerPerformanceService, self).__init__(context, num_nodes, jmx_object_name, jmx_attributes)
         self.kafka = kafka
         self.args = {
             'topic': topic,
@@ -35,12 +35,13 @@ def __init__(self, context, num_nodes, kafka, topic, num_records, record_size, t
         }
         self.settings = settings
         self.intermediate_stats = intermediate_stats
+        self.client_id = client_id
 
     def _worker(self, idx, node):
         args = self.args.copy()
-        args.update({'bootstrap_servers': self.kafka.bootstrap_servers()})
-        cmd = "/opt/kafka/bin/kafka-run-class.sh org.apache.kafka.clients.tools.ProducerPerformance "\
-              "%(topic)s %(num_records)d %(record_size)d %(throughput)d bootstrap.servers=%(bootstrap_servers)s"\
+        args.update({'bootstrap_servers': self.kafka.bootstrap_servers(), 'jmx_port': self.jmx_port, 'client_id': self.client_id})
+        cmd = "JMX_PORT=%(jmx_port)d /opt/kafka/bin/kafka-run-class.sh org.apache.kafka.clients.tools.ProducerPerformance "\
+              "%(topic)s %(num_records)d %(record_size)d %(throughput)d bootstrap.servers=%(bootstrap_servers)s client.id=%(client_id)s"\
               " | tee /mnt/producer-performance.log" % args
 
         for key, value in self.settings.items():
@@ -62,6 +63,7 @@ def parse_stats(line):
             }
         last = None
         for line in node.account.ssh_capture(cmd):
+            self.maybe_start_jmx_tool(idx, node)
             if self.intermediate_stats:
                 try:
                     self.stats[idx-1].append(parse_stats(line))
@@ -74,3 +76,4 @@ def parse_stats(line):
             self.results[idx-1] = parse_stats(last)
         except:
             raise Exception("Unable to parse aggregate performance statistics on node %d: %s" % (idx, last))
+        self.read_jmx_output(idx, node)

tests/kafkatest/services/templates/console_consumer.properties

@@ -17,3 +17,10 @@
 {% if consumer_timeout_ms is defined and consumer_timeout_ms is not none %}
 consumer.timeout.ms={{ consumer_timeout_ms }}
 {% endif %}
+
+group.id={{ group_id|default('test-consumer-group') }}
+
+{% if client_id is defined and client_id is not none %}
+client.id={{ client_id }}
+{% endif %}
+

tests/kafkatest/services/templates/kafka.properties

@@ -39,3 +39,19 @@ log.cleaner.enable=false
 
 zookeeper.connect={{ zk.connect_setting() }}
 zookeeper.connection.timeout.ms=2000
+
+{% if quota_config.quota_producer_default is defined and quota_config.quota_producer_default is not none %}
+quota.producer.default={{ quota_config.quota_producer_default }}
+{% endif %}
+
+{% if quota_config.quota_consumer_default is defined and quota_config.quota_consumer_default is not none %}
+quota.consumer.default={{ quota_config.quota_consumer_default }}
+{% endif %}
+
+{% if quota_config.quota_producer_bytes_per_second_overrides is defined and quota_config.quota_producer_bytes_per_second_overrides is not none %}
+quota.producer.bytes.per.second.overrides={{ quota_config.quota_producer_bytes_per_second_overrides }}
+{% endif %}
+
+{% if quota_config.quota_consumer_bytes_per_second_overrides is defined and quota_config.quota_consumer_bytes_per_second_overrides is not none %}
+quota.consumer.bytes.per.second.overrides={{ quota_config.quota_consumer_bytes_per_second_overrides }}
+{% endif %}