scheduler/rdb_requests.py

# 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)