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