content: # 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.
  #!/usr/bin/env python
  
  # ===================================
  # file: metrics.py
  # ===================================
  
  import inspect
  import json
  import logging
  import os
  import signal
  import subprocess
  import sys
  import traceback
  
  import csv
  import time 
  import matplotlib as mpl
  mpl.use('Agg')
  import matplotlib.pyplot as plt
  from collections import namedtuple
  import numpy
  
  from pyh import *
  import kafka_system_test_utils
  import system_test_utils
  
  logger     = logging.getLogger("namedLogger")
  thisClassName = '(metrics)'
  d = {'name_of_class': thisClassName}
  
  attributeNameToNameInReportedFileMap = {
      'Min': 'min',
      'Max': 'max',
      'Mean': 'mean',
      '50thPercentile': 'median',
      'StdDev': 'stddev',
      '95thPercentile': '95%',
      '99thPercentile': '99%',
      '999thPercentile': '99.9%',
      'Count': 'count',
      'OneMinuteRate': '1 min rate',
      'MeanRate': 'mean rate',
      'FiveMinuteRate': '5 min rate',
      'FifteenMinuteRate': '15 min rate',
      'Value': 'value'
  }
  
  def getCSVFileNameFromMetricsMbeanName(mbeanName):
      return mbeanName.replace(":type=", ".").replace(",name=", ".") + ".csv"
  
  def read_metrics_definition(metricsFile):
      metricsFileData = open(metricsFile, "r").read()
      metricsJsonData = json.loads(metricsFileData)
      allDashboards = metricsJsonData['dashboards']
      allGraphs = []
      for dashboard in allDashboards:
          dashboardName = dashboard['name']
          graphs = dashboard['graphs']
          for graph in graphs:
              bean = graph['bean_name']
              allGraphs.append(graph)
              attributes = graph['attributes']
              #print "Filtering on attributes " + attributes     
      return allGraphs
              
  def get_dashboard_definition(metricsFile, role):
      metricsFileData = open(metricsFile, "r").read()
      metricsJsonData = json.loads(metricsFileData)
      allDashboards = metricsJsonData['dashboards']
      dashboardsForRole = []
      for dashboard in allDashboards:
          if dashboard['role'] == role:
              dashboardsForRole.append(dashboard) 
      return dashboardsForRole
  
  def ensure_valid_headers(headers, attributes):
      if headers[0] != "# time":
          raise Exception("First column should be time")
      for header in headers:
          logger.debug(header, extra=d)
      # there should be exactly one column with a name that matches attributes
      try:
          attributeColumnIndex = headers.index(attributes)
          return attributeColumnIndex
      except ValueError as ve:
          #print "#### attributes : ", attributes
          #print "#### headers    : ", headers
          raise Exception("There should be exactly one column that matches attribute: {0} in".format(attributes) +  
                          " headers: {0}".format(",".join(headers)))
          
  def plot_graphs(inputCsvFiles, labels, title, xLabel, yLabel, attribute, outputGraphFile):
      # create empty plot
      fig=plt.figure()
      fig.subplots_adjust(bottom=0.2)
      ax=fig.add_subplot(111)
      labelx = -0.3  # axes coords
      ax.set_xlabel(xLabel)
      ax.set_ylabel(yLabel)
      ax.grid()
      #ax.yaxis.set_label_coords(labelx, 0.5)
      Coordinates = namedtuple("Coordinates", 'x y')
      plots = []
      coordinates = []
      # read data for all files, organize by label in a dict
      for fileAndLabel in zip(inputCsvFiles, labels):
          inputCsvFile = fileAndLabel[0]
          label = fileAndLabel[1]
          csv_reader = list(csv.reader(open(inputCsvFile, "rb")))
          x,y = [],[]
          xticks_labels = []
          try:
              # read first line as the headers
              headers = csv_reader.pop(0)
              attributeColumnIndex = ensure_valid_headers(headers, attributeNameToNameInReportedFileMap[attribute])
              logger.debug("Column index for attribute {0} is {1}".format(attribute, attributeColumnIndex), extra=d)
              start_time = (int)(os.path.getctime(inputCsvFile) * 1000)
              int(csv_reader[0][0])
              for line in csv_reader:
                  if(len(line) == 0):
                      continue
                  yVal = float(line[attributeColumnIndex])                
                  xVal = int(line[0])
                  y.append(yVal)
                  epoch= start_time + int(line[0])
                  x.append(xVal)
                  xticks_labels.append(time.strftime("%H:%M:%S", time.localtime(epoch)))
                  coordinates.append(Coordinates(xVal, yVal))
              p1 = ax.plot(x,y)
              plots.append(p1)
          except Exception as e:
              logger.error("ERROR while plotting data for {0}: {1}".format(inputCsvFile, e), extra=d)
              traceback.print_exc()
      # find xmin, xmax, ymin, ymax from all csv files
      xmin = min(map(lambda coord: coord.x, coordinates))
      xmax = max(map(lambda coord: coord.x, coordinates))
      ymin = min(map(lambda coord: coord.y, coordinates))
      ymax = max(map(lambda coord: coord.y, coordinates))
      # set x and y axes limits
      plt.xlim(xmin, xmax)
      plt.ylim(ymin, ymax)
      # set ticks accordingly
      xticks = numpy.arange(xmin, xmax, 0.2*xmax)
  #    yticks = numpy.arange(ymin, ymax)
      plt.xticks(xticks,xticks_labels,rotation=17)
  #    plt.yticks(yticks)
      plt.legend(plots,labels, loc=2)
      plt.title(title)
      plt.savefig(outputGraphFile)
  
  def draw_all_graphs(metricsDescriptionFile, testcaseEnv, clusterConfig):
      # go through each role and plot graphs for the role's metrics
      roles = set(map(lambda config: config['role'], clusterConfig))
      for role in roles:
          dashboards = get_dashboard_definition(metricsDescriptionFile, role)
          entities = kafka_system_test_utils.get_entities_for_role(clusterConfig, role)
          for dashboard in dashboards:
              graphs = dashboard['graphs']
              # draw each graph for all entities
              draw_graph_for_role(graphs, entities, role, testcaseEnv)
          
  def draw_graph_for_role(graphs, entities, role, testcaseEnv):
      for graph in graphs:
          graphName = graph['graph_name'] 
          yLabel = graph['y_label']
          inputCsvFiles = []
          graphLegendLabels = []
          for entity in entities:
              entityMetricsDir = kafka_system_test_utils.get_testcase_config_log_dir_pathname(testcaseEnv, role, entity['entity_id'], "metrics")
              entityMetricCsvFile = entityMetricsDir + "/" + getCSVFileNameFromMetricsMbeanName(graph['bean_name'])
              if(not os.path.exists(entityMetricCsvFile)):
                  logger.warn("The file {0} does not exist for plotting".format(entityMetricCsvFile), extra=d)
              else:
                  inputCsvFiles.append(entityMetricCsvFile)
                  graphLegendLabels.append(role + "-" + entity['entity_id'])
  #            print "Plotting graph for metric {0} on entity {1}".format(graph['graph_name'], entity['entity_id'])
          try:
              # plot one graph per mbean attribute
              labels = graph['y_label'].split(',')
              fullyQualifiedAttributeNames = map(lambda attribute: graph['bean_name'] + ':' + attribute, 
                                             graph['attributes'].split(','))
              attributes = graph['attributes'].split(',')
              for labelAndAttribute in zip(labels, fullyQualifiedAttributeNames, attributes):            
                  outputGraphFile = testcaseEnv.testCaseDashboardsDir + "/" + role + "/" + labelAndAttribute[1] + ".svg"            
                  plot_graphs(inputCsvFiles, graphLegendLabels, graph['graph_name'] + '-' + labelAndAttribute[2], 
                              "time", labelAndAttribute[0], labelAndAttribute[2], outputGraphFile)
  #            print "Finished plotting graph for metric {0} on entity {1}".format(graph['graph_name'], entity['entity_id'])
          except Exception as e:
              logger.error("ERROR while plotting graph {0}: {1}".format(outputGraphFile, e), extra=d)
              traceback.print_exc()
  
  def build_all_dashboards(metricsDefinitionFile, testcaseDashboardsDir, clusterConfig):
      metricsHtmlFile = testcaseDashboardsDir + "/metrics.html"
      centralDashboard = PyH('Kafka Metrics Dashboard')
      centralDashboard << h1('Kafka Metrics Dashboard', cl='center')
      roles = set(map(lambda config: config['role'], clusterConfig))
      for role in roles:
          entities = kafka_system_test_utils.get_entities_for_role(clusterConfig, role)
          dashboardPagePath = build_dashboard_for_role(metricsDefinitionFile, role, 
                                                       entities, testcaseDashboardsDir)
          centralDashboard << a(role, href = dashboardPagePath)
          centralDashboard << br()
              
      centralDashboard.printOut(metricsHtmlFile)
  
  def build_dashboard_for_role(metricsDefinitionFile, role, entities, testcaseDashboardsDir):
      # build all dashboards for the input entity's based on its role. It can be one of kafka, zookeeper, producer
      # consumer
      dashboards = get_dashboard_definition(metricsDefinitionFile, role)
      entityDashboard = PyH('Kafka Metrics Dashboard for ' + role)
      entityDashboard << h1('Kafka Metrics Dashboard for ' + role, cl='center')
      entityDashboardHtml = testcaseDashboardsDir + "/" + role + "-dashboards.html"
      for dashboard in dashboards:
          # place the graph svg files in this dashboard
          allGraphs = dashboard['graphs']
          for graph in allGraphs:
              attributes = map(lambda attribute: graph['bean_name'] + ':' + attribute, 
                                             graph['attributes'].split(','))
              for attribute in attributes:                
                  graphFileLocation = testcaseDashboardsDir + "/" + role + "/" + attribute + ".svg"
                  entityDashboard << embed(src = graphFileLocation, type = "image/svg+xml")
      entityDashboard.printOut(entityDashboardHtml)
      return entityDashboardHtml
  
  def start_metrics_collection(jmxHost, jmxPort, role, entityId, systemTestEnv, testcaseEnv):
      logger.info("starting metrics collection on jmx port : " + jmxPort, extra=d)
      jmxUrl = "service:jmx:rmi:///jndi/rmi://" + jmxHost + ":" + jmxPort + "/jmxrmi"
      clusterConfig = systemTestEnv.clusterEntityConfigDictList
      metricsDefinitionFile = systemTestEnv.METRICS_PATHNAME
      entityMetricsDir = kafka_system_test_utils.get_testcase_config_log_dir_pathname(testcaseEnv, role, entityId, "metrics")
      dashboardsForRole = get_dashboard_definition(metricsDefinitionFile, role)
      mbeansForRole = get_mbeans_for_role(dashboardsForRole)
      
      kafkaHome = system_test_utils.get_data_by_lookup_keyval(clusterConfig, "entity_id", entityId, "kafka_home")
      javaHome  = system_test_utils.get_data_by_lookup_keyval(clusterConfig, "entity_id", entityId, "java_home")
      
      for mbean in mbeansForRole:
          outputCsvFile = entityMetricsDir + "/" + mbean + ".csv"
          startMetricsCmdList = ["ssh " + jmxHost,
                                 "'JAVA_HOME=" + javaHome,
                                 "JMX_PORT= " + kafkaHome + "/bin/kafka-run-class.sh kafka.tools.JmxTool",
                                 "--jmx-url " + jmxUrl,
                                 "--object-name " + mbean + " 1> ",
                                  outputCsvFile + " & echo pid:$! > ",
                                  entityMetricsDir + "/entity_pid'"]
  
          startMetricsCommand = " ".join(startMetricsCmdList) 
          logger.debug("executing command: [" + startMetricsCommand + "]", extra=d)
          system_test_utils.async_sys_call(startMetricsCommand)
          time.sleep(1)
  
          pidCmdStr = "ssh " + jmxHost + " 'cat " + entityMetricsDir + "/entity_pid' 2> /dev/null"
          logger.debug("executing command: [" + pidCmdStr + "]", extra=d)
          subproc = system_test_utils.sys_call_return_subproc(pidCmdStr)
  
          # keep track of JMX ppid in a dictionary of entity_id to list of JMX ppid
          # testcaseEnv.entityJmxParentPidDict:
          #   key: entity_id
          #   val: list of JMX ppid associated to that entity_id
          #   { 1: [1234, 1235, 1236], 2: [2234, 2235, 2236], ... }
          for line in subproc.stdout.readlines():
              line = line.rstrip('\n')
              logger.debug("line: [" + line + "]", extra=d)
              if line.startswith("pid"):
                  logger.debug("found pid line: [" + line + "]", extra=d)
                  tokens  = line.split(':')
                  thisPid = tokens[1]
                  if entityId not in testcaseEnv.entityJmxParentPidDict:
                      testcaseEnv.entityJmxParentPidDict[entityId] = []
                  testcaseEnv.entityJmxParentPidDict[entityId].append(thisPid)
                  #print "\n#### testcaseEnv.entityJmxParentPidDict ", testcaseEnv.entityJmxParentPidDict, "\n"
  
  
  def stop_metrics_collection(jmxHost, jmxPort):
      logger.info("stopping metrics collection on " + jmxHost + ":" + jmxPort, extra=d)
      system_test_utils.sys_call("ps -ef | grep JmxTool | grep -v grep | grep " + jmxPort + " | awk '{print $2}' | xargs kill -9")
  
  def get_mbeans_for_role(dashboardsForRole):
      graphs = reduce(lambda x,y: x+y, map(lambda dashboard: dashboard['graphs'], dashboardsForRole))
      return set(map(lambda metric: metric['bean_name'], graphs))
  