1 # Copyright (c) 2014 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 """RDB Host objects. 6 7 RDBHost: Basic host object, capable of retrieving fields of a host that 8 correspond to columns of the host table. 9 10 RDBServerHostWrapper: Server side host adapters that help in making a raw 11 database host object more ameanable to the classes and functions in the rdb 12 and/or rdb clients. 13 14 RDBClientHostWrapper: Scheduler host proxy that converts host information 15 returned by the rdb into a client host object capable of proxying updates 16 back to the rdb. 17 """ 18 19 import logging 20 21 from django.core import exceptions as django_exceptions 22 23 import common 24 from autotest_lib.client.common_lib import utils 25 from autotest_lib.frontend.afe import rdb_model_extensions as rdb_models 26 from autotest_lib.frontend.afe import models as afe_models 27 from autotest_lib.scheduler import rdb_requests 28 from autotest_lib.scheduler import rdb_utils 29 from autotest_lib.server import constants 30 from autotest_lib.utils import labellib 31 32 try: 33 from chromite.lib import metrics 34 except ImportError: 35 metrics = utils.metrics_mock 36 37 38 class RDBHost(object): 39 """A python host object representing a django model for the host.""" 40 41 required_fields = set( 42 rdb_models.AbstractHostModel.get_basic_field_names() + ['id']) 43 44 45 def _update_attributes(self, new_attributes): 46 """Updates attributes based on an input dictionary. 47 48 Since reads are not proxied to the rdb this method caches updates to 49 the host tables as class attributes. 50 51 @param new_attributes: A dictionary of attributes to update. 52 """ 53 for name, value in new_attributes.iteritems(): 54 setattr(self, name, value) 55 56 57 def __init__(self, **kwargs): 58 if self.required_fields - set(kwargs.keys()): 59 raise rdb_utils.RDBException('Creating %s requires %s, got %s ' 60 % (self.__class__, self.required_fields, kwargs.keys())) 61 self._update_attributes(kwargs) 62 63 64 @classmethod 65 def get_required_fields_from_host(cls, host): 66 """Returns all required attributes of the host parsed into a dict. 67 68 Required attributes are defined as the attributes required to 69 create an RDBHost, and mirror the columns of the host table. 70 71 @param host: A host object containing all required fields as attributes. 72 """ 73 required_fields_map = {} 74 try: 75 for field in cls.required_fields: 76 required_fields_map[field] = getattr(host, field) 77 except AttributeError as e: 78 raise rdb_utils.RDBException('Required %s' % e) 79 required_fields_map['id'] = host.id 80 return required_fields_map 81 82 83 def wire_format(self): 84 """Returns information about this host object. 85 86 @return: A dictionary of fields representing the host. 87 """ 88 return RDBHost.get_required_fields_from_host(self) 89 90 91 class RDBServerHostWrapper(RDBHost): 92 """A host wrapper for the base host object. 93 94 This object contains all the attributes of the raw database columns, 95 and a few more that make the task of host assignment easier. It handles 96 the following duties: 97 1. Serialization of the host object and foreign keys 98 2. Conversion of label ids to label names, and retrieval of platform 99 3. Checking the leased bit/status of a host before leasing it out. 100 """ 101 102 def __init__(self, host): 103 """Create an RDBServerHostWrapper. 104 105 @param host: An instance of the Host model class. 106 """ 107 host_fields = RDBHost.get_required_fields_from_host(host) 108 super(RDBServerHostWrapper, self).__init__(**host_fields) 109 self.labels = rdb_utils.LabelIterator(host.labels.all()) 110 self.acls = [aclgroup.id for aclgroup in host.aclgroup_set.all()] 111 self.protection = host.protection 112 platform = host.platform() 113 # Platform needs to be a method, not an attribute, for 114 # backwards compatibility with the rest of the host model. 115 self.platform_name = platform.name if platform else None 116 self.shard_id = host.shard_id 117 118 119 def refresh(self, fields=None): 120 """Refresh the attributes on this instance. 121 122 @param fields: A list of fieldnames to refresh. If None 123 all the required fields of the host are refreshed. 124 125 @raises RDBException: If refreshing a field fails. 126 """ 127 # TODO: This is mainly required for cache correctness. If it turns 128 # into a bottleneck, cache host_ids instead of rdbhosts and rebuild 129 # the hosts once before leasing them out. The important part is to not 130 # trust the leased bit on a cached host. 131 fields = self.required_fields if not fields else fields 132 try: 133 refreshed_fields = afe_models.Host.objects.filter( 134 id=self.id).values(*fields)[0] 135 except django_exceptions.FieldError as e: 136 raise rdb_utils.RDBException('Couldn\'t refresh fields %s: %s' % 137 fields, e) 138 self._update_attributes(refreshed_fields) 139 140 141 def lease(self): 142 """Set the leased bit on the host object, and in the database. 143 144 @raises RDBException: If the host is already leased. 145 """ 146 self.refresh(fields=['leased']) 147 if self.leased: 148 raise rdb_utils.RDBException('Host %s is already leased' % 149 self.hostname) 150 self.leased = True 151 # TODO: Avoid leaking django out of rdb.QueryManagers. This is still 152 # preferable to calling save() on the host object because we're only 153 # updating/refreshing a single indexed attribute, the leased bit. 154 afe_models.Host.objects.filter(id=self.id).update(leased=self.leased) 155 156 157 def wire_format(self, unwrap_foreign_keys=True): 158 """Returns all information needed to scheduler jobs on the host. 159 160 @param unwrap_foreign_keys: If true this method will retrieve and 161 serialize foreign keys of the original host, which are stored 162 in the RDBServerHostWrapper as iterators. 163 164 @return: A dictionary of host information. 165 """ 166 host_info = super(RDBServerHostWrapper, self).wire_format() 167 168 if unwrap_foreign_keys: 169 host_info['labels'] = self.labels.get_label_names() 170 host_info['acls'] = self.acls 171 host_info['platform_name'] = self.platform_name 172 host_info['protection'] = self.protection 173 return host_info 174 175 176 class RDBClientHostWrapper(RDBHost): 177 """A client host wrapper for the base host object. 178 179 This wrapper is used whenever the queue entry needs direct access 180 to the host. 181 """ 182 # Shows more detailed status of what a DUT is doing. 183 _HOST_WORKING_METRIC = 'chromeos/autotest/dut_working' 184 # Shows which hosts are working. 185 _HOST_STATUS_METRIC = 'chromeos/autotest/dut_status' 186 # Maps duts to pools. 187 _HOST_POOL_METRIC = 'chromeos/autotest/dut_pool' 188 # Shows which scheduler machines are using a DUT. 189 _BOARD_SHARD_METRIC = 'chromeos/autotest/shard/board_presence' 190 191 192 def __init__(self, **kwargs): 193 194 # This class is designed to only check for the bare minimum 195 # attributes on a host, so if a client tries accessing an 196 # unpopulated foreign key it will result in an exception. Doing 197 # so makes it easier to add fields to the rdb host without 198 # updating all the clients. 199 super(RDBClientHostWrapper, self).__init__(**kwargs) 200 201 # TODO(beeps): Remove this once we transition to urls 202 from autotest_lib.scheduler import rdb 203 self.update_request_manager = rdb_requests.RDBRequestManager( 204 rdb_requests.UpdateHostRequest, rdb.update_hosts) 205 self.dbg_str = '' 206 self.metadata = {} 207 # We access labels for metrics generation below and it's awkward not 208 # knowing if labels were populated or not. 209 if not hasattr(self, 'labels'): 210 self.labels = () 211 212 213 214 def _update(self, payload): 215 """Send an update to rdb, save the attributes of the payload locally. 216 217 @param: A dictionary representing 'key':value of the update required. 218 219 @raises RDBException: If the update fails. 220 """ 221 logging.info('Host %s in %s updating %s through rdb on behalf of: %s ', 222 self.hostname, self.status, payload, self.dbg_str) 223 self.update_request_manager.add_request(host_id=self.id, 224 payload=payload) 225 for response in self.update_request_manager.response(): 226 if response: 227 raise rdb_utils.RDBException('Host %s unable to perform update ' 228 '%s through rdb on behalf of %s: %s', self.hostname, 229 payload, self.dbg_str, response) 230 super(RDBClientHostWrapper, self)._update_attributes(payload) 231 232 233 def get_metric_fields(self): 234 """Generate default set of fields to include for Monarch. 235 236 @return: Dictionary of default fields. 237 """ 238 fields = { 239 'dut_host_name': self.hostname, 240 'board': self.board, 241 'model': self._model, 242 } 243 244 return fields 245 246 247 def record_pool(self, fields): 248 """Report to Monarch current pool of dut. 249 250 @param fields Dictionary of fields to include. 251 """ 252 pool = '' 253 if len(self.pools) == 1: 254 pool = self.pools[0] 255 if pool in constants.Pools.MANAGED_POOLS: 256 pool = 'managed:' + pool 257 258 metrics.String(self._HOST_POOL_METRIC, 259 reset_after=True).set(pool, fields=fields) 260 261 262 def set_status(self, status): 263 """Proxy for setting the status of a host via the rdb. 264 265 @param status: The new status. 266 """ 267 # Update elasticsearch db. 268 self._update({'status': status}) 269 270 # Update Monarch. 271 fields = self.get_metric_fields() 272 self.record_pool(fields) 273 # As each device switches state, indicate that it is not in any 274 # other state. This allows Monarch queries to avoid double counting 275 # when additional points are added by the Window Align operation. 276 host_status_metric = metrics.Boolean( 277 self._HOST_STATUS_METRIC, reset_after=True) 278 for s in rdb_models.AbstractHostModel.Status.names: 279 fields['status'] = s 280 host_status_metric.set(s == status, fields=fields) 281 282 283 def record_working_state(self, working, timestamp): 284 """Report to Monarch whether we are working or broken. 285 286 @param working Host repair status. `True` means that the DUT 287 is up and expected to pass tests. `False` 288 means the DUT has failed repair and requires 289 manual intervention. 290 @param timestamp Time that the status was recorded. 291 """ 292 fields = self.get_metric_fields() 293 metrics.Boolean( 294 self._HOST_WORKING_METRIC, reset_after=True).set( 295 working, fields=fields) 296 metrics.Boolean(self._BOARD_SHARD_METRIC, reset_after=True).set( 297 True, 298 fields={ 299 'board': self.board, 300 'model': self._model, 301 }, 302 ) 303 self.record_pool(fields) 304 305 306 def update_field(self, fieldname, value): 307 """Proxy for updating a field on the host. 308 309 @param fieldname: The fieldname as a string. 310 @param value: The value to assign to the field. 311 """ 312 self._update({fieldname: value}) 313 314 315 def platform_and_labels(self): 316 """Get the platform and labels on this host. 317 318 @return: A tuple containing a list of label names and the platform name. 319 """ 320 platform = self.platform_name 321 labels = [label for label in self.labels if label != platform] 322 return platform, labels 323 324 325 def platform(self): 326 """Get the name of the platform of this host. 327 328 @return: A string representing the name of the platform. 329 """ 330 return self.platform_name 331 332 333 @property 334 def board(self): 335 """Get the names of the board of this host. 336 337 @return: A string of the name of the board, e.g., lumpy. Returns '' if 338 no board label is found. 339 """ 340 labels = labellib.LabelsMapping(self.labels) 341 return labels.get('board', '') 342 343 344 @property 345 def _model(self): 346 """Get the model this host. 347 348 @return: A string of the name of the model, e.g., robo360. Returns '' if 349 no model label is found. 350 """ 351 labels = labellib.LabelsMapping(self.labels) 352 return labels.get('model', '') 353 354 355 @property 356 def pools(self): 357 """Get the names of the pools of this host. 358 359 @return: A list of pool names that the host is assigned to. 360 """ 361 return [l[len(constants.Labels.POOL_PREFIX):] for l in self.labels 362 if l.startswith(constants.Labels.POOL_PREFIX)] 363 364 365 def get_object_dict(self, **kwargs): 366 """Serialize the attributes of this object into a dict. 367 368 This method is called through frontend code to get a serialized 369 version of this object. 370 371 @param kwargs: 372 extra_fields: Extra fields, outside the columns of a host table. 373 374 @return: A dictionary representing the fields of this host object. 375 """ 376 # TODO(beeps): Implement support for extra fields. Currently nothing 377 # requires them. 378 return self.wire_format() 379 380 381 def save(self): 382 """Save any local data a client of this host object might have saved. 383 384 Setting attributes on a model before calling its save() method is a 385 common django pattern. Most, if not all updates to the host happen 386 either through set status or update_field. Though we keep the internal 387 state of the RDBClientHostWrapper consistent through these updates 388 we need a bulk save method such as this one to save any attributes of 389 this host another model might have set on it before calling its own 390 save method. Eg: 391 task = ST.objects.get(id=12) 392 task.host.status = 'Running' 393 task.save() -> this should result in the hosts status changing to 394 Running. 395 396 Functions like add_host_to_labels will have to update this host object 397 differently, as that is another level of foreign key indirection. 398 """ 399 self._update(self.get_required_fields_from_host(self)) 400 401 402 def return_rdb_host(func): 403 """Decorator for functions that return a list of Host objects. 404 405 @param func: The decorated function. 406 @return: A functions capable of converting each host_object to a 407 rdb_hosts.RDBServerHostWrapper. 408 """ 409 def get_rdb_host(*args, **kwargs): 410 """Takes a list of hosts and returns a list of host_infos. 411 412 @param hosts: A list of hosts. Each host is assumed to contain 413 all the fields in a host_info defined above. 414 @return: A list of rdb_hosts.RDBServerHostWrappers, one per host, or an 415 empty list is no hosts were found.. 416 """ 417 hosts = func(*args, **kwargs) 418 return [RDBServerHostWrapper(host) for host in hosts] 419 return get_rdb_host 420
