| # 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. |
| |
| import json |
| import os |
| |
| from ducktape.services.background_thread import BackgroundThreadService |
| |
| from kafkatest.directory_layout.kafka_path import KafkaPathResolverMixin |
| from kafkatest.services.kafka import TopicPartition |
| from kafkatest.services.kafka.util import get_log4j_config_param, get_log4j_config_for_tools |
| from kafkatest.services.verifiable_client import VerifiableClientMixin |
| from kafkatest.version import DEV_BRANCH |
| |
| class ShareConsumerState: |
| Started = 1 |
| Dead = 2 |
| |
| class ShareConsumerEventHandler(object): |
| |
| def __init__(self, node, idx, state=ShareConsumerState.Dead): |
| self.node = node |
| self.idx = idx |
| self.total_consumed = 0 |
| self.total_acknowledged_successfully = 0 |
| self.total_acknowledged_failed = 0 |
| self.consumed_per_partition = {} |
| self.acknowledged_per_partition = {} |
| self.acknowledged_per_partition_failed = {} |
| self.state = state |
| |
| def handle_shutdown_complete(self, node=None, logger=None): |
| self.state = ShareConsumerState.Dead |
| if node is not None and logger is not None: |
| logger.debug("Shut down %s" % node.account.hostname) |
| |
| def handle_startup_complete(self, node, logger): |
| self.state = ShareConsumerState.Started |
| logger.debug("Started %s" % node.account.hostname) |
| |
| def handle_offsets_acknowledged(self, event, node, logger): |
| if event["success"]: |
| self.total_acknowledged_successfully += event["count"] |
| for share_partition_data in event["partitions"]: |
| topic_partition = TopicPartition(share_partition_data["topic"], share_partition_data["partition"]) |
| self.acknowledged_per_partition[topic_partition] = self.acknowledged_per_partition.get(topic_partition, 0) + share_partition_data["count"] |
| logger.debug("Offsets acknowledged for %s" % (node.account.hostname)) |
| else: |
| self.total_acknowledged_failed += event["count"] |
| for share_partition_data in event["partitions"]: |
| topic_partition = TopicPartition(share_partition_data["topic"], share_partition_data["partition"]) |
| self.acknowledged_per_partition_failed[topic_partition] = self.acknowledged_per_partition_failed.get(topic_partition, 0) + share_partition_data["count"] |
| logger.debug("Offsets acknowledged for %s" % (node.account.hostname)) |
| logger.debug("Offset acknowledgement failed for: %s" % (node.account.hostname)) |
| |
| def handle_records_consumed(self, event, node, logger): |
| self.total_consumed += event["count"] |
| for share_partition_data in event["partitions"]: |
| topic_partition = TopicPartition(share_partition_data["topic"], share_partition_data["partition"]) |
| self.consumed_per_partition[topic_partition] = self.consumed_per_partition.get(topic_partition, 0) + share_partition_data["count"] |
| logger.debug("Offsets consumed for %s" % (node.account.hostname)) |
| |
| |
| def handle_kill_process(self, clean_shutdown): |
| # if the shutdown was clean, then we expect the explicit |
| # shutdown event from the share consumer |
| if not clean_shutdown: |
| self.handle_shutdown_complete() |
| |
| class VerifiableShareConsumer(KafkaPathResolverMixin, VerifiableClientMixin, BackgroundThreadService): |
| """This service wraps org.apache.kafka.tools.VerifiableShareConsumer for use in |
| system testing. |
| |
| NOTE: this class should be treated as a PUBLIC API. Downstream users use |
| this service both directly and through class extension, so care must be |
| taken to ensure compatibility. |
| """ |
| |
| PERSISTENT_ROOT = "/mnt/verifiable_share_consumer" |
| STDOUT_CAPTURE = os.path.join(PERSISTENT_ROOT, "verifiable_share_consumer.stdout") |
| STDERR_CAPTURE = os.path.join(PERSISTENT_ROOT, "verifiable_share_consumer.stderr") |
| LOG_DIR = os.path.join(PERSISTENT_ROOT, "logs") |
| LOG_FILE = os.path.join(LOG_DIR, "verifiable_share_consumer.log") |
| CONFIG_FILE = os.path.join(PERSISTENT_ROOT, "verifiable_share_consumer.properties") |
| |
| logs = { |
| "verifiable_share_consumer_stdout": { |
| "path": STDOUT_CAPTURE, |
| "collect_default": False}, |
| "verifiable_share_consumer_stderr": { |
| "path": STDERR_CAPTURE, |
| "collect_default": False}, |
| "verifiable_share_consumer_log": { |
| "path": LOG_FILE, |
| "collect_default": True} |
| } |
| |
| def __init__(self, context, num_nodes, kafka, topic, group_id, max_messages=-1, |
| acknowledgement_mode="auto", version=DEV_BRANCH, stop_timeout_sec=60, |
| log_level="INFO", jaas_override_variables=None, on_record_consumed=None): |
| """ |
| :param jaas_override_variables: A dict of variables to be used in the jaas.conf template file |
| """ |
| super(VerifiableShareConsumer, self).__init__(context, num_nodes) |
| self.log_level = log_level |
| self.kafka = kafka |
| self.topic = topic |
| self.group_id = group_id |
| self.max_messages = max_messages |
| self.acknowledgement_mode = acknowledgement_mode |
| self.prop_file = "" |
| self.stop_timeout_sec = stop_timeout_sec |
| self.on_record_consumed = on_record_consumed |
| |
| self.event_handlers = {} |
| self.jaas_override_variables = jaas_override_variables or {} |
| |
| self.total_records_consumed = 0 |
| self.total_records_acknowledged = 0 |
| self.total_records_acknowledged_failed = 0 |
| self.consumed_records_offsets = set() |
| self.acknowledged_records_offsets = set() |
| |
| for node in self.nodes: |
| node.version = version |
| |
| def java_class_name(self): |
| return "VerifiableShareConsumer" |
| |
| def create_event_handler(self, idx, node): |
| return ShareConsumerEventHandler(node, idx) |
| |
| def _worker(self, idx, node): |
| with self.lock: |
| self.event_handlers[node] = self.create_event_handler(idx, node) |
| handler = self.event_handlers[node] |
| |
| node.account.ssh("mkdir -p %s" % VerifiableShareConsumer.PERSISTENT_ROOT, allow_fail=False) |
| |
| # Create and upload log properties |
| log_config = self.render(get_log4j_config_for_tools(node), log_file=VerifiableShareConsumer.LOG_FILE) |
| node.account.create_file(get_log4j_config_for_tools(node), log_config) |
| |
| # Create and upload config file |
| self.security_config = self.kafka.security_config.client_config(self.prop_file, node, |
| self.jaas_override_variables) |
| self.security_config.setup_node(node) |
| self.prop_file += str(self.security_config) |
| self.logger.info("verifiable_share_consumer.properties:") |
| self.logger.info(self.prop_file) |
| node.account.create_file(VerifiableShareConsumer.CONFIG_FILE, self.prop_file) |
| self.security_config.setup_node(node) |
| |
| cmd = self.start_cmd(node) |
| self.logger.debug("VerifiableShareConsumer %d command: %s" % (idx, cmd)) |
| |
| for line in node.account.ssh_capture(cmd): |
| event = self.try_parse_json(node, line.strip()) |
| if event is not None: |
| with self.lock: |
| name = event["name"] |
| if name == "shutdown_complete": |
| handler.handle_shutdown_complete(node, self.logger) |
| elif name == "startup_complete": |
| handler.handle_startup_complete(node, self.logger) |
| elif name == "offsets_acknowledged": |
| handler.handle_offsets_acknowledged(event, node, self.logger) |
| self._update_global_acknowledged(event) |
| elif name == "records_consumed": |
| handler.handle_records_consumed(event, node, self.logger) |
| self._update_global_consumed(event) |
| elif name == "record_data" and self.on_record_consumed: |
| self.on_record_consumed(event, node) |
| else: |
| self.logger.debug("%s: ignoring unknown event: %s" % (str(node.account), event)) |
| |
| def _update_global_acknowledged(self, acknowledge_event): |
| if acknowledge_event["success"]: |
| self.total_records_acknowledged += acknowledge_event["count"] |
| else: |
| self.total_records_acknowledged_failed += acknowledge_event["count"] |
| for share_partition_data in acknowledge_event["partitions"]: |
| tpkey = str(share_partition_data["topic"]) + "-" + str(share_partition_data["partition"]) |
| for offset in share_partition_data["offsets"]: |
| key = tpkey + "-" + str(offset) |
| if key not in self.acknowledged_records_offsets: |
| self.acknowledged_records_offsets.add(key) |
| |
| def _update_global_consumed(self, consumed_event): |
| self.total_records_consumed += consumed_event["count"] |
| |
| for share_partition_data in consumed_event["partitions"]: |
| tpkey = str(share_partition_data["topic"]) + "-" + str(share_partition_data["partition"]) |
| for offset in share_partition_data["offsets"]: |
| key = tpkey + "-" + str(offset) |
| if key not in self.consumed_records_offsets: |
| self.consumed_records_offsets.add(key) |
| |
| def start_cmd(self, node): |
| cmd = "" |
| cmd += "export LOG_DIR=%s;" % VerifiableShareConsumer.LOG_DIR |
| cmd += " export KAFKA_OPTS=%s;" % self.security_config.kafka_opts |
| cmd += " export KAFKA_LOG4J_OPTS=\"%s%s\"; " % (get_log4j_config_param(node), get_log4j_config_for_tools(node)) |
| cmd += self.impl.exec_cmd(node) |
| if self.on_record_consumed: |
| cmd += " --verbose" |
| |
| cmd += " --acknowledgement-mode %s" % self.acknowledgement_mode |
| |
| cmd += " --bootstrap-server %s" % self.kafka.bootstrap_servers(self.security_config.security_protocol) |
| |
| cmd += " --group-id %s --topic %s" % (self.group_id, self.topic) |
| |
| if self.max_messages > 0: |
| cmd += " --max-messages %s" % str(self.max_messages) |
| |
| cmd += " --command-config %s" % VerifiableShareConsumer.CONFIG_FILE |
| cmd += " 2>> %s | tee -a %s &" % (VerifiableShareConsumer.STDOUT_CAPTURE, VerifiableShareConsumer.STDOUT_CAPTURE) |
| return cmd |
| |
| def pids(self, node): |
| return self.impl.pids(node) |
| |
| def try_parse_json(self, node, string): |
| """Try to parse a string as json. Return None if not parseable.""" |
| try: |
| return json.loads(string) |
| except ValueError: |
| self.logger.debug("%s: Could not parse as json: %s" % (str(node.account), str(string))) |
| return None |
| |
| def stop_all(self): |
| for node in self.nodes: |
| self.stop_node(node) |
| |
| def kill_node(self, node, clean_shutdown=True, allow_fail=False): |
| sig = self.impl.kill_signal(clean_shutdown) |
| for pid in self.pids(node): |
| node.account.signal(pid, sig, allow_fail) |
| |
| with self.lock: |
| self.event_handlers[node].handle_kill_process(clean_shutdown) |
| |
| def stop_node(self, node, clean_shutdown=True): |
| self.kill_node(node, clean_shutdown=clean_shutdown) |
| |
| stopped = self.wait_node(node, timeout_sec=self.stop_timeout_sec) |
| assert stopped, "Node %s: did not stop within the specified timeout of %s seconds" % \ |
| (str(node.account), str(self.stop_timeout_sec)) |
| |
| def clean_node(self, node): |
| self.kill_node(node, clean_shutdown=False) |
| node.account.ssh("rm -rf " + self.PERSISTENT_ROOT, allow_fail=False) |
| self.security_config.clean_node(node) |
| |
| def total_consumed(self): |
| with self.lock: |
| return self.total_records_consumed |
| |
| def total_unique_consumed(self): |
| with self.lock: |
| return len(self.consumed_records_offsets) |
| |
| def total_unique_acknowledged(self): |
| with self.lock: |
| return len(self.acknowledged_records_offsets) |
| |
| def total_acknowledged(self): |
| with self.lock: |
| return self.total_records_acknowledged + self.total_records_acknowledged_failed |
| |
| def total_acknowledged_successfully(self): |
| with self.lock: |
| return self.total_records_acknowledged |
| |
| def total_failed_acknowledged(self): |
| with self.lock: |
| return self.total_records_acknowledged_failed |
| |
| def total_consumed_for_a_share_consumer(self, node): |
| with self.lock: |
| return self.event_handlers[node].total_consumed |
| |
| def total_acknowledged_for_a_share_consumer(self, node): |
| with self.lock: |
| return self.event_handlers[node].total_acknowledged_successfully + self.event_handlers[node].total_acknowledged_failed |
| |
| def total_acknowledged_sucessfully_for_a_share_consumer(self, node): |
| with self.lock: |
| return self.event_handlers[node].total_acknowledged_successfully |
| |
| def total_failed_acknowledged_for_a_share_consumer(self, node): |
| with self.lock: |
| return self.event_handlers[node].total_acknowledged_failed |
| |
| def dead_nodes(self): |
| with self.lock: |
| return [handler.node for handler in self.event_handlers.values() |
| if handler.state == ShareConsumerState.Dead] |
| |
| def alive_nodes(self): |
| with self.lock: |
| return [handler.node for handler in self.event_handlers.values() |
| if handler.state == ShareConsumerState.Started] |
