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