1# Copyright (c) 2013 The Chromium 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
5import logging
6import random
7
8from time import sleep
9
10import common
11from autotest_lib.client.common_lib import error
12from autotest_lib.client.common_lib import utils
13from autotest_lib.server.cros.ap_configurators import \
14    ap_configurator_factory
15from autotest_lib.client.common_lib.cros.network import ap_constants
16from autotest_lib.server.cros.ap_configurators import ap_cartridge
17
18
19# Max number of retry attempts to lock an ap.
20MAX_RETRIES = 3
21
22
23class ApLocker(object):
24    """Object to keep track of AP lock state.
25
26    @attribute configurator: an APConfigurator object.
27    @attribute to_be_locked: a boolean, True iff ap has not been locked.
28    @attribute retries: an integer, max number of retry attempts to lock ap.
29    """
30
31
32    def __init__(self, configurator, retries):
33        """Initialize.
34
35        @param configurator: an APConfigurator object.
36        @param retries: an integer, max number of retry attempts to lock ap.
37        """
38        self.configurator = configurator
39        self.to_be_locked = True
40        self.retries = retries
41
42
43    def __repr__(self):
44        """@return class name, ap host name, lock status and retries."""
45        return 'class: %s, host name: %s, to_be_locked = %s, retries = %d' % (
46                self.__class__.__name__,
47                self.configurator.host_name,
48                self.to_be_locked,
49                self.retries)
50
51
52def construct_ap_lockers(ap_spec, retries, hostname_matching_only=False,
53                         ap_test_type=ap_constants.AP_TEST_TYPE_CHAOS):
54    """Convert APConfigurator objects to ApLocker objects for locking.
55
56    @param ap_spec: an APSpec object
57    @param retries: an integer, max number of retry attempts to lock ap.
58    @param hostname_matching_only: a boolean, if True matching against
59                                   all other APSpec parameters is not
60                                   performed.
61    @param ap_test_type: Used to determine which type of test we're
62                         currently running (Chaos vs Clique).
63
64    @return a list of ApLocker objects.
65    """
66    ap_lockers_list = []
67    factory = ap_configurator_factory.APConfiguratorFactory(ap_test_type,
68                                                            ap_spec)
69    if hostname_matching_only:
70        for ap in factory.get_aps_by_hostnames(ap_spec.hostnames):
71            ap_lockers_list.append(ApLocker(ap, retries))
72    else:
73        for ap in factory.get_ap_configurators_by_spec(ap_spec):
74            ap_lockers_list.append(ApLocker(ap, retries))
75
76    if not len(ap_lockers_list):
77        raise error.TestError('Found no matching APs to test against.')
78
79    logging.debug('Found %d APs', len(ap_lockers_list))
80    return ap_lockers_list
81
82
83class ApBatchLocker(object):
84    """Object to lock/unlock an APConfigurator.
85
86    @attribute SECONDS_TO_SLEEP: an integer, number of seconds to sleep between
87                                 retries.
88    @attribute ap_spec: an APSpec object
89    @attribute retries: an integer, max number of retry attempts to lock ap.
90                        Defaults to MAX_RETRIES.
91    @attribute aps_to_lock: a list of ApLocker objects.
92    @attribute manager: a HostLockManager object, used to lock/unlock APs.
93    """
94
95
96    MIN_SECONDS_TO_SLEEP = 30
97    MAX_SECONDS_TO_SLEEP = 120
98
99
100    def __init__(self, lock_manager, ap_spec, retries=MAX_RETRIES,
101                 hostname_matching_only=False,
102                 ap_test_type=ap_constants.AP_TEST_TYPE_CHAOS):
103        """Initialize.
104
105        @param ap_spec: an APSpec object
106        @param retries: an integer, max number of retry attempts to lock ap.
107                        Defaults to MAX_RETRIES.
108        @param hostname_matching_only : a boolean, if True matching against
109                                        all other APSpec parameters is not
110                                        performed.
111        @param ap_test_type: Used to determine which type of test we're
112                             currently running (Chaos vs Clique).
113        """
114        self.aps_to_lock = construct_ap_lockers(ap_spec, retries,
115                           hostname_matching_only=hostname_matching_only,
116                           ap_test_type=ap_test_type)
117        self.manager = lock_manager
118        self._locked_aps = []
119
120
121    def has_more_aps(self):
122        """@return True iff there is at least one AP to be locked."""
123        return len(self.aps_to_lock) > 0
124
125
126    def lock_ap_in_afe(self, ap_locker):
127        """Locks an AP host in AFE.
128
129        @param ap_locker: an ApLocker object, AP to be locked.
130        @return a boolean, True iff ap_locker is locked.
131        """
132        if not utils.host_is_in_lab_zone(ap_locker.configurator.host_name):
133            ap_locker.to_be_locked = False
134            return True
135
136        if self.manager.lock([ap_locker.configurator.host_name]):
137            self._locked_aps.append(ap_locker)
138            logging.info('locked %s', ap_locker.configurator.host_name)
139            ap_locker.to_be_locked = False
140            return True
141        else:
142            ap_locker.retries -= 1
143            logging.info('%d retries left for %s',
144                         ap_locker.retries,
145                         ap_locker.configurator.host_name)
146            if ap_locker.retries == 0:
147                logging.info('No more retries left. Remove %s from list',
148                             ap_locker.configurator.host_name)
149                ap_locker.to_be_locked = False
150
151        return False
152
153
154    def get_ap_batch(self, batch_size=ap_cartridge.THREAD_MAX):
155        """Allocates a batch of locked APs.
156
157        @param batch_size: an integer, max. number of aps to lock in one batch.
158                           Defaults to THREAD_MAX in ap_cartridge.py
159        @return a list of APConfigurator objects, locked on AFE.
160        """
161        # We need this while loop to continuously loop over the for loop.
162        # To exit the while loop, we either:
163        #  - locked batch_size number of aps and return them
164        #  - exhausted all retries on all aps in aps_to_lock
165        while len(self.aps_to_lock):
166            ap_batch = []
167
168            for ap_locker in self.aps_to_lock:
169                logging.info('checking %s', ap_locker.configurator.host_name)
170                if self.lock_ap_in_afe(ap_locker):
171                    ap_batch.append(ap_locker.configurator)
172                    if len(ap_batch) == batch_size:
173                        break
174
175            # Remove locked APs from list of APs to process.
176            aps_to_rm = [ap for ap in self.aps_to_lock if not ap.to_be_locked]
177            self.aps_to_lock = list(set(self.aps_to_lock) - set(aps_to_rm))
178            for ap in aps_to_rm:
179                logging.info('Removed %s from self.aps_to_lock',
180                             ap.configurator.host_name)
181            logging.info('Remaining aps to lock = %s',
182                         [ap.configurator.host_name for ap in self.aps_to_lock])
183
184            # Return available APs and retry remaining ones later.
185            if ap_batch:
186                return ap_batch
187
188            # Sleep before next retry.
189            if self.aps_to_lock:
190                seconds_to_sleep = random.randint(self.MIN_SECONDS_TO_SLEEP,
191                                                  self.MAX_SECONDS_TO_SLEEP)
192                logging.info('Sleep %d sec before retry', seconds_to_sleep)
193                sleep(seconds_to_sleep)
194
195        return []
196
197
198    def unlock_one_ap(self, host_name):
199        """Unlock one AP after we're done.
200
201        @param host_name: a string, host name.
202        """
203        for ap_locker in self._locked_aps:
204            if host_name == ap_locker.configurator.host_name:
205                self.manager.unlock(hosts=[host_name])
206                self._locked_aps.remove(ap_locker)
207                return
208
209        logging.error('Tried to unlock a host we have not locked (%s)?',
210                      host_name)
211
212
213    def unlock_aps(self):
214        """Unlock APs after we're done."""
215        # Make a copy of all of the hostnames to process
216        host_names = list()
217        for ap_locker in self._locked_aps:
218            host_names.append(ap_locker.configurator.host_name)
219        for host_name in host_names:
220            self.unlock_one_ap(host_name)
221
222
223    def unlock_and_reclaim_ap(self, host_name):
224        """Unlock an AP but return it to the remaining batch of APs.
225
226        @param host_name: a string, host name.
227        """
228        for ap_locker in self._locked_aps:
229            if host_name == ap_locker.configurator.host_name:
230                self.aps_to_lock.append(ap_locker)
231                self.unlock_one_ap(host_name)
232                return
233
234
235    def unlock_and_reclaim_aps(self):
236        """Unlock APs but return them to the batch of remining APs.
237
238        unlock_aps() will remove the remaining APs from the list of all APs
239        to process.  This method will add the remaining APs back to the pool
240        of unprocessed APs.
241
242        """
243        # Add the APs back into the pool
244        self.aps_to_lock.extend(self._locked_aps)
245        self.unlock_aps()
246