blob: de2ef777b096de156ecd1feba56db7e3ee8af861 [file] [log] [blame] [edit]
# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""RDB request managers and requests.
RDB request managers: Call an rdb api_method with a list of RDBRequests, and
match the requests to the responses returned.
RDB Request classes: Used in conjunction with the request managers. Each class
defines the set of fields the rdb needs to fulfill the request, and a hashable
request object the request managers use to identify a response with a request.
"""
import collections
import common
from autotest_lib.scheduler import rdb_utils
class RDBRequestManager(object):
"""Base request manager for RDB requests.
Each instance of a request manager is associated with one request, and
one api call. All subclasses maintain a queue of unexecuted requests, and
and expose an api to add requests/retrieve the response for these requests.
"""
def __init__(self, request, api_call):
"""
@param request: A subclass of rdb_utls.RDBRequest. The manager can only
manage requests of one type.
@param api_call: The rdb api call this manager is expected to make.
A manager can only send requests of type request, to this api call.
"""
self.request = request
self.api_call = api_call
self.request_queue = []
def add_request(self, **kwargs):
"""Add an RDBRequest to the queue."""
self.request_queue.append(self.request(**kwargs).get_request())
def response(self):
"""Execute the api call and return a response for each request.
The order of responses is the same as the order of requests added
to the queue.
@yield: A response for each request added to the queue after the
last invocation of response.
"""
if not self.request_queue:
raise rdb_utils.RDBException('No requests. Call add_requests '
'with the appropriate kwargs, before calling response.')
result = self.api_call(self.request_queue)
requests = self.request_queue
self.request_queue = []
for request in requests:
yield result.get(request) if result else None
class BaseHostRequestManager(RDBRequestManager):
"""Manager for batched get requests on hosts."""
def response(self):
"""Yields a popped host from the returned host list."""
# As a side-effect of returning a host, this method also removes it
# from the list of hosts matched up against a request. Eg:
# hqes: [hqe1, hqe2, hqe3]
# client requests: [c_r1, c_r2, c_r3]
# generate requests in rdb: [r1 (c_r1 and c_r2), r2]
# and response {r1: [h1, h2], r2:[h3]}
# c_r1 and c_r2 need to get different hosts though they're the same
# request, because they're from different queue_entries.
for hosts in super(BaseHostRequestManager, self).response():
yield hosts.pop() if hosts else None
class RDBRequestMeta(type):
"""Metaclass for constructing rdb requests.
This meta class creates a read-only request template by combining the
request_arguments of all classes in the inheritence hierarchy into a
namedtuple.
"""
def __new__(cls, name, bases, dctn):
for base in bases:
try:
dctn['_request_args'].update(base._request_args)
except AttributeError:
pass
dctn['template'] = collections.namedtuple('template',
dctn['_request_args'])
return type.__new__(cls, name, bases, dctn)
class RDBRequest(object):
"""Base class for an rdb request.
All classes inheriting from RDBRequest will need to specify a list of
request_args necessary to create the request, and will in turn get a
request that the rdb understands.
"""
__metaclass__ = RDBRequestMeta
__slots__ = set(['_request_args', '_request'])
_request_args = set([])
def __init__(self, **kwargs):
for key,value in kwargs.iteritems():
try:
hash(value)
except TypeError as e:
raise rdb_utils.RDBException('All fields of a %s must be. '
'hashable %s: %s, %s failed this test.' %
(self.__class__, key, type(value), value))
try:
self._request = self.template(**kwargs)
except TypeError:
raise rdb_utils.RDBException('Creating %s requires args %s got %s' %
(self.__class__, self.template._fields, kwargs.keys()))
def get_request(self):
"""Returns a request that the rdb understands.
@return: A named tuple with all the fields necessary to make a request.
"""
return self._request
class HashableDict(dict):
"""A hashable dictionary.
This class assumes all values of the input dict are hashable.
"""
def __hash__(self):
return hash(tuple(sorted(self.items())))
class HostRequest(RDBRequest):
"""Basic request for information about a single host.
Eg: HostRequest(host_id=x): Will return all information about host x.
"""
_request_args = set(['host_id'])
class UpdateHostRequest(HostRequest):
"""Defines requests to update hosts.
Eg:
UpdateHostRequest(host_id=x, payload={'afe_hosts_col_name': value}):
Will update column afe_hosts_col_name with the given value, for
the given host_id.
@raises RDBException: If the input arguments don't contain the expected
fields to make the request, or are of the wrong type.
"""
_request_args = set(['payload'])
def __init__(self, **kwargs):
try:
kwargs['payload'] = HashableDict(kwargs['payload'])
except (KeyError, TypeError) as e:
raise rdb_utils.RDBException('Creating %s requires args %s got %s' %
(self.__class__, self.template._fields, kwargs.keys()))
super(UpdateHostRequest, self).__init__(**kwargs)
class AcquireHostRequest(HostRequest):
"""Defines requests to acquire hosts.
Eg:
AcquireHostRequest(host_id=None, deps=[d1, d2], acls=[a1, a2],
priority=None, parent_job_id=None): Will acquire and return a
host that matches the specified deps/acls.
AcquireHostRequest(host_id=x, deps=[d1, d2], acls=[a1, a2]) : Will
acquire and return host x, after checking deps/acls match.
@raises RDBException: If the the input arguments don't contain the expected
fields to make a request, or are of the wrong type.
"""
# TODO(beeps): Priority and parent_job_id shouldn't be a part of the
# core request.
_request_args = set(['priority', 'deps', 'acls', 'parent_job_id'])
def __init__(self, **kwargs):
try:
kwargs['deps'] = frozenset(kwargs['deps'])
kwargs['acls'] = frozenset(kwargs['acls'])
# parent_job_id defaults to NULL but always serializing it as an int
# fits the rdb's type assumptions. Note that job ids are 1 based.
if kwargs['parent_job_id'] is None:
kwargs['parent_job_id'] = 0
except (KeyError, TypeError) as e:
raise rdb_utils.RDBException('Creating %s requires args %s got %s' %
(self.__class__, self.template._fields, kwargs.keys()))
super(AcquireHostRequest, self).__init__(**kwargs)