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 5 """This class defines the Base Label classes.""" 6 7 8 import logging 9 10 import common 11 12 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers 13 14 15 def forever_exists_decorate(exists): 16 """ 17 Decorator for labels that should exist forever once applied. 18 19 We'll check if the label already exists on the host and return True if so. 20 Otherwise we'll check if the label should exist on the host. 21 22 @param exists: The exists method on the label class. 23 """ 24 def exists_wrapper(self, host): 25 """ 26 Wrapper around the label exists method. 27 28 @param self: The label object. 29 @param host: The host object to run methods on. 30 31 @returns True if the label already exists on the host, otherwise run 32 the exists method. 33 """ 34 info = host.host_info_store.get() 35 return (self._NAME in info.labels) or exists(self, host) 36 return exists_wrapper 37 38 39 class BaseLabel(object): 40 """ 41 This class contains the scaffolding for the host-specific labels. 42 43 @property _NAME String that is either the label returned or a prefix of a 44 generated label. 45 @property _LABEL_LIST List of label classes that this label generates its 46 own labels from. This class attribute is primarily 47 for the LabelRetriever class to figure out what 48 labels are generated from this label. In most cases, 49 the _NAME attribute gives us what we want, but in the 50 special case where a label class is actually a 51 collection of label classes, then this attribute 52 comes into play. For the example of 53 testbed_label.ADBDeviceLabels, that class is really a 54 collection of the adb devices' labels in that testbed 55 so _NAME won't cut it. Instead, we use _LABEL_LIST 56 to tell LabelRetriever what list of label classes we 57 are generating and thus are able to have a 58 comprehensive list of the generated labels. 59 """ 60 61 _NAME = None 62 _LABEL_LIST = [] 63 64 def generate_labels(self, host): 65 """ 66 Return the list of labels generated for the host. 67 68 @param host: The host object to check on. Not needed here for base case 69 but could be needed for subclasses. 70 71 @return a list of labels applicable to the host. 72 """ 73 return [self._NAME] 74 75 76 def exists(self, host): 77 """ 78 Checks the host if the label is applicable or not. 79 80 This method is geared for the type of labels that indicate if the host 81 has a feature (bluetooth, touchscreen, etc) and as such require 82 detection logic to determine if the label should be applicable to the 83 host or not. 84 85 @param host: The host object to check on. 86 """ 87 raise NotImplementedError('exists not implemented') 88 89 90 def get(self, host): 91 """ 92 Return the list of labels. 93 94 @param host: The host object to check on. 95 """ 96 if self.exists(host): 97 return self.generate_labels(host) 98 else: 99 return [] 100 101 102 def get_all_labels(self): 103 """ 104 Return all possible labels generated by this label class. 105 106 @returns a tuple of sets, the first set is for labels that are prefixes 107 like 'os:android'. The second set is for labels that are full 108 labels by themselves like 'bluetooth'. 109 """ 110 # Another subclass takes care of prefixed labels so this is empty. 111 prefix_labels = set() 112 full_labels_list = (self._NAME if isinstance(self._NAME, list) else 113 [self._NAME]) 114 full_labels = set(full_labels_list) 115 116 return prefix_labels, full_labels 117 118 119 class StringLabel(BaseLabel): 120 """ 121 This class represents a string label that is dynamically generated. 122 123 This label class is used for the types of label that are always 124 present and will return at least one label out of a list of possible labels 125 (listed in _NAME). It is required that the subclasses implement 126 generate_labels() since the label class will need to figure out which labels 127 to return. 128 129 _NAME must always be overridden by the subclass with all the possible 130 labels that this label detection class can return in order to allow for 131 accurate label updating. 132 """ 133 134 def generate_labels(self, host): 135 raise NotImplementedError('generate_labels not implemented') 136 137 138 def exists(self, host): 139 """Set to true since it is assumed the label is always applicable.""" 140 return True 141 142 143 class StringPrefixLabel(StringLabel): 144 """ 145 This class represents a string label that is dynamically generated. 146 147 This label class is used for the types of label that usually are always 148 present and indicate the os/board/etc type of the host. The _NAME property 149 will be prepended with a colon to the generated labels like so: 150 151 _NAME = 'os' 152 generate_label() returns ['android'] 153 154 The labels returned by this label class will be ['os:android']. 155 It is important that the _NAME attribute be overridden by the 156 subclass; otherwise, all labels returned will be prefixed with 'None:'. 157 """ 158 159 def get(self, host): 160 """Return the list of labels with _NAME prefixed with a colon. 161 162 @param host: The host object to check on. 163 """ 164 if self.exists(host): 165 return ['%s:%s' % (self._NAME, label) 166 for label in self.generate_labels(host)] 167 else: 168 return [] 169 170 171 def get_all_labels(self): 172 """ 173 Return all possible labels generated by this label class. 174 175 @returns a tuple of sets, the first set is for labels that are prefixes 176 like 'os:android'. The second set is for labels that are full 177 labels by themselves like 'bluetooth'. 178 """ 179 # Since this is a prefix label class, we only care about 180 # prefixed_labels. We'll need to append the ':' to the label name to 181 # make sure we only match on prefix labels. 182 full_labels = set() 183 prefix_labels = set(['%s:' % self._NAME]) 184 185 return prefix_labels, full_labels 186 187 188 class LabelRetriever(object): 189 """This class will assist in retrieving/updating the host labels.""" 190 191 def _populate_known_labels(self, label_list): 192 """Create a list of known labels that is created through this class.""" 193 for label_instance in label_list: 194 # If this instance has a list of label, recurse on that list. 195 if label_instance._LABEL_LIST: 196 self._populate_known_labels(label_instance._LABEL_LIST) 197 continue 198 199 prefixed_labels, full_labels = label_instance.get_all_labels() 200 self.label_prefix_names.update(prefixed_labels) 201 self.label_full_names.update(full_labels) 202 203 204 def __init__(self, label_list): 205 self._labels = label_list 206 # These two sets will contain the list of labels we can safely remove 207 # during the update_labels call. 208 self.label_full_names = set() 209 self.label_prefix_names = set() 210 211 212 def get_labels(self, host): 213 """ 214 Retrieve the labels for the host. 215 216 @param host: The host to get the labels for. 217 """ 218 labels = [] 219 for label in self._labels: 220 logging.info('checking label %s', label.__class__.__name__) 221 try: 222 labels.extend(label.get(host)) 223 except Exception as e: 224 logging.exception('error getting label %s: %s', 225 label.__class__.__name__, e) 226 return labels 227 228 229 def _is_known_label(self, label): 230 """ 231 Checks if the label is a label known to the label detection framework. 232 233 We only delete labels that we might have created earlier. There are 234 some labels we should not be removing (e.g. pool:bvt) that we 235 want to keep but won't be part of the new labels detected on the host. 236 To do that we compare the passed in label to our list of known labels 237 and if we get a match, we feel safe knowing we can remove the label. 238 Otherwise we leave that label alone since it was generated elsewhere. 239 240 @param label: The label to check if we want to skip or not. 241 242 @returns True to skip (which means to keep this label, False to remove. 243 """ 244 return (label in self.label_full_names or 245 any([label.startswith(p) for p in self.label_prefix_names])) 246 247 248 def update_labels(self, host): 249 """ 250 Retrieve the labels from the host and update if needed. 251 252 @param host: The host to update the labels for. 253 """ 254 # If we haven't yet grabbed our list of known labels, do so now. 255 if not self.label_full_names and not self.label_prefix_names: 256 self._populate_known_labels(self._labels) 257 258 afe = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10) 259 old_labels = set(host._afe_host.labels) 260 logging.info('existing labels: %s', old_labels) 261 known_labels = set([l for l in old_labels 262 if self._is_known_label(l)]) 263 new_labels = set(self.get_labels(host)) 264 265 # TODO(pprabhu) Replace this update logic using AfeHostInfoBackend. 266 # Remove old labels. 267 labels_to_remove = list(old_labels & (known_labels - new_labels)) 268 if labels_to_remove: 269 logging.info('removing labels: %s', labels_to_remove) 270 afe.run('host_remove_labels', id=host.hostname, 271 labels=labels_to_remove) 272 273 # Add in new labels that aren't already there. 274 labels_to_add = list(new_labels - old_labels) 275 if labels_to_add: 276 logging.info('adding labels: %s', labels_to_add) 277 afe.run('host_add_labels', id=host.hostname, labels=labels_to_add) 278
