1# Copyright 2017 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import logging
6
7import common
8from autotest_lib.frontend.afe.json_rpc import proxy as rpc_proxy
9from autotest_lib.server.hosts import host_info
10from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
11
12class AfeStore(host_info.CachingHostInfoStore):
13    """Directly interact with the (given) AFE for host information."""
14
15    _RETRYING_AFE_TIMEOUT_MIN = 5
16    _RETRYING_AFE_RETRY_DELAY_SEC = 10
17
18    def __init__(self, hostname, afe=None):
19        """
20        @param hostname: The name of the host for which we want to track host
21                information.
22        @param afe: A frontend.AFE object to make RPC calls. Will create one
23                internally if None.
24        """
25        super(AfeStore, self).__init__()
26        self._hostname = hostname
27        self._afe = afe
28        if self._afe is None:
29            self._afe = frontend_wrappers.RetryingAFE(
30                    timeout_min=self._RETRYING_AFE_TIMEOUT_MIN,
31                    delay_sec=self._RETRYING_AFE_RETRY_DELAY_SEC)
32
33
34    def _refresh_impl(self):
35        """Obtains HostInfo directly from the AFE."""
36        try:
37            hosts = self._afe.get_hosts(hostname=self._hostname)
38        except rpc_proxy.JSONRPCException as e:
39            raise host_info.StoreError(e)
40
41        if not hosts:
42            raise host_info.StoreError('No hosts founds with hostname: %s' %
43                                       self._hostname)
44
45        if len(hosts) > 1:
46            logging.warning(
47                    'Found %d hosts with the name %s. Picking the first one.',
48                    len(hosts), self._hostname)
49        host = hosts[0]
50        return host_info.HostInfo(host.labels, host.attributes)
51
52
53    def _commit_impl(self, new_info):
54        """Commits HostInfo back to the AFE.
55
56        @param new_info: The new HostInfo to commit.
57        """
58        # TODO(pprabhu) crbug.com/680322
59        # This method has a potentially malignent race condition. We obtain a
60        # copy of HostInfo from the AFE and then add/remove labels / attribtes
61        # based on that. If another user tries to commit it's changes in
62        # parallel, we'll end up with corrupted labels / attributes.
63        old_info = self._refresh_impl()
64        self._remove_labels_on_afe(set(old_info.labels) - set(new_info.labels))
65        self._add_labels_on_afe(set(new_info.labels) - set(old_info.labels))
66
67        # TODO(pprabhu) Also commit attributes when we first replace a direct
68        # AFE call for attribute update.
69        if old_info.attributes != new_info.attributes:
70            logging.warning(
71                    'Updating attributes is currently not supported. '
72                    'attributes update skipped. old attributes: %s, committed '
73                    'attributes: %s',
74                    old_info.attributes, new_info.attributes)
75
76
77    def _remove_labels_on_afe(self, labels):
78        """Requests the AFE to remove the given labels.
79
80        @param labels: Remove these.
81        """
82        if not labels:
83            return
84
85        logging.debug('removing labels: %s', labels)
86        try:
87            self._afe.run('host_remove_labels', id=self._hostname,
88                          labels=labels)
89        except rpc_proxy.JSONRPCException as e:
90            raise host_info.StoreError(e)
91
92
93    def _add_labels_on_afe(self, labels):
94        """Requests the AFE to add the given labels.
95
96        @param labels: Add these.
97        """
98        if not labels:
99            return
100
101        logging.info('adding labels: %s', labels)
102        try:
103            self._afe.run('host_add_labels', id=self._hostname, labels=labels)
104        except rpc_proxy.JSONRPCException as e:
105            raise host_info.StoreError(e)
106