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