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 utilities. 6 7Do not import rdb or autotest modules here to avoid cyclic dependencies. 8""" 9 10import collections 11 12import common 13from autotest_lib.client.common_lib.cros.graphite import autotest_stats 14from autotest_lib.client.common_lib import priorities 15 16RDB_STATS_KEY = 'rdb' 17 18class RDBException(Exception): 19 """Generic RDB exception.""" 20 21 def wire_format(self, **kwargs): 22 """Convert the exception to a format better suited to an rpc response. 23 """ 24 return str(self) 25 26 27class CacheMiss(RDBException): 28 """Generic exception raised for a cache miss in the rdb.""" 29 pass 30 31 32class LabelIterator(object): 33 """An Iterator for labels. 34 35 Within the rdb any label/dependency comparisons are performed based on label 36 ids. However, the host object returned needs to contain label names instead. 37 This class returns label ids for iteration, but a list of all label names 38 when accessed through get_label_names. 39 """ 40 41 def __init__(self, labels): 42 self.labels = labels 43 44 45 def __iter__(self): 46 return iter(label.id for label in self.labels) 47 48 49 def get_label_names(self): 50 """Get all label names of the labels associated with this class. 51 52 @return: A list of label names. 53 """ 54 return [label.name for label in self.labels] 55 56 57class RequestAccountant(object): 58 """A helper class that count requests and manages min_duts requirement. 59 60 On initialization, this object takes a list of host requests. 61 It will batch the requests by grouping similar requests together 62 and generate a mapping from unique request-> count of the request. 63 It will also generates a mapping from suite_job_id -> min_duts. 64 65 RDB does a two-round of host aquisition. The first round attempts 66 to get min_duts for each suite. The second round attemps to satisfy 67 the rest requests. RDB calls get_min_duts and get_rest to 68 figure out how many duts it should attempt to get for a unique 69 request in the first and second round respectively. 70 71 Assume we have two distinct requests 72 R1 (parent_job_id: 10, need hosts: 2) 73 R2 (parent_job_id: 10, need hosts: 4) 74 And parent job P (job_id:10) has min dut requirement of 3. So we got 75 requests_to_counts = {R1: 2, R2: 4} 76 min_duts_map = {P: 3} 77 78 First round acquiring: 79 Call get_min_duts(R1) 80 return 2, because P hasn't reach its min dut limit (3) yet 81 requests_to_counts -> {R1: 2-2=0, R2: 4} 82 min_duts_map -> {P: 3-2=1} 83 84 Call get_min_duts(R2) 85 return 1, because although R2 needs 4 duts, P's min dut limit is now 1 86 requests_to_counts -> {R1: 0, R2: 4-1=3} 87 min_duts_map -> {P: 1-1=0} 88 89 Second round acquiring: 90 Call get_rest(R1): 91 return 0, requests_to_counts[R1] 92 Call get_rest(R2): 93 return 3, requests_to_counts[R2] 94 95 Note it is possible that in the first round acquiring, although 96 R1 requested 2 duts, it may only get 1 or None. However get_rest 97 doesn't need to care whether the first round succeeded or not, as 98 in the case when the first round failed, regardless how many duts 99 get_rest requests, it will not be fullfilled anyway. 100 """ 101 102 _gauge = autotest_stats.Gauge(RDB_STATS_KEY) 103 104 105 def __init__(self, host_requests): 106 """Initialize. 107 108 @param host_requests: A list of request to acquire hosts. 109 """ 110 self.requests_to_counts = {} 111 # The order matters, it determines which request got fullfilled first. 112 self.requests = [] 113 for request, count in self._batch_requests(host_requests): 114 self.requests.append(request) 115 self.requests_to_counts[request] = count 116 self.min_duts_map = dict( 117 (r.parent_job_id, r.suite_min_duts) 118 for r in self.requests_to_counts.iterkeys() if r.parent_job_id) 119 120 121 @classmethod 122 def _batch_requests(cls, requests): 123 """ Group similar requests, sort by priority and parent_job_id. 124 125 @param requests: A list or unsorted, unordered requests. 126 127 @return: A list of tuples of the form (request, number of occurances) 128 formed by counting the number of requests with the same acls/deps/ 129 priority in the input list of requests, and sorting by priority. 130 The order of this list ensures against priority inversion. 131 """ 132 sort_function = lambda request: (request[0].priority, 133 -request[0].parent_job_id) 134 return sorted(collections.Counter(requests).items(), key=sort_function, 135 reverse=True) 136 137 138 def get_min_duts(self, host_request): 139 """Given a distinct host request figure out min duts to request for. 140 141 @param host_request: A request. 142 @returns: The minimum duts that should be requested. 143 """ 144 parent_id = host_request.parent_job_id 145 count = self.requests_to_counts[host_request] 146 if parent_id: 147 min_duts = self.min_duts_map.get(parent_id, 0) 148 to_acquire = min(count, min_duts) 149 self.min_duts_map[parent_id] = max(0, min_duts - to_acquire) 150 else: 151 to_acquire = 0 152 self.requests_to_counts[host_request] -= to_acquire 153 return to_acquire 154 155 156 def get_duts(self, host_request): 157 """Return the number of duts host_request still need. 158 159 @param host_request: A request. 160 @returns: The number of duts need to be requested. 161 """ 162 return self.requests_to_counts[host_request] 163 164 165 def record_acquire_min_duts(cls, host_request, hosts_required, 166 acquired_host_count): 167 """Send stats to graphite about host acquisition. 168 169 @param host_request: A request. 170 @param hosts_required: Number of hosts required to satisfy request. 171 @param acquired_host_count: Number of host acquired. 172 """ 173 try: 174 priority = priorities.Priority.get_string(host_request.priority) 175 except ValueError: 176 return 177 key = ('min_duts_acquisition_rate.%s') % priority 178 cls._gauge.send(key, acquired_host_count/float(hosts_required)) 179