server/hosts/afe_store.py

# Copyright 2017 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.

import logging

import common
from autotest_lib.frontend.afe.json_rpc import proxy as rpc_proxy
from autotest_lib.server.hosts import host_info
from autotest_lib.server.cros.dynamic_suite import frontend_wrappers

class AfeStore(host_info.CachingHostInfoStore):
    """Directly interact with the (given) AFE for host information."""

    _RETRYING_AFE_TIMEOUT_MIN = 5
    _RETRYING_AFE_RETRY_DELAY_SEC = 10

    def __init__(self, hostname, afe=None):
        """
        @param hostname: The name of the host for which we want to track host
                information.
        @param afe: A frontend.AFE object to make RPC calls. Will create one
                internally if None.
        """
        super(AfeStore, self).__init__()
        self._hostname = hostname
        self._afe = afe
        if self._afe is None:
            self._afe = frontend_wrappers.RetryingAFE(
                    timeout_min=self._RETRYING_AFE_TIMEOUT_MIN,
                    delay_sec=self._RETRYING_AFE_RETRY_DELAY_SEC)


    def _refresh_impl(self):
        """Obtains HostInfo directly from the AFE."""
        try:
            hosts = self._afe.get_hosts(hostname=self._hostname)
        except rpc_proxy.JSONRPCException as e:
            raise host_info.StoreError(e)

        if not hosts:
            raise host_info.StoreError('No hosts founds with hostname: %s' %
                                       self._hostname)

        if len(hosts) > 1:
            logging.warning(
                    'Found %d hosts with the name %s. Picking the first one.',
                    len(hosts), self._hostname)
        host = hosts[0]
        return host_info.HostInfo(host.labels, host.attributes)


    def _commit_impl(self, new_info):
        """Commits HostInfo back to the AFE.

        @param new_info: The new HostInfo to commit.
        """
        # TODO(pprabhu) crbug.com/680322
        # This method has a potentially malignent race condition. We obtain a
        # copy of HostInfo from the AFE and then add/remove labels / attribtes
        # based on that. If another user tries to commit it's changes in
        # parallel, we'll end up with corrupted labels / attributes.
        old_info = self._refresh_impl()
        self._remove_labels_on_afe(set(old_info.labels) - set(new_info.labels))
        self._add_labels_on_afe(set(new_info.labels) - set(old_info.labels))

        # TODO(pprabhu) Also commit attributes when we first replace a direct
        # AFE call for attribute update.
        if old_info.attributes != new_info.attributes:
            logging.warning(
                    'Updating attributes is currently not supported. '
                    'attributes update skipped. old attributes: %s, committed '
                    'attributes: %s',
                    old_info.attributes, new_info.attributes)


    def _remove_labels_on_afe(self, labels):
        """Requests the AFE to remove the given labels.

        @param labels: Remove these.
        """
        if not labels:
            return

        logging.debug('removing labels: %s', labels)
        try:
            self._afe.run('host_remove_labels', id=self._hostname,
                          labels=labels)
        except rpc_proxy.JSONRPCException as e:
            raise host_info.StoreError(e)


    def _add_labels_on_afe(self, labels):
        """Requests the AFE to add the given labels.

        @param labels: Add these.
        """
        if not labels:
            return

        logging.info('adding labels: %s', labels)
        try:
            self._afe.run('host_add_labels', id=self._hostname, labels=labels)
        except rpc_proxy.JSONRPCException as e:
            raise host_info.StoreError(e)