~veebers/juju-ci-tools/explore-perf-bulk-model-destruction

#!/usr/bin/env python
from __future__ import print_function
from argparse import ArgumentParser
import contextlib
from copy import (
    copy,
    deepcopy,
)
import logging
import re
import os
import subprocess
import sys
import tempfile
from textwrap import dedent
import time

from deploy_stack import (
    BootstrapManager,
    get_random_string,
    update_env,
)
from jujupy import (
    KVM_MACHINE,
    LXC_MACHINE,
    LXD_MACHINE,
)
from utility import (
    add_basic_testing_arguments,
    configure_logging,
    wait_for_port,
)


__metaclass__ = type


log = logging.getLogger("assess_container_networking")


# This needs refactored out to utility
class JujuAssertionError(AssertionError):
    """Exception for juju assertion failures."""


def parse_args(argv=None):
    """Parse all arguments."""

    description = dedent("""\
    Test container address allocation.
    For LXC and KVM, create machines of each type and test the network
    between LXC <--> LXC, KVM <--> KVM and LXC <--> KVM. Also test machine
    to outside world, DNS and that these tests still pass after a reboot. In
    case of failure pull logs and configuration files from the machine that
    we detected a problem on for later analysis.
    """)
    parser = add_basic_testing_arguments(ArgumentParser(
        description=description
    ))
    parser.add_argument(
        '--machine-type',
        help='Which virtual machine/container type to test. Defaults to all.',
        choices=[KVM_MACHINE, LXC_MACHINE, LXD_MACHINE])
    parser.add_argument(
        '--clean-environment', action='store_true', help=dedent("""\
        Attempts to re-use an existing environment rather than destroying it
        and creating a new one.

        On launch, if an environment exists, clean out services and machines
        from it rather than destroying it. If an environment doesn't exist,
        create one and use it.

        At termination, clean out services and machines from the environment
        rather than destroying it."""))
    args = parser.parse_args(argv)
    # Passing --clean-environment implies --keep-env
    if args.clean_environment:
        args.keep_env = True
    return args


def ssh(client, machine, cmd):
    """Convenience function: run a juju ssh command and get back the output
    :param client: A Juju client
    :param machine: ID of the machine on which to run a command
    :param cmd: the command to run
    :return: text output of the command
    """
    back_off = 2
    attempts = 4
    for attempt in range(attempts):
        try:
            return client.get_juju_output('ssh', '--proxy', machine, cmd)
        except subprocess.CalledProcessError as e:
            # If the connection to the host failed, try again in a couple of
            # seconds. This is usually due to heavy load.
            if(attempt < attempts - 1 and
                re.search('ssh_exchange_identification: '
                          'Connection closed by remote host', e.stderr)):
                time.sleep(back_off)
                back_off *= 2
            else:
                raise


def clean_environment(client, services_only=False):
    """Remove all the services and, optionally, machines from an environment.

    Use as an alternative to destroying an environment and creating a new one
    to save a bit of time.

    :param client: a Juju client
    """
    # A short timeout is used for get_status here because if we don't get a
    # response from  get_status quickly then the environment almost
    # certainly doesn't exist or needs recreating.
    try:
        status = client.get_status(5)
    except Exception as e:
        # TODO(gz): get_status should return a more specific error type.
        log.info("Could not clean existing env: %s", e)
        return False

    for service in status.get_applications():
        client.remove_service(service)

    if not services_only:
        # First remove all containers; we can't remove a machine that is
        # hosting containers.
        for m, _ in status.iter_machines(containers=True, machines=False):
            client.juju('remove-machine', m)

        client.wait_for('containers', 'none')

        for m, _ in status.iter_machines(containers=False, machines=True):
            if m != '0':
                try:
                    client.juju('remove-machine', m)
                except subprocess.CalledProcessError:
                    # Sometimes this fails because while we have asked Juju
                    # to remove a container and it says that it has, when we
                    # ask it to remove the host Juju thinks it still has
                    # containers on it. Normally a small pause and trying
                    # again is all that is needed to resolve this issue.
                    time.sleep(2)
                    client.wait_for_started()
                    client.juju('remove-machine', m)

        client.wait_for('machines-not-0', 'none')

    client.wait_for_started()
    return True


def make_machines(client, container_types):
    """Make a test environment consisting of:
       Two host machines.
       Two of each container_type on one host machine.
       One of each container_type on one host machine.
    :param client: An EnvJujuClient
    :param container_types: list of containers to create
    :return: hosts (list), containers {container_type}{host}[containers]
    """
    # Find existing host machines
    old_hosts = client.get_status().status['machines']
    machines_to_add = 2 - len(old_hosts)

    # Allocate more hosts as needed
    if machines_to_add > 0:
        client.juju('add-machine', ('-n', str(machines_to_add)))
    status = client.wait_for_started()
    hosts = sorted(status.status['machines'].keys())[:2]

    # Find existing containers
    required = dict(zip(hosts, [copy(container_types) for h in hosts]))
    required[hosts[0]] += container_types
    for c in status.iter_machines(containers=True, machines=False):
        host, type, id = c[0].split('/')
        if type in required[host]:
            required[host].remove(type)

    # Start any new containers we need
    for host, containers in required.iteritems():
        for container in containers:
            client.juju('add-machine', ('{}:{}'.format(container, host)))

    status = client.wait_for_started()

    # Build a list of containers, now they have all started
    tmp = dict(zip(hosts, [[] for h in hosts]))
    containers = dict(zip(container_types,
                          [deepcopy(tmp) for t in container_types]))
    for c in status.iter_machines(containers=True, machines=False):
        host, type, id = c[0].split('/')
        if type in containers and host in containers[type]:
            containers[type][host].append(c[0])

    return hosts, containers


def find_network(client, machine, addr):
    """Find a connected subnet containing the given address.

    When using this to find the subnet of a container, don't use the container
    as the machine to run the ip route show command on ("machine"), use a real
    box because lxc will just send everything to its host machine, so it is on
    a subnet containing itself. Not much use.
    :param client: A Juju client
    :param machine: ID of the machine on which to run a command
    :param addr: find the connected subnet containing this address
    :return: CIDR containing the address if found, else, None
    """
    ip_cmd = ' '.join(['ip', 'route', 'show', 'to', 'match', addr])
    routes = ssh(client, machine, ip_cmd)

    for route in re.findall(r'^(\S+).*[\d\.]+/\d+', routes, re.MULTILINE):
        if route != 'default':
            return route

    raise ValueError("Unable to find route to %r" % addr)


def assess_network_traffic(client, targets):
    """Test that all containers in target can talk to target[0]
    :param client: Juju client
    :param targets: machine IDs of machines to test
    :return: None;
    """
    status = client.wait_for_started().status
    log.info('Assessing network traffic.')
    source = targets[0]
    dests = targets[1:]

    with tempfile.NamedTemporaryFile(delete=False) as f:
        f.write('tmux new-session -d -s test "nc -l 6778 > nc_listen.out"')
    client.juju('scp', ('--proxy', f.name, source + ':/home/ubuntu/listen.sh'))
    os.remove(f.name)

    # Containers are named 'x/type/y' where x is the host of the container. We
    host = source.split('/')[0]
    address = status['machines'][host]['containers'][source]['dns-name']

    for dest in dests:
        log.info('Assessing network traffic for {}.'.format(dest))
        msg = get_random_string()
        ssh(client, source, 'rm nc_listen.out; bash ./listen.sh')
        ssh(client, dest,
            'echo "{msg}" | nc {addr} 6778'.format(msg=msg, addr=address))
        result = ssh(client, source, 'more nc_listen.out')
        if result.rstrip() != msg:
            raise ValueError("Wrong or missing message: %r" % result.rstrip())
        log.info('SUCCESS.')


def private_address(client, host):
    default_route = ssh(client, host, 'ip -4 -o route list 0/0')
    log.info("Default route from {}: {}".format(host, default_route))
    route_match = re.search(r'([\w-]+)\s*$', default_route)
    if route_match is None:
        raise JujuAssertionError(
            "Failed to find device in {}".format(default_route))
    device = route_match.group(1)
    log.info("Fetching the device IP of {}".format(device))
    device_ip = ssh(client, host, 'ip -4 -o addr show {}'.format(device))
    log.info("Device IP for {}: {}".format(host, device_ip))
    ip_match = re.search(r'inet\s+(\S+)/\d+\s', device_ip)
    if ip_match is None:
        raise JujuAssertionError(
            "Failed to find ip for device: {}".format(device))
    return ip_match.group(1)


def assess_address_range(client, targets):
    """Test that two containers are in the same subnet as their host
    :param client: Juju client
    :param targets: machine IDs of machines to test
    :return: None; raises ValueError on failure
    """
    log.info('Assessing address range.')
    status = client.wait_for_started().status

    host_subnet_cache = {}

    for target in targets:
        log.info('Assessing address range for {}.'.format(target))
        host = target.split('/')[0]

        if host in host_subnet_cache:
            host_subnet = host_subnet_cache[host]
        else:
            host_address = private_address(client, host)
            host_subnet = find_network(client, host, host_address)
            host_subnet_cache[host] = host_subnet

        addr = status['machines'][host]['containers'][target]['dns-name']
        subnet = find_network(client, host, addr)
        if host_subnet != subnet:
            raise ValueError(
                '{} ({}) not on the same subnet as {} ({})'.format(
                    target, subnet, host, host_subnet))
        log.info('SUCCESS.')


def assess_internet_connection(client, targets):
    """Test that targets can ping their default route
    :param client: Juju client
    :param targets: machine IDs of machines to test
    :return: None; raises ValueError on failure
    """
    log.info('Assessing internet connection.')
    for target in targets:
        log.info("Assessing internet connection for {}".format(target))
        routes = ssh(client, target, 'ip route show')

        d = re.search(r'^default\s+via\s+([\d\.]+)\s+', routes, re.MULTILINE)
        if d:
            rc = client.juju('ssh', ('--proxy', target,
                                     'ping -c1 -q ' + d.group(1)), check=False)
            if rc != 0:
                raise ValueError('%s unable to ping default route' % target)
        else:
            raise ValueError("Default route not found")
        log.info("SUCCESS")


def _assessment_iteration(client, containers):
    """Run the network tests on this collection of machines and containers
    :param client: Juju client
    :param hosts: list of hosts of containers
    :param containers: list of containers to run tests between
    :return: None
    """
    assess_internet_connection(client, containers)
    assess_address_range(client, containers)
    assess_network_traffic(client, containers)


def _assess_container_networking(client, types, hosts, containers):
    """Run _assessment_iteration on all useful combinations of containers
    :param client: Juju client
    :param args: Parsed command line arguments
    :return: None
    """
    for container_type in types:
        # Test with two containers on the same host
        _assessment_iteration(client, containers[container_type][hosts[0]])

        # Now test with two containers on two different hosts
        test_containers = [
            containers[container_type][hosts[0]][0],
            containers[container_type][hosts[1]][0],
        ]
        _assessment_iteration(client, test_containers)

    if KVM_MACHINE in types and LXC_MACHINE in types:
        test_containers = [
            containers[LXC_MACHINE][hosts[0]][0],
            containers[KVM_MACHINE][hosts[0]][0],
        ]
        _assessment_iteration(client, test_containers)

        # Test with an LXC and a KVM on different machines
        test_containers = [
            containers[LXC_MACHINE][hosts[0]][0],
            containers[KVM_MACHINE][hosts[1]][0],
        ]
        _assessment_iteration(client, test_containers)


def get_uptime(client, host):
    uptime_pattern = re.compile(r'.*(\d+)')
    uptime_output = ssh(client, host, 'uptime -p')
    log.info('uptime -p: {}'.format(uptime_output))
    match = uptime_pattern.match(uptime_output)
    if match:
        return int(match.group(1))
    else:
        return 0


def assess_container_networking(client, types):
    """Runs _assess_address_allocation, reboots hosts, repeat.

    :param client: Juju client
    :param types: Container types to test
    :return: None
    """
    log.info("Setting up test.")
    hosts, containers = make_machines(client, types)
    status = client.wait_for_started().status
    log.info("Setup complete.")
    log.info("Test started.")

    _assess_container_networking(client, types, hosts, containers)

    # Reboot all hosted modelled machines then the controller.
    log.info("Instrumenting reboot of all machines.")
    try:
        for host in hosts:
            log.info("Restarting hosted machine: {}".format(host))
            client.juju(
                'run', ('--machine', host, 'sudo shutdown -r now'))
        client.juju('show-action-status', ('--name', 'juju-run'))

        log.info("Restarting controller machine 0")
        controller_client = client.get_controller_client()
        controller_status = controller_client.get_status()
        controller_host = controller_status.status['machines']['0']['dns-name']
        first_uptime = get_uptime(controller_client, '0')
        ssh(controller_client, '0', 'sudo shutdown -r now')
    except subprocess.CalledProcessError as e:
        logging.info(
            "Error running shutdown:\nstdout: %s\nstderr: %s",
            e.output, getattr(e, 'stderr', None))
        raise

    # Wait for the controller to shut down if it has not yet restarted.
    # This ensure the call to wait_for_started happens after each host
    # has restarted.
    second_uptime = get_uptime(controller_client, '0')
    if second_uptime > first_uptime:
        wait_for_port(controller_host, 22, closed=True, timeout=300)
    client.wait_for_started()

    # Once Juju is up it can take a little while before ssh responds.
    for host in hosts:
        hostname = status['machines'][host]['dns-name']
        wait_for_port(hostname, 22, timeout=240)
    log.info("Reboot complete and all hosts ready for retest.")

    _assess_container_networking(client, types, hosts, containers)
    log.info("PASS")


class _CleanedContext:

    def __init__(self, client):
        self.client = client
        self.return_code = None


@contextlib.contextmanager
def cleaned_bootstrap_context(bs_manager, args):
    ctx = _CleanedContext(bs_manager.client)
    client = ctx.client
    # TODO(gz): Having to manipulate client env state here to get the temp env
    #           is ugly, would ideally be captured in an explicit scope.
    update_env(client.env, bs_manager.temp_env_name, series=bs_manager.series,
               agent_url=bs_manager.agent_url,
               agent_stream=bs_manager.agent_stream, region=bs_manager.region)
    with bs_manager.top_context() as machines:
        bootstrap_required = True
        if args.clean_environment and clean_environment(client):
            bootstrap_required = False
        if bootstrap_required:
            with bs_manager.bootstrap_context(machines):
                client.bootstrap(args.upload_tools)
        with bs_manager.runtime_context(machines):
            yield ctx
        ctx.return_code = 0
        if args.clean_environment and not clean_environment(client):
            ctx.return_code = 1


def _get_container_types(client, machine_type):
    """
    Give list of container types to run testing against.

    If a machine_type was explictly specified, only test against those kind
    of containers. Otherwise, test all possible containers for the given
    juju version.
    """
    if machine_type:
        if machine_type not in client.supported_container_types:
            raise Exception(
                "no {} support on juju {}".format(machine_type,
                                                  client.version))
        return [machine_type]
    # TODO(gz): Only include LXC for 1.X clients
    types = list(client.supported_container_types)
    types.sort()
    return types


def main(argv=None):
    args = parse_args(argv)
    configure_logging(args.verbose)
    bs_manager = BootstrapManager.from_args(args)
    client = bs_manager.client
    machine_types = _get_container_types(client, args.machine_type)
    with cleaned_bootstrap_context(bs_manager, args) as ctx:
        assess_container_networking(bs_manager.client, machine_types)
    return ctx.return_code


if __name__ == '__main__':
    sys.exit(main())