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"""This module manages translation between monitor_db and the rdb. """
6
7import common
8from autotest_lib.scheduler import rdb
9from autotest_lib.scheduler import rdb_hosts
10from autotest_lib.scheduler import rdb_requests
11from autotest_lib.server.cros import provision
12
13
14# Adapters for scheduler specific objects: Convert job information to a
15# format more ameanable to the rdb/rdb request managers.
16class JobQueryManager(object):
17    """A caching query manager for all job related information."""
18    def __init__(self, queue_entries, suite_min_duts=None):
19        """Initialize.
20
21        @param queue_entries: A list of HostQueueEntry objects.
22        @param suite_min_duts: A dictionary where the key is suite job id,
23                and the value is the value of 'suite_min_dut' in the suite's
24                job keyvals. It should cover all the suite jobs which
25                the jobs (associated with the queue_entries) belong to.
26        """
27        # TODO(beeps): Break this dependency on the host_query_manager,
28        # crbug.com/336934.
29        from autotest_lib.scheduler import query_managers
30        self.query_manager = query_managers.AFEHostQueryManager()
31        jobs = [queue_entry.job_id for queue_entry in queue_entries]
32        self._job_acls = self.query_manager._get_job_acl_groups(jobs)
33        self._job_deps = self.query_manager._get_job_dependencies(jobs)
34        self._labels = self.query_manager._get_labels(self._job_deps)
35        self._suite_min_duts = suite_min_duts or {}
36
37
38    def get_job_info(self, queue_entry):
39        """Extract job information from a queue_entry/host-scheduler.
40
41        @param queue_entry: The queue_entry for which we need job information.
42
43        @return: A dictionary representing job related information.
44        """
45        job_id = queue_entry.job_id
46        job_deps, job_preferred_deps = [], []
47        for dep in self._job_deps.get(job_id, []):
48           if not provision.is_for_special_action(self._labels[dep].name):
49               job_deps.append(dep)
50           elif provision.Provision.acts_on(self._labels[dep].name):
51               job_preferred_deps.append(dep)
52
53        job_acls = self._job_acls.get(job_id, [])
54        parent_id = queue_entry.job.parent_job_id
55        min_duts = self._suite_min_duts.get(parent_id, 0) if parent_id else 0
56
57        return {'deps': job_deps, 'acls': job_acls,
58                'preferred_deps': job_preferred_deps,
59                'host_id': queue_entry.host_id,
60                'parent_job_id': queue_entry.job.parent_job_id,
61                'priority': queue_entry.job.priority,
62                'suite_min_duts': min_duts}
63
64
65def acquire_hosts(queue_entries, suite_min_duts=None):
66    """Acquire hosts for the list of queue_entries.
67
68    The act of acquisition involves leasing a host from the rdb.
69
70    @param queue_entries: A list of queue_entries that need hosts.
71    @param suite_min_duts: A dictionary that maps suite job id to the minimum
72                           number of duts required.
73
74    @yield: An rdb_hosts.RDBClientHostWrapper for each host acquired on behalf
75        of a queue_entry, or None if a host wasn't found.
76
77    @raises RDBException: If something goes wrong making the request.
78    """
79    job_query_manager = JobQueryManager(queue_entries, suite_min_duts)
80    request_manager = rdb_requests.BaseHostRequestManager(
81            rdb_requests.AcquireHostRequest, rdb.rdb_host_request_dispatcher)
82    for entry in queue_entries:
83        request_manager.add_request(**job_query_manager.get_job_info(entry))
84
85    for host in request_manager.response():
86        yield (rdb_hosts.RDBClientHostWrapper(**host)
87               if host else None)
88
89
90def get_hosts(host_ids):
91    """Get information about the hosts with ids in host_ids.
92
93    get_hosts is different from acquire_hosts in that it is completely
94    oblivious to the leased state of a host.
95
96    @param host_ids: A list of host_ids.
97
98    @return: A list of rdb_hosts.RDBClientHostWrapper objects.
99
100    @raises RDBException: If something goes wrong in making the request.
101    """
102    request_manager = rdb_requests.BaseHostRequestManager(
103            rdb_requests.HostRequest, rdb.get_hosts)
104    for host_id in host_ids:
105        request_manager.add_request(host_id=host_id)
106
107    hosts = []
108    for host in request_manager.response():
109        hosts.append(rdb_hosts.RDBClientHostWrapper(**host)
110                     if host else None)
111    return hosts
112