KAFKA-2527; System Test for Quotas in Ducktape

tests/kafkatest/services/console_consumer.py

@@ -15,11 +15,13 @@
 
 from ducktape.services.background_thread import BackgroundThreadService
 from ducktape.utils.util import wait_until
+from kafkatest.services.performance.jmx_mixin import JmxMixin
+from kafkatest.services.performance import PerformanceService
 from kafkatest.utils.security_config import SecurityConfig
 
 import os
 import subprocess
-
+import itertools
 
 def is_int(msg):
     """Default method used to check whether text pulled from console consumer is a message.
@@ -72,7 +74,7 @@ def is_int(msg):
 """
 
 
-class ConsoleConsumer(BackgroundThreadService):
+class ConsoleConsumer(JmxMixin, PerformanceService):
     # Root directory for persistent output
     PERSISTENT_ROOT = "/mnt/console_consumer"
     STDOUT_CAPTURE = os.path.join(PERSISTENT_ROOT, "console_consumer.stdout")
@@ -94,7 +96,8 @@ class ConsoleConsumer(BackgroundThreadService):
             "collect_default": True}
         }
 
-    def __init__(self, context, num_nodes, kafka, topic, security_protocol=None, new_consumer=None, message_validator=None, from_beginning=True, consumer_timeout_ms=None):
+    def __init__(self, context, num_nodes, kafka, topic, security_protocol=None, new_consumer=None, message_validator=None,
+                 from_beginning=True, consumer_timeout_ms=None, client_id="console-consumer", jmx_object_names=None, jmx_attributes=[]):
         """
         Args:
             context:                    standard context
@@ -110,7 +113,8 @@ def __init__(self, context, num_nodes, kafka, topic, security_protocol=None, new
                                         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)
+        JmxMixin.__init__(self, num_nodes, jmx_object_names, jmx_attributes)
+        PerformanceService.__init__(self, context, num_nodes)
         self.kafka = kafka
         self.new_consumer = new_consumer
         self.args = {
@@ -122,9 +126,10 @@ def __init__(self, context, num_nodes, kafka, topic, security_protocol=None, new
         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
 
         # Process client configuration
-        self.prop_file = self.render('console_consumer.properties', consumer_timeout_ms=self.consumer_timeout_ms)
+        self.prop_file = self.render('console_consumer.properties', consumer_timeout_ms=self.consumer_timeout_ms, client_id=self.client_id)
 
         # Add security properties to the config. If security protocol is not specified,
         # use the default in the template properties.
@@ -143,10 +148,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" \
+        cmd += " JMX_PORT=%(jmx_port)d /opt/kafka/bin/kafka-console-consumer.sh --topic %(topic)s" \
             " --consumer.config %(config_file)s" % args
 
         if self.new_consumer:
@@ -173,6 +179,7 @@ def alive(self, node):
     def _worker(self, idx, node):
         node.account.ssh("mkdir -p %s" % ConsoleConsumer.PERSISTENT_ROOT, allow_fail=False)
 
+        # Create and upload config file
         self.logger.info("console_consumer.properties:")
         self.logger.info(self.prop_file)
         node.account.create_file(ConsoleConsumer.CONFIG_FILE, self.prop_file)
@@ -185,26 +192,32 @@ def _worker(self, idx, node):
         # Run and capture output
         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):
+
+        consumer_output = node.account.ssh_capture(cmd, allow_fail=False)
+        first_line = consumer_output.next()
+        self.start_jmx_tool(idx, node)
+        for line in itertools.chain([first_line], consumer_output):
             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)
+        PerformanceService.start_node(self, 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)
+        JmxMixin.clean_node(self, node)
+        PerformanceService.clean_node(self, node)
         node.account.ssh("rm -rf %s" % ConsoleConsumer.PERSISTENT_ROOT, allow_fail=False)
         self.security_config.clean_node(node)
-

tests/kafkatest/services/kafka.py

@@ -15,15 +15,15 @@
 
 from ducktape.services.service import Service
 from ducktape.utils.util import wait_until
+from kafkatest.services.performance.jmx_mixin import JmxMixin
 from kafkatest.utils.security_config import SecurityConfig
-
 import json
 import re
 import signal
 import time
 
 
-class KafkaService(Service):
+class KafkaService(JmxMixin, Service):
 
     logs = {
         "kafka_log": {
@@ -34,13 +34,15 @@ class KafkaService(Service):
             "collect_default": False}
     }
 
-    def __init__(self, context, num_nodes, zk, security_protocol=SecurityConfig.PLAINTEXT, interbroker_security_protocol=SecurityConfig.PLAINTEXT, topics=None):
+    def __init__(self, context, num_nodes, zk, security_protocol=SecurityConfig.PLAINTEXT, interbroker_security_protocol=SecurityConfig.PLAINTEXT,
+                 topics=None, quota_config=None, jmx_object_names=None, jmx_attributes=[]):
         """
         :type context
         :type zk: ZookeeperService
         :type topics: dict
         """
-        super(KafkaService, self).__init__(context, num_nodes)
+        Service.__init__(self, context, num_nodes)
+        JmxMixin.__init__(self, num_nodes, jmx_object_names, jmx_attributes)
         self.zk = zk
         if security_protocol == SecurityConfig.SSL or interbroker_security_protocol == SecurityConfig.SSL:
             self.security_config = SecurityConfig(SecurityConfig.SSL)
@@ -50,9 +52,10 @@ def __init__(self, context, num_nodes, zk, security_protocol=SecurityConfig.PLAI
         self.interbroker_security_protocol = interbroker_security_protocol
         self.port = 9092 if security_protocol == SecurityConfig.PLAINTEXT else 9093
         self.topics = topics
+        self.quota_config = quota_config
 
     def start(self):
-        super(KafkaService, self).start()
+        Service.start(self)
 
         # Create topics if necessary
         if self.topics is not None:
@@ -65,18 +68,19 @@ def start(self):
 
     def start_node(self, node):
         props_file = self.render('kafka.properties', node=node, broker_id=self.idx(node),
-            port = self.port, security_protocol = self.security_protocol, 
+            port = self.port, security_protocol = self.security_protocol, quota_config=self.quota_config,
             interbroker_security_protocol=self.interbroker_security_protocol)
         self.logger.info("kafka.properties:")
         self.logger.info(props_file)
         node.account.create_file("/mnt/kafka.properties", props_file)
         self.security_config.setup_node(node)
 
-        cmd = "/opt/kafka/bin/kafka-server-start.sh /mnt/kafka.properties 1>> /mnt/kafka.log 2>> /mnt/kafka.log & echo $! > /mnt/kafka.pid"
+        cmd = "JMX_PORT=%d /opt/kafka/bin/kafka-server-start.sh /mnt/kafka.properties 1>> /mnt/kafka.log 2>> /mnt/kafka.log & echo $! > /mnt/kafka.pid" % self.jmx_port
         self.logger.debug("Attempting to start KafkaService on %s with command: %s" % (str(node.account), cmd))
         with node.account.monitor_log("/mnt/kafka.log") as monitor:
             node.account.ssh(cmd)
             monitor.wait_until("Kafka Server.*started", timeout_sec=30, err_msg="Kafka server didn't finish startup")
+        self.start_jmx_tool(self.idx(node), node)
         if len(self.pids(node)) == 0:
             raise Exception("No process ids recorded on node %s" % str(node))
 
@@ -106,6 +110,7 @@ def stop_node(self, node, clean_shutdown=True):
         node.account.ssh("rm -f /mnt/kafka.pid", allow_fail=False)
 
     def clean_node(self, node):
+        JmxMixin.clean_node(self, node)
         node.account.kill_process("kafka", clean_shutdown=False, allow_fail=True)
         node.account.ssh("rm -rf /mnt/kafka-logs /mnt/kafka.properties /mnt/kafka.log /mnt/kafka.pid", allow_fail=False)
         self.security_config.clean_node(node)
@@ -242,3 +247,7 @@ def bootstrap_servers(self):
         """Get the broker list to connect to Kafka using the specified security protocol
         """
         return ','.join([node.account.hostname + ":" + `self.port` for node in self.nodes])
+
+    def read_jmx_output_all_nodes(self):
+        for node in self.nodes:
+            self.read_jmx_output(self.idx(node), node)

tests/kafkatest/services/performance/jmx_mixin.py

@@ -0,0 +1,81 @@
+# 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.
+
+class JmxMixin(object):
+
+    def __init__(self, num_nodes, jmx_object_names=None, jmx_attributes=[]):
+        self.jmx_object_names = jmx_object_names
+        self.jmx_attributes = jmx_attributes
+        self.jmx_port = 9192
+
+        self.started = [False] * num_nodes
+        self.jmx_stats = [{} for x in range(num_nodes)]
+        self.maximum_jmx_value = {}  # map from object_attribute_name to maximum value observed over time
+        self.average_jmx_value = {}  # map from object_attribute_name to average value observed over time
+
+    def clean_node(self, 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 start_jmx_tool(self, idx, node):
+        if self.started[idx-1] == True or self.jmx_object_names == 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" % self.jmx_port
+        for jmx_object_name in self.jmx_object_names:
+            cmd += " --object-name %s" % jmx_object_name
+        for jmx_attribute in self.jmx_attributes:
+            cmd += " --attributes %s" % jmx_attribute
+        cmd += " | tee -a /mnt/jmx_tool.log"
+
+        self.logger.debug("Start JmxTool %d command: %s", idx, cmd)
+        jmx_output = node.account.ssh_capture(cmd, allow_fail=False)
+        jmx_output.next()
+
+    def read_jmx_output(self, idx, node):
+        if self.started[idx-1] == False:
+            return
+        self.maximum_jmx_value = {}
+        self.average_jmx_value = {}
+        object_attribute_names = []
+
+        cmd = "cat /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):
+            if "time" in line:
+                object_attribute_names = line.strip()[1:-1].split("\",\"")[1:]
+                continue
+            stats = [float(field) for field in line.split(',')]
+            time_sec = int(stats[0]/1000)
+            self.jmx_stats[idx-1][time_sec] = {name : stats[i+1] for i, name in enumerate(object_attribute_names)}
+
+        # do not calculate average and maximum of jmx stats until we have read output from all nodes
+        if any(len(time_to_stats)==0 for time_to_stats in self.jmx_stats):
+            return
+
+        start_time_sec = min([min(time_to_stats.keys()) for time_to_stats in self.jmx_stats])
+        end_time_sec = max([max(time_to_stats.keys()) for time_to_stats in self.jmx_stats])
+
+        for name in object_attribute_names:
+            aggregates_per_time = []
+            for time_sec in xrange(start_time_sec, end_time_sec+1):
+                # assume that value is 0 if it is not read by jmx tool at the given time. This is appropriate for metrics such as bandwidth
+                values_per_node = [time_to_stats.get(time_sec, {}).get(name, 0) for time_to_stats in self.jmx_stats]
+                # assume that value is aggregated across nodes by sum. This is appropriate for metrics such as bandwidth
+                aggregates_per_time.append(sum(values_per_node))
+            self.average_jmx_value[name] = sum(aggregates_per_time)/len(aggregates_per_time)
+            self.maximum_jmx_value[name] = max(aggregates_per_time)

tests/kafkatest/services/performance/producer_performance.py

@@ -13,20 +13,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from kafkatest.services.performance.jmx_mixin import JmxMixin
 from kafkatest.services.performance import PerformanceService
+import itertools
 from kafkatest.utils.security_config import SecurityConfig
 
-
-class ProducerPerformanceService(PerformanceService):
+class ProducerPerformanceService(JmxMixin, PerformanceService):
 
     logs = {
         "producer_performance_log": {
             "path": "/mnt/producer-performance.log",
             "collect_default": True},
     }
 
-    def __init__(self, context, num_nodes, kafka, security_protocol, topic, num_records, record_size, throughput, settings={}, intermediate_stats=False):
-        super(ProducerPerformanceService, self).__init__(context, num_nodes)
+    def __init__(self, context, num_nodes, kafka, security_protocol, topic, num_records, record_size, throughput, settings={},
+                 intermediate_stats=False, client_id="producer-performance", jmx_object_names=None, jmx_attributes=[]):
+        JmxMixin.__init__(self, num_nodes, jmx_object_names, jmx_attributes)
+        PerformanceService.__init__(self, context, num_nodes)
         self.kafka = kafka
         self.security_config = SecurityConfig(security_protocol)
         self.security_protocol = security_protocol
@@ -38,12 +41,13 @@ def __init__(self, context, num_nodes, kafka, security_protocol, topic, num_reco
         }
         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
+        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" % args
 
         self.security_config.setup_node(node)
         if self.security_protocol == SecurityConfig.SSL:
@@ -68,7 +72,10 @@ def parse_stats(line):
                 'latency_999th_ms': float(parts[7].split()[0]),
             }
         last = None
-        for line in node.account.ssh_capture(cmd):
+        producer_output = node.account.ssh_capture(cmd)
+        first_line = producer_output.next()
+        self.start_jmx_tool(idx, node)
+        for line in itertools.chain([first_line], producer_output):
             if self.intermediate_stats:
                 try:
                     self.stats[idx-1].append(parse_stats(line))
@@ -81,3 +88,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

@@ -47,6 +47,22 @@ 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 %}
+
 security.inter.broker.protocol={{ interbroker_security_protocol }}
 ssl.keystore.location=/mnt/ssl/test.keystore.jks
 ssl.keystore.password=test-ks-passwd