adjust

#!/usr/bin/env python3
"""
Optune 'adjust' driver for Tomcat service.
(C) 2018, Opsani.
use:

TBD

adjust --version
adjust --info
adjust --query app_name
adjust app_name <settings-file.json
This script requires a configuration file placed at a fixed location (see
CFG_FILE constant below). The file should contain the complete formatted
reply to be sent in response to '--query' (see the TDR External Commands draft
document), including default values for all settings (which should match the
value the application would use in case the setting is not provided explicitly
or is empty). The key under 'components' is the name service to be adjusted

Example CFG_FILE:

tomcat:
  consul_url: http://example.com/foo
  health_check_url: "http://{}:8080/bar"
  health_check_timeout: 10
  health_check_ok_string: "<STATUS>OK</STATUS>"
  start_file: /tmp/tomcat.start.new
  restart_cmd: /tmp/tomcat.start.new 0<&- >/dev/null 2>&1 &
  ssh_opts: "-o user=root"

  components:
    gui:
      settings:
        InitialHeapSize: {}
        MaxHeapSize: {}
        MaxGCPauseMillis: {}
        MaxGCMinorPauseMillis: {}
        UseParallelOldGC: {}
        GCTimeRatio: {}
"""


from __future__ import print_function    # py2 compatibility
from concurrent.futures import ThreadPoolExecutor, as_completed

import json
import os
import re
import requests
import subprocess
import tempfile
import time
import yaml

from adjust import Adjust

DESC = "Tomcat adjust driver for Opsani Optune"
VERSION = "1.0.0"
HAS_CANCEL = False

CFG_FILE = "./config.yaml"
DRIVER_NAME = "tomcat"
JAVA_PRINT_FLAGS = "-Xms10m -XX:+PrintFlagsFinal -version"
JAVA_EXEC_CMD = '${JAVA_CMD} ${APP_ARGS}'

HCK_SLEEP_S = 5

# Regexp to be use to update tomcat command line in otrder to get java settings
JAVA_CMD_RE = r'^(.*)(org.apache.catalina.startup.Bootstrap \\)$'


OPTUNE_DRIVER_ANCHOR = "__OPTUNE_DRIVER_ANCHOR__"


class TomcatDriver(Adjust):

    def __init__(self, version, cli_desc, supports_cancel):
        super().__init__(version, cli_desc, supports_cancel)
        self.adjust_failed = False

    # Raise an exception if adjust for a host has already failed. This can be
    # used to stall adjusting of other hosts
    def _check_for_failed_adjust(self, host):
        if not self.adjust_failed:
            return

        msg = "Aborting adjust for host %s as another host has already failed" \
            % host

        self.debug(msg)
        raise Exception(msg)

    def _run_cmd(self, cmd):
        '''
        Takes a command (as a list of strings), runs it and returns exit_code,
        stdout (list of lines) and stderr (list of lines)
        '''

        out = []
        err = []

        proc = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        for line in proc.stdout.readlines():
            out.append(line)

        for line in proc.stderr.readlines():
            err.append(line)

        return_code = proc.wait()

        return return_code, out, err

    def _parse_jvm_settings(self, lines):
        '''
        Takes a list of strings, i.e. the output of java with flags
        `-XX:+PrintFlagsFinal -version`

        Returns a dict where the keys are the flag names and the values are
        dicts with two keys: type and value of the flag
        '''

        java_curr_settings = {}
        anchor_found = False
        for line in lines[1:]:
            line = line.decode()

            if OPTUNE_DRIVER_ANCHOR in line:
                anchor_found = True

            if not anchor_found:
                continue

            m = re.search(r'^\s*(\S+)\s+(\S+)\s+:?=\s+(\S+)\s+{.+}$', line)
            if m:
                t = m.group(1)
                name = m.group(2)
                value = m.group(3)

                if t == "intx" or t == "uint64_t" or t == "uintx":
                    value = int(value)
                elif t == "double":
                    value = float(value)
                elif t == "bool":
                    value = True if value == "true" else False

                java_curr_settings[name] = {
                    "type": t,
                    "value": value,
                }

        return java_curr_settings

    def _parse_cfg(self):
        '''
        load the user-defined application descriptor
        '''

        try:
            f = open(CFG_FILE)
            d = yaml.load(f)
        except IOError as e:
            raise Exception(
                "cannot read configuration from {}:{}".format(CFG_FILE, e.strerror))
        except yaml.error.YAMLError as e:
            raise Exception("syntax error in {}: {}".format(CFG_FILE, str(e)))
        # everything else: raise

        # Valchk
        assert(DRIVER_NAME in d), \
            "Missing driver configuration: missing {} section".format(
                DRIVER_NAME)

        cfg = d[DRIVER_NAME]

        if not "components" in cfg:
            raise Exception(
                "Missing driver configuration: Missing components key.")

        # Valchk
        if len(cfg["components"].keys()) != 1:
            raise Exception(
                "App descriptor error: must have exactly one component")

        # Make sure Healthcheck config is present
        assert("health_check_url" in cfg and
               isinstance(cfg["health_check_url"], str)), \
            "Missing or invalid driver configuration: health_check_url"

        assert("health_check_timeout" in cfg and
               isinstance(cfg["health_check_timeout"], int)), \
            "Missing or invalid driver configuration: health_check_timeout"

        assert("health_check_ok_string" in cfg and
               isinstance(cfg["health_check_ok_string"], str)), \
            "Missing or invalid driver configuration: health_check_ok_string"

        # Make sure we have consul config or list of instances
        assert(
            ("consul_url" in cfg and isinstance(cfg["consul_url"], str)) or
            ("instances" in cfg and isinstance(cfg["instances"], dict))
        ), \
            "Missing or invalid driver configuration: provide either 'consul_url' or 'instances'"

        # Make sure Healthcheck config is present
        assert("start_file" in cfg and
               isinstance(cfg["start_file"], str)), \
            "Missing or invalid driver configuration: start_file"

        # Make sure Healthcheck config is present
        assert("restart_cmd" in cfg and
               isinstance(cfg["restart_cmd"], str)), \
            "Missing or invalid driver configuration: restart_cmd"

        return cfg

    def _get_remote_file(self, host, remote_path, local_path, ssh_opts):

        cmd = ['scp']
        cmd.extend(ssh_opts)
        cmd.append(host + ":" + remote_path)
        cmd.append(local_path)

        exit_code, out, err = self._run_cmd(cmd)

        if exit_code != 0:
            msg = "Failed to get remote file {}, from host {}. Command {} " + \
                "returned exit code {}. STDOUT {}, STDERR {}"
            msg = msg.format(remote_path, host, str(
                cmd), exit_code, str(out), str(err))
            raise Exception(msg)

    def _put_remote_file(self, host, local_path, remote_path, ssh_opts):

        cmd = ['scp']
        cmd.extend(ssh_opts)
        cmd.append(local_path)
        cmd.append(host + ":" + remote_path)

        exit_code, out, err = self._run_cmd(cmd)

        if exit_code != 0:
            msg = "Failed to put file {}, to host {}. Command {} " + \
                "returned exit code {}. STDOUT {}, STDERR {}"
            msg = msg.format(remote_path, host, str(
                cmd), exit_code, str(out), str(err))
            raise Exception(msg)

        # Make file executable
        cmd = ['ssh']
        cmd.extend(ssh_opts)
        cmd.extend([host, "chmod", "+x", remote_path])

        exit_code, out, err = self._run_cmd(cmd)

        if exit_code != 0:
            msg = "Failed to make file {} executable on host {}. Command {} " + \
                "returned exit code {}. STDOUT {}, STDERR {}"
            msg = msg.format(remote_path, host, str(
                cmd), exit_code, str(out), str(err))
            raise Exception(msg)

    def _run_remote_cmd(self, host, cmd, ssh_opts):
        r_cmd = ["ssh"]
        r_cmd.extend(ssh_opts)
        r_cmd.append(host)
        r_cmd.extend(cmd)

        exit_code, out, err = self._run_cmd(r_cmd)

        if exit_code != 0:
            msg = "Failed to run command {} executable on host {}, it " + \
                "returned exit code {}. STDOUT {}, STDERR {}"
            msg = msg.format(cmd, host, exit_code, str(out), str(err))
            raise Exception(msg)

        return out

    def _get_ssh_opts_from_desc(self, cfg):
        try:
            ssh_opts = cfg["ssh_opts"].split()
        except Exception as e:
            ssh_opts = []

        return ssh_opts

    def _update_remote_tomcat_start_script(self, host, get_file, put_file, ssh_opts, line_xfrom_fn):
        with tempfile.NamedTemporaryFile() as fh:
            temp_file = fh.name

            # Get current tomcat config
            self._get_remote_file(host, get_file, temp_file, ssh_opts)

            # Modify tomcat file so it can query settings
            lines = []
            # Read file and update in memory
            for line in fh.readlines():
                line = line.decode()

                updated_line, is_last = line_xfrom_fn(line)
                lines.append(updated_line)
                if is_last:
                    break

        # Create new temp file with modified content and upload to remote host
        _, temp_file = tempfile.mkstemp()
        try:
            with open(temp_file, 'w') as f:
                for line in lines:
                    f.write(line)
            # Upload file
            self._put_remote_file(host, temp_file, put_file, ssh_opts)
        except Exception as e:
            raise Exception(
                "Failed to update tomcat file {}: {}".format(temp_file, str(e))
            )
        finally:
            os.unlink(temp_file)

    def _get_java_settings(self, host, tomcat_start_path, ssh_opts, extra_settings):

        tomcat_start_path_tmp = tomcat_start_path + '.tmp'

        def line_xfrom_fn(line):
            m = re.search(JAVA_CMD_RE, line)
            if not m:
                return line, False
            else:
                s = extra_settings if extra_settings else m.group(1)
                ret_line = s + " " + JAVA_PRINT_FLAGS + '"\n\n'
                ret_line += 'echo ' + OPTUNE_DRIVER_ANCHOR + '\n'
                ret_line += JAVA_EXEC_CMD + '\n'
                ret_line += 'exit 0\n'
                return ret_line, True

        self._update_remote_tomcat_start_script(
            host, tomcat_start_path, tomcat_start_path_tmp, ssh_opts,
            line_xfrom_fn)

        # Run file remotely and capture output
        out = self._run_remote_cmd(host, [tomcat_start_path_tmp], ssh_opts)

        # Parse JVM settings
        settings = self._parse_jvm_settings(out)
        assert(settings), "Failed to get Java settings, output was: {}".format(out)

        return settings

    def _get_curr_java_settings(self, host, tomcat_start_path, ssh_opts):
        return self._get_java_settings(host, tomcat_start_path, ssh_opts, "")

    def _get_instances(self, cfg):
        svc_name = list(cfg["components"].keys())[0]

        if "consul_url" in cfg:
            cert_path = cfg.get("consul_cert", None)
            self.debug("Getting list of instances from consul at",
                       cfg["consul_url"], cert_path)

            r = requests.get(
                cfg["consul_url"],
                cert=cert_path,
                verify=False)

            hosts = list(map(lambda x: x["Node"]["Address"], r.json()))
        else:
            self.debug("Getting list of instances from descriptor")
            hosts = cfg["instances"][svc_name]

        self.debug("Using instances:", hosts)

        return hosts

    def _adjust_host(
        self,
        host,
        settings,
        hchk_url,
        hchk_timeout,
        hchk_ok_string,
        hck_init_sleep,
        tomcat_start_path,
        tomcat_restart_cmd,
        ssh_opts):

        self._check_for_failed_adjust(host)

        self.debug("Updating host", host)

        java_curr_settings = self._get_curr_java_settings(
            host, tomcat_start_path, ssh_opts)

        # Convert Optune seetings to java settings
        java_settings = ""
        for k in settings.keys():
            v = settings[k]["value"]
            curr = java_curr_settings[k]
            if curr["type"] == "bool":
                java_settings += " -XX:{}{}".format("+" if v else "-", k)
            elif curr["type"] == "intx" or curr["type"] == "uint64_t" or curr["type"] == "uintx":
                java_settings += " -XX:{}={}".format(k, int(v))
            else:
                java_settings += " -XX:{}={}".format(k, v)

        self._check_for_failed_adjust(host)

        # Do a dry run with the new settings
        try:
            self._get_java_settings(
                host, tomcat_start_path, ssh_opts, java_settings)
        except Exception as e:
            raise Exception(
                "Failed to do a dry run for adjustment: {}".format(str(e))
            )

        def line_xfrom_fn(line):
            m = re.search(JAVA_CMD_RE, line)

            if not m:
                return line, False
            else:
                ret_line = java_settings + " " + m.group(2) + '\n'
                return ret_line, False

        self._check_for_failed_adjust(host)

        # Update tomcat config
        self._update_remote_tomcat_start_script(
            host, tomcat_start_path, tomcat_start_path, ssh_opts,
            line_xfrom_fn)

        self._check_for_failed_adjust(host)

        # Restart tomcat
        self._run_remote_cmd(host, tomcat_restart_cmd, ssh_opts)


        self._check_for_failed_adjust(host)

        # Run healtcheck
        if not hchk_url:
            self.debug("Skipping healthcheck for host", host)
            return

        if hck_init_sleep:
            self.debug(
                "Sleeping for {} seconds before doing a healthcheck on host {}"
                .format(hck_init_sleep, host))

            time.sleep(hck_init_sleep)

        run_until_ts = time.time() + hchk_timeout
        host_healthy = False

        while time.time() < run_until_ts:
            self._check_for_failed_adjust(host)

            self.debug("Running healthcheck for url", hchk_url)
            try:
                r = requests.get(hchk_url, verify=False)
                assert(r.status_code == 200), \
                    "Healthcheck returned status " + str(r.status_code)
                assert(hchk_ok_string in r.text), \
                    "Did not find string '{}' in healthcheck response".format(
                        hchk_ok_string)

                host_healthy = True
                self.debug("Healthcheck succeeded for host", host)
                break
            except Exception as e:
                self.debug("Healthcheck try failed, error was:", e)
                time.sleep(HCK_SLEEP_S)

        if not host_healthy:
            raise Exception(
                "Host {} did not return healthcheck reponse {} in {} seconds".format(
                    host, hchk_ok_string, hchk_timeout))

    # Overwritten
    def query(self):
        # Parse app descriptor
        cfg = self._parse_cfg()

        svc_name = list(cfg["components"].keys())[0]

        ssh_opts = self._get_ssh_opts_from_desc(cfg)

        # TODO: do for each instance and make sure results match
        host = self._get_instances(cfg)[0]

        tomcat_start_path = cfg["start_file"]

        java_curr_settings = self._get_curr_java_settings(
            host, tomcat_start_path, ssh_opts)

        # Populate settings from the app descriptor with their actual values
        settings = cfg["components"][svc_name]["settings"]
        for k in settings.keys():
            if not k in java_curr_settings:
                raise Exception(
                    "Counld not find setting {} in current JVM settings".format(k))

            curr = java_curr_settings[k]
            settings[k]["value"] = curr["value"]
            settings[k]["type"] = "range"
            settings[k]["step"] = 1

            # Convert bools to an int with range 0-1
            if curr["type"] == "bool":
                settings[k]["value"] = 1 if curr["value"] else 0
                settings[k]["min"] = 0
                settings[k]["max"] = 1

        ret = {
            "components": {
                svc_name: {
                    "settings": settings
                }
            }
        }

        return ret

    # Overwritten
    def adjust(self):
        # Parse app descriptor
        cfg = self._parse_cfg()

        svc_name = list(cfg["components"].keys())[0]

        ssh_opts = self._get_ssh_opts_from_desc(cfg)

        tomcat_start_path = cfg["start_file"]

        tomcat_restart_cmd = cfg["restart_cmd"].split()

        hosts = self._get_instances(cfg)

        # Validate input
        if not "application" in self.input_data:
            raise Exception("Invalid input: missing application key")

        if not "components" in self.input_data["application"]:
            raise Exception(
                "Invalid input: missing application.components key")

        if not svc_name in self.input_data["application"]["components"]:
            raise Exception(
                "Invalid input: component {} not found in input data".format(
                    svc_name
                ))

        if not "settings" in self.input_data["application"]["components"][svc_name]:
            raise Exception(
                "Invalid input: settings key not present for component " + svc_name
            )

        settings = self.input_data["application"]["components"][svc_name]["settings"]

        input_user_data = self.input_data.get(
            "control", {}).get("userdata", {})

        if not isinstance(settings, dict):
            raise Exception(
                "Invalid input: settings for component {} must be a dict".format(
                    svc_name
                ))

        # Get health_check_initial_sleep from input data or config
        hck_init_sleep = int(input_user_data.get(
            "health_check_initial_sleep",
            cfg.get('health_check_initial_sleep', 0)
        ))

        hchk_timeout = int(input_user_data.get(
            "health_check_timeout",
            cfg.get("health_check_timeout", 0)
        ))
        hchk_ok_string = cfg["health_check_ok_string"]

        concurrent_hosts = int(cfg.get("n_concurrent_updates", 1))
        futures = {}

        self.debug("Updating hosts, %d at a time" % concurrent_hosts)
        with ThreadPoolExecutor(max_workers=concurrent_hosts) as executor:
            # submit work
            for host in hosts:
                if 'health_check_url' in cfg:
                    hchk_url = cfg['health_check_url'].format(host)
                else:
                    hchk_url = None

                future = executor.submit(
                    self._adjust_host,host,
                    settings,
                    hchk_url,
                    hchk_timeout,
                    hchk_ok_string,
                    hck_init_sleep,
                    tomcat_start_path,
                    tomcat_restart_cmd,
                    ssh_opts)

                futures[future] = host

            # Get results
            completed_hosts = 0
            total_hosts = len(hosts)

            self.debug("Waiting for host updates")

            for future in as_completed(futures):
                host = futures[future]

                try:
                    result = future.result()
                except Exception as e:
                    self.debug('host %s failed to update: %s' % (host, e))
                    future.cancel()
                    self.adjust_failed = True
                    raise
                else:
                    self.debug('host %s successfully updated' % host)
                    completed_hosts += 1
                    progress = int(100*completed_hosts/total_hosts)
                    prog_data = {
                        'progress': progress,
                        'message': 'updated host %s' % host,
                    }
                    print(json.dumps(prog_data), flush=True)

        self.debug("Sucessfully updated settings on all hosts")


if __name__ == '__main__':
    driver = TomcatDriver(VERSION, DESC, HAS_CANCEL)
    driver.run()