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