1# Copyright 2016 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 abc
6import copy
7import logging
8
9import common
10from autotest_lib.server.cros import provision
11
12
13class HostInfo(object):
14    """Holds label/attribute information about a host as understood by infra.
15
16    This class is the source of truth of label / attribute information about a
17    host for the test runner (autoserv) and the tests, *from the point of view
18    of the infrastructure*.
19
20    Typical usage:
21        store = AfeHostInfoStore(...)
22        host_info = store.get()
23        update_somehow(host_info)
24        store.commit(host_info)
25
26    Besides the @property listed below, the following rw variables are part of
27    the public API:
28        labels: The list of labels for this host.
29        attributes: The list of attributes for this host.
30    """
31
32    __slots__ = ['labels', 'attributes']
33
34    # Constants related to exposing labels as more semantic properties.
35    _BOARD_PREFIX = 'board'
36    _OS_PREFIX = 'os'
37    _POOL_PREFIX = 'pool'
38
39    def __init__(self, labels=None, attributes=None):
40        """
41        @param labels: (optional list) labels to set on the HostInfo.
42        @param attributes: (optional dict) attributes to set on the HostInfo.
43        """
44        self.labels = labels if labels is not None else []
45        self.attributes = attributes if attributes is not None else {}
46
47
48    @property
49    def build(self):
50        """Retrieve the current build for the host.
51
52        TODO(pprabhu) Make provision.py depend on this instead of the other way
53        around.
54
55        @returns The first build label for this host (if there are multiple).
56                None if no build label is found.
57        """
58        for label_prefix in [provision.CROS_VERSION_PREFIX,
59                            provision.ANDROID_BUILD_VERSION_PREFIX,
60                            provision.TESTBED_BUILD_VERSION_PREFIX]:
61            build_labels = self._get_stripped_labels_with_prefix(label_prefix)
62            if build_labels:
63                return build_labels[0]
64        return None
65
66
67    @property
68    def board(self):
69        """Retrieve the board label value for the host.
70
71        @returns: The (stripped) board label, or None if no label is found.
72        """
73        return self.get_label_value(self._BOARD_PREFIX)
74
75
76    @property
77    def os(self):
78        """Retrieve the os for the host.
79
80        @returns The os (str) or None if no os label exists. Returns the first
81                matching os if mutiple labels are found.
82        """
83        return self.get_label_value(self._OS_PREFIX)
84
85
86    @property
87    def pools(self):
88        """Retrieve the set of pools for the host.
89
90        @returns: set(str) of pool values.
91        """
92        return set(self._get_stripped_labels_with_prefix(self._POOL_PREFIX))
93
94
95    def get_label_value(self, prefix):
96        """Retrieve the value stored as a label with a well known prefix.
97
98        @param prefix: The prefix of the desired label.
99        @return: For the first label matching 'prefix:value', returns value.
100                Returns '' if no label matches the given prefix.
101        """
102        values = self._get_stripped_labels_with_prefix(prefix)
103        return values[0] if values else ''
104
105
106    def _get_stripped_labels_with_prefix(self, prefix):
107        """Search for labels with the prefix and remove the prefix.
108
109        e.g.
110            prefix = blah
111            labels = ['blah:a', 'blahb', 'blah:c', 'doo']
112            returns: ['a', 'c']
113
114        @returns: A list of stripped labels. [] in case of no match.
115        """
116        full_prefix = prefix + ':'
117        prefix_len = len(full_prefix)
118        return [label[prefix_len:] for label in self.labels
119                if label.startswith(full_prefix)]
120
121
122    def __str__(self):
123        return ('HostInfo [Labels: %s, Attributes: %s'
124                % (self.labels, self.attributes))
125
126
127class StoreError(Exception):
128    """Raised when a CachingHostInfoStore operation fails."""
129
130
131class CachingHostInfoStore(object):
132    """Abstract class to obtain and update host information from the infra.
133
134    This class describes the API used to retrieve host information from the
135    infrastructure. The actual, uncached implementation to obtain / update host
136    information is delegated to the concrete store classes.
137
138    We use two concrete stores:
139        AfeHostInfoStore: Directly obtains/updates the host information from
140                the AFE.
141        LocalHostInfoStore: Obtains/updates the host information from a local
142                file.
143    An extra store is provided for unittests:
144        InMemoryHostInfoStore: Just store labels / attributes in-memory.
145    """
146
147    __metaclass__ = abc.ABCMeta
148
149    def __init__(self):
150        self._private_cached_info = None
151
152
153    def get(self, force_refresh=False):
154        """Obtain (possibly cached) host information.
155
156        @param force_refresh: If True, forces the cached HostInfo to be
157                refreshed from the store.
158        @returns: A HostInfo object.
159        """
160        if force_refresh:
161            return self._get_uncached()
162
163        # |_cached_info| access is costly, so do it only once.
164        info = self._cached_info
165        if info is None:
166            return self._get_uncached()
167        return info
168
169
170    def commit(self, info):
171        """Update host information in the infrastructure.
172
173        @param info: A HostInfo object with the new information to set. You
174                should obtain a HostInfo object using the |get| or
175                |get_uncached| methods, update it as needed and then commit.
176        """
177        logging.debug('Committing HostInfo to store %s', self)
178        try:
179            self._commit_impl(info)
180            self._cached_info = info
181            logging.debug('HostInfo updated to: %s', info)
182        except Exception:
183            self._cached_info = None
184            raise
185
186
187    @abc.abstractmethod
188    def _refresh_impl(self):
189        """Actual implementation to refresh host_info from the store.
190
191        Concrete stores must implement this function.
192        @returns: A HostInfo object.
193        """
194        raise NotImplementedError
195
196
197    @abc.abstractmethod
198    def _commit_impl(self, host_info):
199        """Actual implementation to commit host_info to the store.
200
201        Concrete stores must implement this function.
202        @param host_info: A HostInfo object.
203        """
204        raise NotImplementedError
205
206
207    def _get_uncached(self):
208        """Obtain freshly synced host information.
209
210        @returns: A HostInfo object.
211        """
212        logging.debug('Refreshing HostInfo using store %s', self)
213        logging.debug('Old host_info: %s', self._cached_info)
214        try:
215            info = self._refresh_impl()
216            self._cached_info = info
217        except Exception:
218            self._cached_info = None
219            raise
220
221        logging.debug('New host_info: %s', info)
222        return info
223
224
225    @property
226    def _cached_info(self):
227        """Access the cached info, enforcing a deepcopy."""
228        return copy.deepcopy(self._private_cached_info)
229
230
231    @_cached_info.setter
232    def _cached_info(self, info):
233        """Update the cached info, enforcing a deepcopy.
234
235        @param info: The new info to update from.
236        """
237        self._private_cached_info = copy.deepcopy(info)
238
239
240class InMemoryHostInfoStore(CachingHostInfoStore):
241    """A simple store that gives unittests direct access to backing data.
242
243    Unittests can access the |info| attribute to obtain the backing HostInfo.
244    """
245
246    def __init__(self, info=None):
247        """Seed object with initial data.
248
249        @param info: Initial backing HostInfo object.
250        """
251        super(InMemoryHostInfoStore, self).__init__()
252        self.info = info if info is not None else HostInfo()
253
254
255    def _refresh_impl(self):
256        """Return a copy of the private HostInfo."""
257        return copy.deepcopy(self.info)
258
259
260    def _commit_impl(self, info):
261        """Copy HostInfo data to in-memory store.
262
263        @param info: The HostInfo object to commit.
264        """
265        self.info = copy.deepcopy(info)
266
267
268def get_store_from_machine(machine):
269    """Obtain the host_info_store object stuffed in the machine dict.
270
271    The machine argument to jobs can be a string (a hostname) or a dict because
272    of legacy reasons. If we can't get a real store, return a dummy.
273    """
274    if isinstance(machine, dict):
275        return machine['host_info_store']
276    else:
277        return InMemoryHostInfoStore()
278