• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Lint as: python2, python3
2# Copyright 2016 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6from __future__ import absolute_import
7from __future__ import division
8from __future__ import print_function
9
10import abc
11import copy
12import json
13import logging
14
15import common
16from autotest_lib.server.cros import provision
17import six
18
19
20class HostInfo(object):
21    """Holds label/attribute information about a host as understood by infra.
22
23    This class is the source of truth of label / attribute information about a
24    host for the test runner (autoserv) and the tests, *from the point of view
25    of the infrastructure*.
26
27    Typical usage:
28        store = AfeHostInfoStore(...)
29        host_info = store.get()
30        update_somehow(host_info)
31        store.commit(host_info)
32
33    Besides the @property listed below, the following rw variables are part of
34    the public API:
35        labels: The list of labels for this host.
36        attributes: The list of attributes for this host.
37    """
38
39    __slots__ = ['labels', 'attributes', 'stable_versions']
40
41    # Constants related to exposing labels as more semantic properties.
42    _BOARD_PREFIX = 'board'
43    _MODEL_PREFIX = 'model'
44    # sku was already used for perf labeling, but it's a human readable
45    # string (gen'd from HWID) and not the raw sku value, so avoiding collision
46    # with device-sku instead.
47    _DEVICE_SKU_PREFIX = 'device-sku'
48    _BRAND_CODE_PREFIX = 'brand-code'
49    _OS_PREFIX = 'os'
50    _POOL_PREFIX = 'pool'
51    # stable version constants
52    _SV_CROS_KEY = "cros"
53    _SV_FAFT_KEY = "faft"
54    _SV_FIRMWARE_KEY = "firmware"
55    _SV_SERVO_CROS_KEY = "servo-cros"
56
57    _OS_VERSION_LABELS = (
58            provision.CROS_VERSION_PREFIX,
59            provision.CROS_ANDROID_VERSION_PREFIX,
60    )
61
62    _VERSION_LABELS = _OS_VERSION_LABELS + (
63            provision.FW_RO_VERSION_PREFIX,
64            provision.FW_RW_VERSION_PREFIX,
65    )
66
67    def __init__(self, labels=None, attributes=None, stable_versions=None):
68        """
69        @param labels: (optional list) labels to set on the HostInfo.
70        @param attributes: (optional dict) attributes to set on the HostInfo.
71        @param stable_versions: (optional dict) stable version information to set on the HostInfo.
72        """
73        self.labels = labels if labels is not None else []
74        self.attributes = attributes if attributes is not None else {}
75        self.stable_versions = stable_versions if stable_versions is not None else {}
76
77
78    @property
79    def build(self):
80        """Retrieve the current build for the host.
81
82        TODO(pprabhu) Make provision.py depend on this instead of the other way
83        around.
84
85        @returns The first build label for this host (if there are multiple).
86                None if no build label is found.
87        """
88        for label_prefix in self._OS_VERSION_LABELS:
89            build_labels = self._get_stripped_labels_with_prefix(label_prefix)
90            if build_labels:
91                return build_labels[0]
92        return None
93
94
95    @property
96    def board(self):
97        """Retrieve the board label value for the host.
98
99        @returns: The (stripped) board label, or the empty string if no
100        label is found.
101        """
102        return self.get_label_value(self._BOARD_PREFIX)
103
104
105    @property
106    def model(self):
107        """Retrieve the model label value for the host.
108
109        @returns: The (stripped) model label, or the empty string if no
110        label is found.
111        """
112        return self.get_label_value(self._MODEL_PREFIX)
113
114
115    @property
116    def device_sku(self):
117        """Retrieve the device_sku label value for the host.
118
119        @returns: The (stripped) device_sku label, or the empty string if no
120        label is found.
121        """
122        return self.get_label_value(self._DEVICE_SKU_PREFIX)
123
124    @property
125    def brand_code(self):
126        """Retrieve the brand_code label value for the host.
127
128        @returns: The (stripped) brand_code label, or the empty string if no
129        label is found.
130        """
131        return self.get_label_value(self._BRAND_CODE_PREFIX)
132
133    @property
134    def os(self):
135        """Retrieve the os for the host.
136
137        @returns The os (str) or the empty string if no os label
138                exists. Returns the first matching os if mutiple labels
139                are found.
140        """
141        return self.get_label_value(self._OS_PREFIX)
142
143
144    @property
145    def pools(self):
146        """Retrieve the set of pools for the host.
147
148        @returns: set(str) of pool values.
149        """
150        return set(self._get_stripped_labels_with_prefix(self._POOL_PREFIX))
151
152
153    @property
154    def cros_stable_version(self):
155        """Retrieve the cros stable version
156        """
157        return self.stable_versions.get(self._SV_CROS_KEY)
158
159    @property
160    def faft_stable_version(self):
161        """Retrieve the faft stable version
162        """
163        return self.stable_versions.get(self._SV_FAFT_KEY)
164
165    @property
166    def firmware_stable_version(self):
167        """Retrieve the firmware stable version
168        """
169        return self.stable_versions.get(self._SV_FIRMWARE_KEY)
170
171    @property
172    def servo_cros_stable_version(self):
173        """Retrieve the servo cros stable verion
174        """
175        return self.stable_versions.get(self._SV_SERVO_CROS_KEY)
176
177    def get_label_value(self, prefix):
178        """Retrieve the value stored as a label with a well known prefix.
179
180        @param prefix: The prefix of the desired label.
181        @return: For the first label matching 'prefix:value', returns value.
182                Returns '' if no label matches the given prefix.
183        """
184        values = self._get_stripped_labels_with_prefix(prefix)
185        return values[0] if values else ''
186
187    def has_label(self, name):
188        """Check if label is present.
189
190        @param name: The name of the desired label.
191        @return: bool, True if present.
192        """
193        for label in self.labels:
194            if label == name or label.startswith(name + ':'):
195                return True
196        return False
197
198    def clear_version_labels(self, version_prefix=None):
199        """Clear all or a particular version label(s) for the host.
200
201        @param version_prefix: The prefix label which needs to be cleared.
202                               If this is set to None, all version labels will
203                               be cleared.
204        """
205        version_labels = ([version_prefix] if version_prefix else
206                          self._VERSION_LABELS)
207        self.labels = [
208                label for label in self.labels if
209                not any(label.startswith(prefix + ':')
210                        for prefix in version_labels)]
211
212
213    def set_version_label(self, version_prefix, version):
214        """Sets the version label for the host.
215
216        If a label with version_prefix exists, this updates the value for that
217        label, else appends a new label to the end of the label list.
218
219        @param version_prefix: The prefix to use (without the infix ':').
220        @param version: The version label value to set.
221        """
222        full_prefix = _to_label_prefix(version_prefix)
223        new_version_label = full_prefix + version
224        for index, label in enumerate(self.labels):
225            if label.startswith(full_prefix):
226                self.labels[index] = new_version_label
227                return
228        else:
229            self.labels.append(new_version_label)
230
231
232    def _get_stripped_labels_with_prefix(self, prefix):
233        """Search for labels with the prefix and remove the prefix.
234
235        e.g.
236            prefix = blah
237            labels = ['blah:a', 'blahb', 'blah:c', 'doo']
238            returns: ['a', 'c']
239
240        @returns: A list of stripped labels. [] in case of no match.
241        """
242        full_prefix = prefix + ':'
243        prefix_len = len(full_prefix)
244        return [label[prefix_len:] for label in self.labels
245                if label.startswith(full_prefix)]
246
247
248    def __str__(self):
249        return ('%s[Labels: %s, Attributes: %s, StableVersions: %s]'
250                % (type(self).__name__, self.labels, self.attributes, self.stable_versions))
251
252
253    def __eq__(self, other):
254        if isinstance(other, type(self)):
255            return all([
256                    self.labels == other.labels,
257                    self.attributes == other.attributes,
258                    self.stable_versions == other.stable_versions,
259            ])
260        else:
261            return NotImplemented
262
263
264    def __ne__(self, other):
265        return not (self == other)
266
267
268class StoreError(Exception):
269    """Raised when a CachingHostInfoStore operation fails."""
270
271
272class CachingHostInfoStore(six.with_metaclass(abc.ABCMeta, object)):
273    """Abstract class to obtain and update host information from the infra.
274
275    This class describes the API used to retrieve host information from the
276    infrastructure. The actual, uncached implementation to obtain / update host
277    information is delegated to the concrete store classes.
278
279    We use two concrete stores:
280        AfeHostInfoStore: Directly obtains/updates the host information from
281                the AFE.
282        LocalHostInfoStore: Obtains/updates the host information from a local
283                file.
284    An extra store is provided for unittests:
285        InMemoryHostInfoStore: Just store labels / attributes in-memory.
286    """
287
288    def __init__(self):
289        self._private_cached_info = None
290
291
292    def get(self, force_refresh=False):
293        """Obtain (possibly cached) host information.
294
295        @param force_refresh: If True, forces the cached HostInfo to be
296                refreshed from the store.
297        @returns: A HostInfo object.
298        """
299        if force_refresh:
300            return self._get_uncached()
301
302        # |_cached_info| access is costly, so do it only once.
303        info = self._cached_info
304        if info is None:
305            return self._get_uncached()
306        return info
307
308
309    def commit(self, info):
310        """Update host information in the infrastructure.
311
312        @param info: A HostInfo object with the new information to set. You
313                should obtain a HostInfo object using the |get| or
314                |get_uncached| methods, update it as needed and then commit.
315        """
316        logging.debug('Committing HostInfo to store %s', self)
317        try:
318            self._commit_impl(info)
319            self._cached_info = info
320            logging.debug('HostInfo updated to: %s', info)
321        except Exception:
322            self._cached_info = None
323            raise
324
325
326    @abc.abstractmethod
327    def _refresh_impl(self):
328        """Actual implementation to refresh host_info from the store.
329
330        Concrete stores must implement this function.
331        @returns: A HostInfo object.
332        """
333        raise NotImplementedError
334
335
336    @abc.abstractmethod
337    def _commit_impl(self, host_info):
338        """Actual implementation to commit host_info to the store.
339
340        Concrete stores must implement this function.
341        @param host_info: A HostInfo object.
342        """
343        raise NotImplementedError
344
345
346    def _get_uncached(self):
347        """Obtain freshly synced host information.
348
349        @returns: A HostInfo object.
350        """
351        logging.debug('Refreshing HostInfo using store %s', self)
352        logging.debug('Old host_info: %s', self._cached_info)
353        try:
354            info = self._refresh_impl()
355            self._cached_info = info
356        except Exception:
357            self._cached_info = None
358            raise
359
360        logging.debug('New host_info: %s', info)
361        return info
362
363
364    @property
365    def _cached_info(self):
366        """Access the cached info, enforcing a deepcopy."""
367        return copy.deepcopy(self._private_cached_info)
368
369
370    @_cached_info.setter
371    def _cached_info(self, info):
372        """Update the cached info, enforcing a deepcopy.
373
374        @param info: The new info to update from.
375        """
376        self._private_cached_info = copy.deepcopy(info)
377
378
379class InMemoryHostInfoStore(CachingHostInfoStore):
380    """A simple store that gives unittests direct access to backing data.
381
382    Unittests can access the |info| attribute to obtain the backing HostInfo.
383    """
384
385    def __init__(self, info=None):
386        """Seed object with initial data.
387
388        @param info: Initial backing HostInfo object.
389        """
390        super(InMemoryHostInfoStore, self).__init__()
391        self.info = info if info is not None else HostInfo()
392
393
394    def __str__(self):
395        return '%s[%s]' % (type(self).__name__, self.info)
396
397    def _refresh_impl(self):
398        """Return a copy of the private HostInfo."""
399        return copy.deepcopy(self.info)
400
401
402    def _commit_impl(self, info):
403        """Copy HostInfo data to in-memory store.
404
405        @param info: The HostInfo object to commit.
406        """
407        self.info = copy.deepcopy(info)
408
409
410def get_store_from_machine(machine):
411    """Obtain the host_info_store object stuffed in the machine dict.
412
413    The machine argument to jobs can be a string (a hostname) or a dict because
414    of legacy reasons. If we can't get a real store, return a dummy.
415    """
416    if isinstance(machine, dict):
417        return machine['host_info_store']
418    else:
419        return InMemoryHostInfoStore()
420
421
422class DeserializationError(Exception):
423    """Raised when deserialization fails due to malformed input."""
424
425
426# Default serialzation version. This should be uprevved whenever a change to
427# HostInfo is backwards incompatible, i.e. we can no longer correctly
428# deserialize a previously serialized HostInfo. An example of such change is if
429# a field in the HostInfo object is dropped.
430_CURRENT_SERIALIZATION_VERSION = 1
431
432
433def json_serialize(info, file_obj, version=_CURRENT_SERIALIZATION_VERSION):
434    """Serialize the given HostInfo.
435
436    @param info: A HostInfo object to serialize.
437    @param file_obj: A file like object to serialize info into.
438    @param version: Use a specific serialization version. Should mostly use the
439            default.
440    """
441    info_json = {
442            'serializer_version': version,
443            'labels': info.labels,
444            'attributes': info.attributes,
445            'stable_versions': info.stable_versions,
446    }
447    return json.dump(info_json, file_obj, sort_keys=True, indent=4,
448                     separators=(',', ': '))
449
450
451def json_deserialize(file_obj):
452    """Deserialize a HostInfo from the given file.
453
454    @param file_obj: a file like object containing a json_serialized()ed
455            HostInfo.
456    @returns: The deserialized HostInfo object.
457    """
458    try:
459        deserialized_json = json.load(file_obj)
460    except ValueError as e:
461        raise DeserializationError(e)
462
463    try:
464        return HostInfo(deserialized_json['labels'],
465                        deserialized_json['attributes'],
466                        deserialized_json.get('stable_versions', {}))
467    except KeyError as e:
468        raise DeserializationError('Malformed serialized host_info: %r' % e)
469
470
471def _to_label_prefix(prefix):
472    """Ensure that prefix has the expected format for label prefixes.
473
474    @param prefix: The (str) prefix to sanitize.
475    @returns: The sanitized (str) prefix.
476    """
477    return prefix if prefix.endswith(':') else prefix + ':'
478