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# 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 = 10 96 MAX_SECONDS_TO_SLEEP = 20 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(self, ap_locker): 126 """Locks an AP host in DataStore. 127 128 @param ap_locker: an ApLocker object, AP to be locked. 129 @return a boolean, True if ap_locker is locked. 130 """ 131 132 if self.manager.lock([ap_locker.configurator.host_name]): 133 self._locked_aps.append(ap_locker) 134 logging.info('locked %s', ap_locker.configurator.host_name) 135 ap_locker.to_be_locked = False 136 return True 137 else: 138 ap_locker.retries -= 1 139 logging.info('%d retries left for %s', 140 ap_locker.retries, 141 ap_locker.configurator.host_name) 142 if ap_locker.retries == 0: 143 logging.info('No more retries left. Remove %s from list', 144 ap_locker.configurator.host_name) 145 ap_locker.to_be_locked = False 146 147 return False 148 149 150 151 def get_ap_batch(self, batch_size=ap_cartridge.THREAD_MAX): 152 """Allocates a batch of locked APs. 153 154 @param batch_size: an integer, max. number of aps to lock in one batch. 155 Defaults to THREAD_MAX in ap_cartridge.py 156 @return a list of APConfigurator objects, locked in datastore. 157 """ 158 # We need this while loop to continuously loop over the for loop. 159 # To exit the while loop, we either: 160 # - locked batch_size number of aps and return them 161 # - exhausted all retries on all aps in aps_to_lock 162 while len(self.aps_to_lock): 163 ap_batch = [] 164 165 for ap_locker in self.aps_to_lock: 166 logging.info('checking %s', ap_locker.configurator.host_name) 167 # Lock AP in DataStore 168 if self.lock_ap(ap_locker): 169 ap_batch.append(ap_locker.configurator) 170 if len(ap_batch) == batch_size: 171 break 172 173 # Remove locked APs from list of APs to process. 174 aps_to_rm = [ap for ap in self.aps_to_lock if not ap.to_be_locked] 175 self.aps_to_lock = list(set(self.aps_to_lock) - set(aps_to_rm)) 176 for ap in aps_to_rm: 177 logging.info('Removed %s from self.aps_to_lock', 178 ap.configurator.host_name) 179 logging.info('Remaining aps to lock = %s', 180 [ap.configurator.host_name for ap in self.aps_to_lock]) 181 182 # Return available APs and retry remaining ones later. 183 if ap_batch: 184 return ap_batch 185 186 # Sleep before next retry. 187 if self.aps_to_lock: 188 seconds_to_sleep = random.randint(self.MIN_SECONDS_TO_SLEEP, 189 self.MAX_SECONDS_TO_SLEEP) 190 logging.info('Sleep %d sec before retry', seconds_to_sleep) 191 sleep(seconds_to_sleep) 192 193 return [] 194 195 196 def unlock_one_ap(self, host_name): 197 """ 198 Unlock one AP from datastore after we're done. 199 200 @param host_name: a string, AP host name. 201 202 @raise TestError: when unable to unlock AP in datastore. 203 204 """ 205 for ap_locker in self._locked_aps: 206 if host_name == ap_locker.configurator.host_name: 207 self.manager.unlock(hosts=[host_name]) 208 self._locked_aps.remove(ap_locker) 209 return 210 211 logging.error('Tried to unlock a host we have not locked (%s)?', 212 host_name) 213 214 215 def unlock_aps(self): 216 """Unlock APs after we're done.""" 217 # Make a copy of all of the hostnames to process 218 219 host_names = list() 220 for ap_locker in self._locked_aps: 221 host_names.append(ap_locker.configurator.host_name) 222 for host_name in host_names: 223 self.unlock_one_ap(host_name) 224 225 226 def unlock_and_reclaim_ap(self, host_name): 227 """Unlock an AP but return it to the remaining batch of APs. 228 229 @param host_name: a string, host name. 230 """ 231 for ap_locker in self._locked_aps: 232 if host_name == ap_locker.configurator.host_name: 233 self.aps_to_lock.append(ap_locker) 234 self.unlock_one_ap(host_name) 235 return 236 237 238 def unlock_and_reclaim_aps(self): 239 """Unlock APs but return them to the batch of remining APs. 240 241 unlock_aps() will remove the remaining APs from the list of all APs 242 to process. This method will add the remaining APs back to the pool 243 of unprocessed APs. 244 245 """ 246 # Add the APs back into the pool 247 self.aps_to_lock.extend(self._locked_aps) 248 self.unlock_aps() 249