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