scheduler/rdb_hosts.py

# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""RDB Host objects.

RDBHost: Basic host object, capable of retrieving fields of a host that
correspond to columns of the host table.

RDBServerHostWrapper: Server side host adapters that help in making a raw
database host object more ameanable to the classes and functions in the rdb
and/or rdb clients.

RDBClientHostWrapper: Scheduler host proxy that converts host information
returned by the rdb into a client host object capable of proxying updates
back to the rdb.
"""

import logging

from django.core import exceptions as django_exceptions

import common
from autotest_lib.client.common_lib import utils
from autotest_lib.frontend.afe import rdb_model_extensions as rdb_models
from autotest_lib.frontend.afe import models as afe_models
from autotest_lib.scheduler import rdb_requests
from autotest_lib.scheduler import rdb_utils
from autotest_lib.server import constants
from autotest_lib.utils import labellib

try:
    from chromite.lib import metrics
except ImportError:
    metrics = utils.metrics_mock


class RDBHost(object):
    """A python host object representing a django model for the host."""

    required_fields = set(
            rdb_models.AbstractHostModel.get_basic_field_names() + ['id'])


    def _update_attributes(self, new_attributes):
        """Updates attributes based on an input dictionary.

        Since reads are not proxied to the rdb this method caches updates to
        the host tables as class attributes.

        @param new_attributes: A dictionary of attributes to update.
        """
        for name, value in new_attributes.iteritems():
            setattr(self, name, value)


    def __init__(self, **kwargs):
        if self.required_fields - set(kwargs.keys()):
            raise rdb_utils.RDBException('Creating %s requires %s, got %s '
                    % (self.__class__, self.required_fields, kwargs.keys()))
        self._update_attributes(kwargs)


    @classmethod
    def get_required_fields_from_host(cls, host):
        """Returns all required attributes of the host parsed into a dict.

        Required attributes are defined as the attributes required to
        create an RDBHost, and mirror the columns of the host table.

        @param host: A host object containing all required fields as attributes.
        """
        required_fields_map = {}
        try:
            for field in cls.required_fields:
                required_fields_map[field] = getattr(host, field)
        except AttributeError as e:
            raise rdb_utils.RDBException('Required %s' % e)
        required_fields_map['id'] = host.id
        return required_fields_map


    def wire_format(self):
        """Returns information about this host object.

        @return: A dictionary of fields representing the host.
        """
        return RDBHost.get_required_fields_from_host(self)


class RDBServerHostWrapper(RDBHost):
    """A host wrapper for the base host object.

    This object contains all the attributes of the raw database columns,
    and a few more that make the task of host assignment easier. It handles
    the following duties:
        1. Serialization of the host object and foreign keys
        2. Conversion of label ids to label names, and retrieval of platform
        3. Checking the leased bit/status of a host before leasing it out.
    """

    def __init__(self, host):
        """Create an RDBServerHostWrapper.

        @param host: An instance of the Host model class.
        """
        host_fields = RDBHost.get_required_fields_from_host(host)
        super(RDBServerHostWrapper, self).__init__(**host_fields)
        self.labels = rdb_utils.LabelIterator(host.labels.all())
        self.acls = [aclgroup.id for aclgroup in host.aclgroup_set.all()]
        self.protection = host.protection
        platform = host.platform()
        # Platform needs to be a method, not an attribute, for
        # backwards compatibility with the rest of the host model.
        self.platform_name = platform.name if platform else None
        self.shard_id = host.shard_id


    def refresh(self, fields=None):
        """Refresh the attributes on this instance.

        @param fields: A list of fieldnames to refresh. If None
            all the required fields of the host are refreshed.

        @raises RDBException: If refreshing a field fails.
        """
        # TODO: This is mainly required for cache correctness. If it turns
        # into a bottleneck, cache host_ids instead of rdbhosts and rebuild
        # the hosts once before leasing them out. The important part is to not
        # trust the leased bit on a cached host.
        fields = self.required_fields if not fields else fields
        try:
            refreshed_fields = afe_models.Host.objects.filter(
                    id=self.id).values(*fields)[0]
        except django_exceptions.FieldError as e:
            raise rdb_utils.RDBException('Couldn\'t refresh fields %s: %s' %
                    fields, e)
        self._update_attributes(refreshed_fields)


    def lease(self):
        """Set the leased bit on the host object, and in the database.

        @raises RDBException: If the host is already leased.
        """
        self.refresh(fields=['leased'])
        if self.leased:
            raise rdb_utils.RDBException('Host %s is already leased' %
                                         self.hostname)
        self.leased = True
        # TODO: Avoid leaking django out of rdb.QueryManagers. This is still
        # preferable to calling save() on the host object because we're only
        # updating/refreshing a single indexed attribute, the leased bit.
        afe_models.Host.objects.filter(id=self.id).update(leased=self.leased)


    def wire_format(self, unwrap_foreign_keys=True):
        """Returns all information needed to scheduler jobs on the host.

        @param unwrap_foreign_keys: If true this method will retrieve and
            serialize foreign keys of the original host, which are stored
            in the RDBServerHostWrapper as iterators.

        @return: A dictionary of host information.
        """
        host_info = super(RDBServerHostWrapper, self).wire_format()

        if unwrap_foreign_keys:
            host_info['labels'] = self.labels.get_label_names()
            host_info['acls'] = self.acls
            host_info['platform_name'] = self.platform_name
            host_info['protection'] = self.protection
        return host_info


class RDBClientHostWrapper(RDBHost):
    """A client host wrapper for the base host object.

    This wrapper is used whenever the queue entry needs direct access
    to the host.
    """
    # Shows more detailed status of what a DUT is doing.
    _HOST_WORKING_METRIC = 'chromeos/autotest/dut_working'
    # Shows which hosts are working.
    _HOST_STATUS_METRIC = 'chromeos/autotest/dut_status'
    # Maps duts to pools.
    _HOST_POOL_METRIC = 'chromeos/autotest/dut_pool'
    # Shows which scheduler machines are using a DUT.
    _BOARD_SHARD_METRIC = 'chromeos/autotest/shard/board_presence'


    def __init__(self, **kwargs):

        # This class is designed to only check for the bare minimum
        # attributes on a host, so if a client tries accessing an
        # unpopulated foreign key it will result in an exception. Doing
        # so makes it easier to add fields to the rdb host without
        # updating all the clients.
        super(RDBClientHostWrapper, self).__init__(**kwargs)

        # TODO(beeps): Remove this once we transition to urls
        from autotest_lib.scheduler import rdb
        self.update_request_manager = rdb_requests.RDBRequestManager(
                rdb_requests.UpdateHostRequest, rdb.update_hosts)
        self.dbg_str = ''
        self.metadata = {}
        # We access labels for metrics generation below and it's awkward not
        # knowing if labels were populated or not.
        if not hasattr(self, 'labels'):
            self.labels = ()



    def _update(self, payload):
        """Send an update to rdb, save the attributes of the payload locally.

        @param: A dictionary representing 'key':value of the update required.

        @raises RDBException: If the update fails.
        """
        logging.info('Host %s in %s updating %s through rdb on behalf of: %s ',
                     self.hostname, self.status, payload, self.dbg_str)
        self.update_request_manager.add_request(host_id=self.id,
                payload=payload)
        for response in self.update_request_manager.response():
            if response:
                raise rdb_utils.RDBException('Host %s unable to perform update '
                        '%s through rdb on behalf of %s: %s',  self.hostname,
                        payload, self.dbg_str, response)
        super(RDBClientHostWrapper, self)._update_attributes(payload)


    def get_metric_fields(self):
        """Generate default set of fields to include for Monarch.

        @return: Dictionary of default fields.
        """
        fields = {
            'dut_host_name': self.hostname,
            'board': self.board,
            'model': self._model,
        }

        return fields


    def record_pool(self, fields):
        """Report to Monarch current pool of dut.

        @param fields   Dictionary of fields to include.
        """
        pool = ''
        if len(self.pools) == 1:
            pool = self.pools[0]
        if pool in constants.Pools.MANAGED_POOLS:
            pool = 'managed:' + pool

        metrics.String(self._HOST_POOL_METRIC,
                       reset_after=True).set(pool, fields=fields)


    def set_status(self, status):
        """Proxy for setting the status of a host via the rdb.

        @param status: The new status.
        """
        # Update elasticsearch db.
        self._update({'status': status})

        # Update Monarch.
        fields = self.get_metric_fields()
        self.record_pool(fields)
        # As each device switches state, indicate that it is not in any
        # other state.  This allows Monarch queries to avoid double counting
        # when additional points are added by the Window Align operation.
        host_status_metric = metrics.Boolean(
                self._HOST_STATUS_METRIC, reset_after=True)
        for s in rdb_models.AbstractHostModel.Status.names:
            fields['status'] = s
            host_status_metric.set(s == status, fields=fields)


    def record_working_state(self, working, timestamp):
        """Report to Monarch whether we are working or broken.

        @param working    Host repair status. `True` means that the DUT
                          is up and expected to pass tests.  `False`
                          means the DUT has failed repair and requires
                          manual intervention.
        @param timestamp  Time that the status was recorded.
        """
        fields = self.get_metric_fields()
        metrics.Boolean(
                self._HOST_WORKING_METRIC, reset_after=True).set(
                        working, fields=fields)
        metrics.Boolean(self._BOARD_SHARD_METRIC, reset_after=True).set(
                True,
                fields={
                        'board': self.board,
                        'model': self._model,
                },
        )
        self.record_pool(fields)


    def update_field(self, fieldname, value):
        """Proxy for updating a field on the host.

        @param fieldname: The fieldname as a string.
        @param value: The value to assign to the field.
        """
        self._update({fieldname: value})


    def platform_and_labels(self):
        """Get the platform and labels on this host.

        @return: A tuple containing a list of label names and the platform name.
        """
        platform = self.platform_name
        labels = [label for label in self.labels if label != platform]
        return platform, labels


    def platform(self):
        """Get the name of the platform of this host.

        @return: A string representing the name of the platform.
        """
        return self.platform_name


    @property
    def board(self):
        """Get the names of the board of this host.

        @return: A string of the name of the board, e.g., lumpy. Returns '' if
                no board label is found.
        """
        labels = labellib.LabelsMapping(self.labels)
        return labels.get('board', '')


    @property
    def _model(self):
        """Get the model this host.

        @return: A string of the name of the model, e.g., robo360. Returns '' if
                no model label is found.
        """
        labels = labellib.LabelsMapping(self.labels)
        return labels.get('model', '')


    @property
    def pools(self):
        """Get the names of the pools of this host.

        @return: A list of pool names that the host is assigned to.
        """
        return [l[len(constants.Labels.POOL_PREFIX):] for l in self.labels
                if l.startswith(constants.Labels.POOL_PREFIX)]


    def get_object_dict(self, **kwargs):
        """Serialize the attributes of this object into a dict.

        This method is called through frontend code to get a serialized
        version of this object.

        @param kwargs:
            extra_fields: Extra fields, outside the columns of a host table.

        @return: A dictionary representing the fields of this host object.
        """
        # TODO(beeps): Implement support for extra fields. Currently nothing
        # requires them.
        return self.wire_format()


    def save(self):
        """Save any local data a client of this host object might have saved.

        Setting attributes on a model before calling its save() method is a
        common django pattern. Most, if not all updates to the host happen
        either through set status or update_field. Though we keep the internal
        state of the RDBClientHostWrapper consistent through these updates
        we need a bulk save method such as this one to save any attributes of
        this host another model might have set on it before calling its own
        save method. Eg:
            task = ST.objects.get(id=12)
            task.host.status = 'Running'
            task.save() -> this should result in the hosts status changing to
            Running.

        Functions like add_host_to_labels will have to update this host object
        differently, as that is another level of foreign key indirection.
        """
        self._update(self.get_required_fields_from_host(self))


def return_rdb_host(func):
    """Decorator for functions that return a list of Host objects.

    @param func: The decorated function.
    @return: A functions capable of converting each host_object to a
        rdb_hosts.RDBServerHostWrapper.
    """
    def get_rdb_host(*args, **kwargs):
        """Takes a list of hosts and returns a list of host_infos.

        @param hosts: A list of hosts. Each host is assumed to contain
            all the fields in a host_info defined above.
        @return: A list of rdb_hosts.RDBServerHostWrappers, one per host, or an
            empty list is no hosts were found..
        """
        hosts = func(*args, **kwargs)
        return [RDBServerHostWrapper(host) for host in hosts]
    return get_rdb_host