#!/usr/bin/env python3
# Copyright 2017 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Custom swarming triggering script.

This script does custom swarming triggering logic, to enable device affinity
for our bots, while lumping all trigger calls under one logical step.

For the perf use case of device affinity, this script now enables soft device
affinity.  This means that it tries to smartly allocate jobs to bots based
on what is currently alive and what bot the task was last triggered on,
preferring that last triggered bot if available.  If the
--multiple-trigger-configs flag is specified than this script overrides
the soft device affinity functionality in favor of the provided ids.

The algorithm is roughly the following:

Find eligible bots, healthy or not.
  * Query swarming for eligible bots based on the dimensions passed in
    on the swarming call.  Determine their health status based on
    is not quarantied and is not is_dead

Of the eligible bots determine what bot id to run the shard on.
(Implementation in _select_config_indices_with_soft_affinity)
  * First query swarming for the last task that ran that shard with
    given dimensions.  Assuming they are returned with most recent first.
  * Check if the bot id that ran that task is alive, if so trigger
    on that bot again.
  * If that bot isn't alive, allocate to another alive bot or if no
    other alive bots exist, trigger on the same dead one.

Scripts inheriting must have roughly the same command line interface as
swarming.py trigger. It modifies it in the following ways:

 * Intercepts the dump-json argument, and creates its own by combining the
   results from each trigger call.
 * Intercepts the dimensions from the swarming call and determines what bots
   are healthy based on the above device affinity algorithm, and triggers
 * Adds a tag to the swarming trigger job with the shard so we know the last
   bot that ran this shard.

This script is normally called from the swarming recipe module in tools/build.

"""

import argparse
import copy
import os
import subprocess
import sys
import logging
import random

import base_test_triggerer

SRC_DIR = os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
sys.path.append(os.path.join(SRC_DIR, 'tools', 'perf'))

# //tools/perf imports.
import generate_perf_sharding
from core import bot_platforms


class Bot(object):  # pylint: disable=useless-object-inheritance
    """Eligible bots to run the task."""

    def __init__(self, bot_id, is_alive):
        self._bot_id = bot_id
        self._is_alive = is_alive

    def id(self):
        return self._bot_id

    def is_alive(self):
        return self._is_alive

    def as_json_config(self):
        return {'id': self._bot_id}


class PerfDeviceTriggerer(base_test_triggerer.BaseTestTriggerer):
    def __init__(self, args, swarming_args):
        # pylint: disable=super-with-arguments
        super(PerfDeviceTriggerer, self).__init__()
        # pylint: enable=super-with-arguments
        self._sharded_query_failed = False

        if not args.multiple_trigger_configs:
            # Represents the list of current dimensions requested
            # by the parent swarming job.
            self._dimensions = _get_swarming_dimensions(swarming_args)

            # Store what swarming server we need and whether or not we need
            # to send down authentication with it
            self._swarming_server = _get_swarming_server(swarming_args)

            # Map of all existing bots in swarming that satisfy the current
            # set of dimensions indexed by bot id.
            # Note: this assumes perf bot dimensions are unique between
            # configurations.
            self._eligible_bots_by_ids = (
                self._query_swarming_for_eligible_bot_configs(self._dimensions)
            )

        if args.multiple_dimension_script_verbose:
            logging.basicConfig(level=logging.DEBUG)

    def generate_shard_map(self, args, buildername, selected_config):
        shard_map = None
        num_of_shards = len(selected_config)
        builder = bot_platforms.find_bot_platform(buildername)
        if args.use_dynamic_shards and builder and num_of_shards:
            logging.info(
                'Generating dynamic shardmap for builder: %s with %d shards',
                buildername,
                num_of_shards,
            )
            shard_map = generate_perf_sharding.GenerateShardMap(
                builder=builder, num_of_shards=num_of_shards
            )
            for shard_index, bot_index in selected_config:
                bot_id = self._bot_configs[bot_index]['id']
                shard_map['extra_infos']['bot #%s' % shard_index] = bot_id
        return shard_map

    def append_additional_args(self, args, shard_index):
        # Append a tag to the swarming task with the shard number
        # so we can query for the last bot that ran a specific shard.
        tag = 'shard:%d' % shard_index
        shard_tag = ['--tag', tag]
        # Need to append this before the dash if present so it gets fed to
        # the swarming task itself.
        if '--' in args:
            dash_ind = args.index('--')
            return args[:dash_ind] + shard_tag + args[dash_ind:]
        return args + shard_tag

    def parse_bot_configs(self, args):
        if args.multiple_trigger_configs:
            # pylint: disable=super-with-arguments
            super(PerfDeviceTriggerer, self).parse_bot_configs(args)
            # pylint: enable=super-with-arguments
        else:
            self._bot_configs = []
            # For each eligible bot, append the dimension
            # to the eligible bot_configs
            for _, bot in self._eligible_bots_by_ids.items():
                self._bot_configs.append(bot.as_json_config())

    def select_config_indices(self, args):
        if args.multiple_trigger_configs:
            configs = []
            # If specific bot ids were passed in, we want to trigger a job for
            # every valid config regardless of health status since
            # each config represents exactly one bot in the perf swarming pool.
            for index in range(
                len(base_test_triggerer.indices_to_trigger(args))
            ):
                configs.append((index, index))
        if args.use_dynamic_shards:
            return self._select_config_indices_with_dynamic_sharding()
        return self._select_config_indices_with_soft_affinity(args)

    def _select_config_indices_with_dynamic_sharding(self):
        alive_bot_ids = [
            bot_id
            for bot_id, b in self._eligible_bots_by_ids.items()
            if b.is_alive()
        ]
        trigger_count = len(alive_bot_ids)

        indexes = list(range(trigger_count))
        random.shuffle(indexes)
        selected_config = [
            (indexes[i], self._find_bot_config_index(alive_bot_ids[i]))
            for i in range(trigger_count)
        ]
        selected_config.sort()

        for shard_index, bot_index in selected_config:
            logging.info(
                'Shard %d\n\tBot: %s',
                shard_index,
                self._bot_configs[bot_index]['id'],
            )

        return selected_config

    def _select_config_indices_with_soft_affinity(self, args):
        trigger_count = len(base_test_triggerer.indices_to_trigger(args))
        # First make sure the number of shards doesn't exceed the
        # number of eligible bots. This means there is a config error somewhere.
        if trigger_count > len(self._eligible_bots_by_ids):
            _print_device_affinity_info(
                {}, {}, self._eligible_bots_by_ids, trigger_count
            )
            raise ValueError(
                'Not enough available machines exist in swarming '
                'pool.  Shards requested (%d) exceeds available bots '
                '(%d).' % (trigger_count, len(self._eligible_bots_by_ids))
            )

        shard_to_bot_assignment_map = {}
        unallocated_bots_by_ids = copy.deepcopy(self._eligible_bots_by_ids)
        for shard_index in base_test_triggerer.indices_to_trigger(args):
            bot_id = self._query_swarming_for_last_shard_id(shard_index)
            if bot_id and bot_id in unallocated_bots_by_ids:
                bot = unallocated_bots_by_ids[bot_id]
                shard_to_bot_assignment_map[shard_index] = bot
                unallocated_bots_by_ids.pop(bot_id)
            else:
                shard_to_bot_assignment_map[shard_index] = None

        # Maintain the current map for debugging purposes
        existing_shard_bot_to_shard_map = copy.deepcopy(
            shard_to_bot_assignment_map
        )
        # Now create sets of remaining healthy and bad bots
        unallocated_healthy_bots = {
            b for b in unallocated_bots_by_ids.values() if b.is_alive()
        }
        unallocated_bad_bots = {
            b for b in unallocated_bots_by_ids.values() if not b.is_alive()
        }

        # Try assigning healthy bots for new shards first.
        for shard_index, bot in sorted(shard_to_bot_assignment_map.items()):
            if not bot and unallocated_healthy_bots:
                shard_to_bot_assignment_map[shard_index] = (
                    unallocated_healthy_bots.pop()
                )
                logging.info(
                    'First time shard %d has been triggered', shard_index
                )
            elif not bot:
                shard_to_bot_assignment_map[shard_index] = (
                    unallocated_bad_bots.pop()
                )

        # Handle the rest of shards that were assigned dead bots:
        for shard_index, bot in sorted(shard_to_bot_assignment_map.items()):
            if not bot.is_alive() and unallocated_healthy_bots:
                dead_bot = bot
                healthy_bot = unallocated_healthy_bots.pop()
                shard_to_bot_assignment_map[shard_index] = healthy_bot
                logging.info(
                    'Device affinity broken for shard #%d. bot %s is dead,'
                    ' new mapping to bot %s',
                    shard_index,
                    dead_bot.id(),
                    healthy_bot.id(),
                )

        # Now populate the indices into the bot_configs array
        selected_configs = []
        for shard_index in base_test_triggerer.indices_to_trigger(args):
            selected_configs.append(
                (
                    shard_index,
                    self._find_bot_config_index(
                        shard_to_bot_assignment_map[shard_index].id()
                    ),
                )
            )
        _print_device_affinity_info(
            shard_to_bot_assignment_map,
            existing_shard_bot_to_shard_map,
            self._eligible_bots_by_ids,
            trigger_count,
        )
        return selected_configs

    def _query_swarming_for_eligible_bot_configs(self, dimensions):
        """Query Swarming to figure out which bots are available.

        Returns: a dictionary in which the keys are the bot id and
        the values are Bot object that indicate the health status
        of the bots.
        """

        query_result = self.list_bots(dimensions, server=self._swarming_server)
        perf_bots = {}
        for bot in query_result:
            # Device maintenance is usually quick, and we can wait for it to
            # finish. However, if the device is too hot, it can take a long time
            # for it to cool down, so check for 'Device temperature' in
            # maintenance_msg.
            alive = (
                not bot.get('is_dead')
                and not bot.get('quarantined')
                and 'Device temperature' not in bot.get('maintenance_msg', '')
            )
            perf_bots[bot['bot_id']] = Bot(bot['bot_id'], alive)
        return perf_bots

    def _find_bot_config_index(self, bot_id):
        # Find the index into the bot_config map that
        # maps to the bot id in question
        for i, dimensions in enumerate(self._bot_configs):
            if dimensions['id'] == bot_id:
                return i
        return None

    def _query_swarming_for_last_shard_id(self, shard_index):
        """Per shard, query swarming for the last bot that ran the task.

        Example: swarming.py query -S server-url.com --limit 1 \\
          'tasks/list?tags=os:Windows&tags=pool:chrome.tests.perf&tags=shard:12'
        """
        values = ['%s:%s' % (k, v) for k, v in self._dimensions.items()]
        values.sort()

        # Append the shard as a tag
        values_with_shard = list(values)
        values_with_shard.append('%s:%s' % ('shard', str(shard_index)))
        values_with_shard.sort()

        # TODO(eyaich): For now we are ignoring the state of the returned
        # task (ie completed, timed_out, bot_died, etc) as we are just
        # answering the question "What bot did we last trigger this shard on?"
        # Evaluate if this is the right decision going forward.

        # Query for the last task that ran with these dimensions and this shard.
        # Query with the shard param first. This will sometimes time out for
        # queries we've never done before, so try querying without it if that
        # happens.
        try:
            if not self._sharded_query_failed:
                tasks = self.list_tasks(
                    values_with_shard, limit='1', server=self._swarming_server
                )
        except subprocess.CalledProcessError:
            self._sharded_query_failed = True
        if self._sharded_query_failed:
            tasks = self.list_tasks(
                values, limit='1', server=self._swarming_server
            )

        if tasks:
            # We queried with a limit of 1 so we could only get back
            # the most recent which is what we care about.
            task = tasks[0]
            if 'bot_id' in task:
                return task['bot_id']
            for tag in task['tags']:
                if tag.startswith('id:'):
                    return tag[len('id:') :]
        # No eligible shard for this bot
        return None


def _print_device_affinity_info(new_map, existing_map, health_map, num_shards):
    logging.info('')
    for shard_index in range(num_shards):
        existing = existing_map.get(shard_index, None)
        new = new_map.get(shard_index, None)
        existing_id = ''
        if existing:
            existing_id = existing.id()
        new_id = ''
        if new:
            new_id = new.id()
        logging.info(
            'Shard %d\n\tprevious: %s\n\tnew: %s',
            shard_index,
            existing_id,
            new_id,
        )

    healthy_bots = []
    dead_bots = []
    for _, b in health_map.items():
        if b.is_alive():
            healthy_bots.append(b.id())
        else:
            dead_bots.append(b.id())
    logging.info('Shards needed: %d', num_shards)
    logging.info(
        'Total bots (dead + healthy): %d', len(dead_bots) + len(healthy_bots)
    )
    logging.info('Healthy bots, %d: %s', len(healthy_bots), healthy_bots)
    logging.info('Dead Bots, %d: %s', len(dead_bots), dead_bots)
    logging.info('')


def _get_swarming_dimensions(args):
    dimensions = {}
    for i in range(len(args) - 2):
        if args[i] == '--dimension':
            dimensions[args[i + 1]] = args[i + 2]
    return dimensions


def _get_swarming_server(args):
    for i, argument in enumerate(args):
        if '--swarming' in argument:
            server = args[i + 1]
            slashes_index = server.index('//') + 2
            # Strip out the protocol
            return server[slashes_index:]
    return None


def main():
    logging.basicConfig(
        level=logging.INFO,
        format='(%(levelname)s) %(asctime)s pid=%(process)d'
        '  %(module)s.%(funcName)s:%(lineno)d  %(message)s',
    )
    # Setup args for common contract of base class
    parser = base_test_triggerer.BaseTestTriggerer.setup_parser_contract(
        argparse.ArgumentParser(description=__doc__)
    )
    parser.add_argument(
        '--use-dynamic-shards',
        action='store_true',
        required=False,
        help='Ignore --shards and the existing shard map. Will '
        'generate a shard map at run time and use as much '
        'device as possible.',
    )
    args, remaining = parser.parse_known_args()

    triggerer = PerfDeviceTriggerer(args, remaining)
    return triggerer.trigger_tasks(args, remaining)


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