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