1#pylint: disable-msg=C0111
2
3import cPickle
4import logging
5import os
6import time
7
8import common
9from autotest_lib.scheduler import drone_utility, email_manager
10from autotest_lib.client.bin import local_host
11from autotest_lib.client.common_lib import error, global_config, utils
12from autotest_lib.client.common_lib.cros.graphite import autotest_stats
13
14
15AUTOTEST_INSTALL_DIR = global_config.global_config.get_config_value('SCHEDULER',
16                                                 'drone_installation_directory')
17DEFAULT_CONTAINER_PATH = global_config.global_config.get_config_value(
18        'AUTOSERV', 'container_path')
19
20class DroneUnreachable(Exception):
21    """The drone is non-sshable."""
22    pass
23
24
25class _BaseAbstractDrone(object):
26    """
27    Attributes:
28    * allowed_users: set of usernames allowed to use this drone.  if None,
29            any user can use this drone.
30    """
31    def __init__(self, timestamp_remote_calls=True):
32        """Instantiate an abstract drone.
33
34        @param timestamp_remote_calls: If true, drone_utility is invoked with
35            the --call_time option and the current time. Currently this is only
36            used for testing.
37        """
38        self._calls = []
39        self.hostname = None
40        self.enabled = True
41        self.max_processes = 0
42        self.active_processes = 0
43        self.allowed_users = None
44        self._autotest_install_dir = AUTOTEST_INSTALL_DIR
45        self._host = None
46        self.timestamp_remote_calls = timestamp_remote_calls
47        # If drone supports server-side packaging. The property support_ssp will
48        # init self._support_ssp later.
49        self._support_ssp = None
50
51
52    def shutdown(self):
53        pass
54
55
56    @property
57    def _drone_utility_path(self):
58        return os.path.join(self._autotest_install_dir,
59                            'scheduler', 'drone_utility.py')
60
61
62    def used_capacity(self):
63        """Gets the capacity used by this drone
64
65        Returns a tuple of (percentage_full, -max_capacity). This is to aid
66        direct comparisons, so that a 0/10 drone is considered less heavily
67        loaded than a 0/2 drone.
68
69        This value should never be used directly. It should only be used in
70        direct comparisons using the basic comparison operators, or using the
71        cmp() function.
72        """
73        if self.max_processes == 0:
74            return (1.0, 0)
75        return (float(self.active_processes) / self.max_processes,
76                -self.max_processes)
77
78
79    def usable_by(self, user):
80        if self.allowed_users is None:
81            return True
82        return user in self.allowed_users
83
84
85    def _execute_calls_impl(self, calls):
86        if not self._host:
87            raise ValueError('Drone cannot execute calls without a host.')
88        drone_utility_cmd = self._drone_utility_path
89        if self.timestamp_remote_calls:
90            drone_utility_cmd = '%s --call_time %s' % (
91                    drone_utility_cmd, time.time())
92        logging.info("Running drone_utility on %s", self.hostname)
93        result = self._host.run('python %s' % drone_utility_cmd,
94                                stdin=cPickle.dumps(calls), stdout_tee=None,
95                                connect_timeout=300)
96        try:
97            return cPickle.loads(result.stdout)
98        except Exception: # cPickle.loads can throw all kinds of exceptions
99            logging.critical('Invalid response:\n---\n%s\n---', result.stdout)
100            raise
101
102
103    def _execute_calls(self, calls):
104        autotest_stats.Gauge('drone_execute_call_count').send(
105                    self.hostname.replace('.', '_'), len(calls))
106        return_message = self._execute_calls_impl(calls)
107        for warning in return_message['warnings']:
108            subject = 'Warning from drone %s' % self.hostname
109            logging.warning(subject + '\n' + warning)
110            email_manager.manager.enqueue_notify_email(subject, warning)
111        return return_message['results']
112
113
114    def get_calls(self):
115        """Returns the calls queued against this drone.
116
117        @return: A list of calls queued against the drone.
118        """
119        return self._calls
120
121
122    def call(self, method, *args, **kwargs):
123        return self._execute_calls(
124            [drone_utility.call(method, *args, **kwargs)])
125
126
127    def queue_call(self, method, *args, **kwargs):
128        self._calls.append(drone_utility.call(method, *args, **kwargs))
129
130
131    def clear_call_queue(self):
132        self._calls = []
133
134
135    def execute_queued_calls(self):
136        if not self._calls:
137            return
138        results = self._execute_calls(self._calls)
139        self.clear_call_queue()
140        return results
141
142
143    def set_autotest_install_dir(self, path):
144        pass
145
146
147    @property
148    def support_ssp(self):
149        """Check if the drone supports server-side packaging with container.
150
151        @return: True if the drone supports server-side packaging with container
152        """
153        if not self._host:
154            raise ValueError('Can not determine if drone supports server-side '
155                             'packaging before host is set.')
156        if self._support_ssp is None:
157            try:
158                # TODO(crbug.com/471316): We need a better way to check if drone
159                # supports container, and install/upgrade base container. The
160                # check of base container folder is not reliable and shall be
161                # obsoleted once that bug is fixed.
162                self._host.run('which lxc-start')
163                # Test if base container is setup.
164                base_container = os.path.join(DEFAULT_CONTAINER_PATH, 'base')
165                # SSP uses privileged containers, sudo access is required. If
166                # the process can't run sudo command without password, SSP can't
167                # work properly. sudo command option -n will avoid user input.
168                # If password is required, the command will fail and raise
169                # AutoservRunError exception.
170                self._host.run('sudo -n ls "%s"' %  base_container)
171                self._support_ssp = True
172            except (error.AutoservRunError, error.AutotestHostRunError):
173                # Local drone raises AutotestHostRunError, while remote drone
174                # raises AutoservRunError.
175                self._support_ssp = False
176        return self._support_ssp
177
178
179SiteDrone = utils.import_site_class(
180   __file__, 'autotest_lib.scheduler.site_drones',
181   '_SiteAbstractDrone', _BaseAbstractDrone)
182
183
184class _AbstractDrone(SiteDrone):
185    pass
186
187
188class _LocalDrone(_AbstractDrone):
189    def __init__(self, timestamp_remote_calls=True):
190        super(_LocalDrone, self).__init__(
191                timestamp_remote_calls=timestamp_remote_calls)
192        self.hostname = 'localhost'
193        self._host = local_host.LocalHost()
194        self._drone_utility = drone_utility.DroneUtility()
195
196
197    def send_file_to(self, drone, source_path, destination_path,
198                     can_fail=False):
199        if drone.hostname == self.hostname:
200            self.queue_call('copy_file_or_directory', source_path,
201                            destination_path)
202        else:
203            self.queue_call('send_file_to', drone.hostname, source_path,
204                            destination_path, can_fail)
205
206
207class _RemoteDrone(_AbstractDrone):
208    def __init__(self, hostname, timestamp_remote_calls=True):
209        super(_RemoteDrone, self).__init__(
210                timestamp_remote_calls=timestamp_remote_calls)
211        self.hostname = hostname
212        self._host = drone_utility.create_host(hostname)
213        if not self._host.is_up():
214            logging.error('Drone %s is unpingable, kicking out', hostname)
215            raise DroneUnreachable
216
217
218    def set_autotest_install_dir(self, path):
219        self._autotest_install_dir = path
220
221
222    def shutdown(self):
223        super(_RemoteDrone, self).shutdown()
224        self._host.close()
225
226
227    def send_file_to(self, drone, source_path, destination_path,
228                     can_fail=False):
229        if drone.hostname == self.hostname:
230            self.queue_call('copy_file_or_directory', source_path,
231                            destination_path)
232        elif isinstance(drone, _LocalDrone):
233            drone.queue_call('get_file_from', self.hostname, source_path,
234                             destination_path)
235        else:
236            self.queue_call('send_file_to', drone.hostname, source_path,
237                            destination_path, can_fail)
238
239
240def get_drone(hostname):
241    """
242    Use this factory method to get drone objects.
243    """
244    if hostname == 'localhost':
245        return _LocalDrone()
246    try:
247        return _RemoteDrone(hostname)
248    except DroneUnreachable:
249        return None
250