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 request managers and requests. 6 7 RDB request managers: Call an rdb api_method with a list of RDBRequests, and 8 match the requests to the responses returned. 9 10 RDB Request classes: Used in conjunction with the request managers. Each class 11 defines the set of fields the rdb needs to fulfill the request, and a hashable 12 request object the request managers use to identify a response with a request. 13 """ 14 15 import collections 16 17 import common 18 from autotest_lib.scheduler import rdb_utils 19 20 21 class RDBRequestManager(object): 22 """Base request manager for RDB requests. 23 24 Each instance of a request manager is associated with one request, and 25 one api call. All subclasses maintain a queue of unexecuted requests, and 26 and expose an api to add requests/retrieve the response for these requests. 27 """ 28 29 30 def __init__(self, request, api_call): 31 """ 32 @param request: A subclass of rdb_utls.RDBRequest. The manager can only 33 manage requests of one type. 34 @param api_call: The rdb api call this manager is expected to make. 35 A manager can only send requests of type request, to this api call. 36 """ 37 self.request = request 38 self.api_call = api_call 39 self.request_queue = [] 40 41 42 def add_request(self, **kwargs): 43 """Add an RDBRequest to the queue.""" 44 self.request_queue.append(self.request(**kwargs).get_request()) 45 46 47 def response(self): 48 """Execute the api call and return a response for each request. 49 50 The order of responses is the same as the order of requests added 51 to the queue. 52 53 @yield: A response for each request added to the queue after the 54 last invocation of response. 55 """ 56 if not self.request_queue: 57 raise rdb_utils.RDBException('No requests. Call add_requests ' 58 'with the appropriate kwargs, before calling response.') 59 60 result = self.api_call(self.request_queue) 61 requests = self.request_queue 62 self.request_queue = [] 63 for request in requests: 64 yield result.get(request) if result else None 65 66 67 class BaseHostRequestManager(RDBRequestManager): 68 """Manager for batched get requests on hosts.""" 69 70 71 def response(self): 72 """Yields a popped host from the returned host list.""" 73 74 # As a side-effect of returning a host, this method also removes it 75 # from the list of hosts matched up against a request. Eg: 76 # hqes: [hqe1, hqe2, hqe3] 77 # client requests: [c_r1, c_r2, c_r3] 78 # generate requests in rdb: [r1 (c_r1 and c_r2), r2] 79 # and response {r1: [h1, h2], r2:[h3]} 80 # c_r1 and c_r2 need to get different hosts though they're the same 81 # request, because they're from different queue_entries. 82 for hosts in super(BaseHostRequestManager, self).response(): 83 yield hosts.pop() if hosts else None 84 85 86 class RDBRequestMeta(type): 87 """Metaclass for constructing rdb requests. 88 89 This meta class creates a read-only request template by combining the 90 request_arguments of all classes in the inheritence hierarchy into a 91 namedtuple. 92 """ 93 def __new__(cls, name, bases, dctn): 94 for base in bases: 95 try: 96 dctn['_request_args'].update(base._request_args) 97 except AttributeError: 98 pass 99 dctn['template'] = collections.namedtuple('template', 100 dctn['_request_args']) 101 return type.__new__(cls, name, bases, dctn) 102 103 104 class RDBRequest(object): 105 """Base class for an rdb request. 106 107 All classes inheriting from RDBRequest will need to specify a list of 108 request_args necessary to create the request, and will in turn get a 109 request that the rdb understands. 110 """ 111 __metaclass__ = RDBRequestMeta 112 __slots__ = set(['_request_args', '_request']) 113 _request_args = set([]) 114 115 116 def __init__(self, **kwargs): 117 for key,value in kwargs.iteritems(): 118 try: 119 hash(value) 120 except TypeError as e: 121 raise rdb_utils.RDBException('All fields of a %s must be. ' 122 'hashable %s: %s, %s failed this test.' % 123 (self.__class__, key, type(value), value)) 124 try: 125 self._request = self.template(**kwargs) 126 except TypeError: 127 raise rdb_utils.RDBException('Creating %s requires args %s got %s' % 128 (self.__class__, self.template._fields, kwargs.keys())) 129 130 131 def get_request(self): 132 """Returns a request that the rdb understands. 133 134 @return: A named tuple with all the fields necessary to make a request. 135 """ 136 return self._request 137 138 139 class HashableDict(dict): 140 """A hashable dictionary. 141 142 This class assumes all values of the input dict are hashable. 143 """ 144 145 def __hash__(self): 146 return hash(tuple(sorted(self.items()))) 147 148 149 class HostRequest(RDBRequest): 150 """Basic request for information about a single host. 151 152 Eg: HostRequest(host_id=x): Will return all information about host x. 153 """ 154 _request_args = set(['host_id']) 155 156 157 class UpdateHostRequest(HostRequest): 158 """Defines requests to update hosts. 159 160 Eg: 161 UpdateHostRequest(host_id=x, payload={'afe_hosts_col_name': value}): 162 Will update column afe_hosts_col_name with the given value, for 163 the given host_id. 164 165 @raises RDBException: If the input arguments don't contain the expected 166 fields to make the request, or are of the wrong type. 167 """ 168 _request_args = set(['payload']) 169 170 171 def __init__(self, **kwargs): 172 try: 173 kwargs['payload'] = HashableDict(kwargs['payload']) 174 except (KeyError, TypeError) as e: 175 raise rdb_utils.RDBException('Creating %s requires args %s got %s' % 176 (self.__class__, self.template._fields, kwargs.keys())) 177 super(UpdateHostRequest, self).__init__(**kwargs) 178 179 180 class AcquireHostRequest(HostRequest): 181 """Defines requests to acquire hosts. 182 183 Eg: 184 AcquireHostRequest(host_id=None, deps=[d1, d2], acls=[a1, a2], 185 priority=None, parent_job_id=None): Will acquire and return a 186 host that matches the specified deps/acls. 187 AcquireHostRequest(host_id=x, deps=[d1, d2], acls=[a1, a2]) : Will 188 acquire and return host x, after checking deps/acls match. 189 190 @raises RDBException: If the the input arguments don't contain the expected 191 fields to make a request, or are of the wrong type. 192 """ 193 # TODO(beeps): Priority and parent_job_id shouldn't be a part of the 194 # core request. 195 _request_args = set(['priority', 'deps', 'preferred_deps', 'acls', 196 'parent_job_id', 'suite_min_duts']) 197 198 199 def __init__(self, **kwargs): 200 try: 201 kwargs['deps'] = frozenset(kwargs['deps']) 202 kwargs['preferred_deps'] = frozenset(kwargs['preferred_deps']) 203 kwargs['acls'] = frozenset(kwargs['acls']) 204 205 # parent_job_id defaults to NULL but always serializing it as an int 206 # fits the rdb's type assumptions. Note that job ids are 1 based. 207 if kwargs['parent_job_id'] is None: 208 kwargs['parent_job_id'] = 0 209 except (KeyError, TypeError) as e: 210 raise rdb_utils.RDBException('Creating %s requires args %s got %s' % 211 (self.__class__, self.template._fields, kwargs.keys())) 212 super(AcquireHostRequest, self).__init__(**kwargs) 213 214 215
