# 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