1 # Copyright 2015 The Chromium OS 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 5 import collections 6 import logging 7 import random 8 9 from time import sleep 10 11 import common 12 from autotest_lib.client.common_lib import error 13 from autotest_lib.server import hosts 14 from autotest_lib.server import frontend 15 from autotest_lib.server import site_utils 16 from autotest_lib.server.cros.dynamic_suite import constants 17 from autotest_lib.server.cros.network import wifi_client 18 19 # Max number of retry attempts to lock a DUT. 20 MAX_RETRIES = 3 21 22 # Tuple containing the DUT objects 23 DUTObject = collections.namedtuple('DUTObject', ['host', 'wifi_client']) 24 25 class DUTSpec(): 26 """Object to specify the DUT spec. 27 28 @attribute board_name: String representing the board name corresponding to 29 the board. 30 @attribute host_name: String representing the host name corresponding to 31 the machine. 32 """ 33 34 def __init__(self, board_name=None, host_name=None): 35 """Initialize. 36 37 @param board_name: String representing the board name corresponding to 38 the board. 39 @param host_name: String representing the host name corresponding to 40 the machine. 41 """ 42 self.board_name = board_name 43 self.host_name = host_name 44 45 def __repr__(self): 46 """@return class name, dut host name, lock status and retries.""" 47 return 'class: %s, Board name: %s, Num DUTs = %d' % ( 48 self.__class__.__name__, 49 self.board_name, 50 self.host_name) 51 52 53 class DUTSetSpec(list): 54 """Object to specify the DUT set spec. It's a list of DUTSpec objects.""" 55 56 def __init__(self): 57 """Initialize.""" 58 super(DUTSetSpec, self) 59 60 61 class DUTPoolSpec(list): 62 """Object to specify the DUT pool spec.It's a list of DUTSetSpec objects.""" 63 64 def __init__(self): 65 """Initialize.""" 66 super(DUTPoolSpec, self) 67 68 69 class DUTLocker(object): 70 """Object to keep track of DUT lock state. 71 72 @attribute dut_spec: an DUTSpec object corresponding to the DUT we need. 73 @attribute retries: an integer, max number of retry attempts to lock DUT. 74 @attribute to_be_locked: a boolean, True iff DUT has not been locked. 75 """ 76 77 78 def __init__(self, dut_spec, retries): 79 """Initialize. 80 81 @param dut_spec: a DUTSpec object corresponding to the spec of the DUT 82 to be locked. 83 @param retries: an integer, max number of retry attempts to lock DUT. 84 """ 85 self.dut_spec = dut_spec 86 self.retries = retries 87 self.to_be_locked = True 88 89 def __repr__(self): 90 """@return class name, dut host name, lock status and retries.""" 91 return 'class: %s, host name: %s, to_be_locked = %s, retries = %d' % ( 92 self.__class__.__name__, 93 self.dut.host.hostname, 94 self.to_be_locked, 95 self.retries) 96 97 98 class CliqueDUTBatchLocker(object): 99 """Object to lock/unlock an DUT. 100 101 @attribute SECONDS_TO_SLEEP: an integer, number of seconds to sleep between 102 retries. 103 @attribute duts_to_lock: a list of DUTLocker objects. 104 @attribute locked_duts: a list of DUTObject's corresponding to DUT's which 105 have already been allocated. 106 @attribute manager: a HostLockManager object, used to lock/unlock DUTs. 107 """ 108 109 MIN_SECONDS_TO_SLEEP = 30 110 MAX_SECONDS_TO_SLEEP = 120 111 112 def __init__(self, lock_manager, dut_pool_spec, retries=MAX_RETRIES): 113 """Initialize. 114 115 @param lock_manager: a HostLockManager object, used to lock/unlock DUTs. 116 @param dut_pool_spec: A DUTPoolSpec object corresponding to the DUT's in 117 the pool. 118 @param retries: Number of times to retry the locking of DUT's. 119 120 """ 121 self.lock_manager = lock_manager 122 self.duts_to_lock = self._construct_dut_lockers(dut_pool_spec, retries) 123 self.locked_duts = [] 124 125 @staticmethod 126 def _construct_dut_lockers(dut_pool_spec, retries): 127 """Convert DUTObject objects to DUTLocker objects for locking. 128 129 @param dut_pool_spec: A DUTPoolSpec object corresponding to the DUT's in 130 the pool. 131 @param retries: an integer, max number of retry attempts to lock DUT. 132 133 @return a list of DUTLocker objects. 134 """ 135 dut_lockers_list = [] 136 for dut_set_spec in dut_pool_spec: 137 dut_set_lockers_list = [] 138 for dut_spec in dut_set_spec: 139 dut_locker = DUTLocker(dut_spec, retries) 140 dut_set_lockers_list.append(dut_locker) 141 dut_lockers_list.append(dut_set_lockers_list) 142 return dut_lockers_list 143 144 def _allocate_dut(self, host_name=None, board_name=None): 145 """Allocates a machine to the DUT pool for running the test. 146 147 Locks the allocated machine if the machine was discovered via AFE 148 to prevent tests stomping on each other. 149 150 @param host_name: Host name for the DUT. 151 @param board_name: Board name Label to use for finding the DUT. 152 153 @return: hostname of the device locked in AFE. 154 """ 155 hostname = None 156 if host_name: 157 if self.lock_manager.lock([host_name]): 158 logging.info('Locked device %s.', host_name) 159 hostname = host_name 160 else: 161 logging.error('Unable to lock device %s.', host_name) 162 else: 163 afe = frontend.AFE(debug=True, 164 server=site_utils.get_global_afe_hostname()) 165 labels = [] 166 labels.append(constants.BOARD_PREFIX + board_name) 167 labels.append('clique_dut') 168 try: 169 hostname = site_utils.lock_host_with_labels( 170 afe, self.lock_manager, labels=labels) + '.cros' 171 except error.NoEligibleHostException as e: 172 raise error.TestError("Unable to find a suitable device.") 173 except error.TestError as e: 174 logging.error(e) 175 return hostname 176 177 @staticmethod 178 def _create_dut_object(host_name): 179 """Create the DUTObject tuple for the DUT. 180 181 @param host_name: Host name for the DUT. 182 183 @return: Tuple of Host and Wifi client objects representing DUTObject 184 for invoking RPC calls. 185 """ 186 dut_host = hosts.create_host(host_name) 187 dut_wifi_client = wifi_client.WiFiClient(dut_host, './debug', False) 188 return DUTObject(dut_host, dut_wifi_client) 189 190 def _lock_dut_in_afe(self, dut_locker): 191 """Locks an DUT host in AFE. 192 193 @param dut_locker: an DUTLocker object, DUT to be locked. 194 @return a hostname iff dut_locker is locked, else returns None. 195 """ 196 logging.debug('Trying to find a device with spec (%s, %s)', 197 dut_locker.dut_spec.host_name, 198 dut_locker.dut_spec.board_name) 199 host_name = self._allocate_dut( 200 dut_locker.dut_spec.host_name, dut_locker.dut_spec.board_name) 201 if host_name: 202 logging.info('Locked %s', host_name) 203 dut_locker.to_be_locked = False 204 else: 205 dut_locker.retries -= 1 206 logging.info('%d retries left for (%s, %s)', 207 dut_locker.retries, 208 dut_locker.dut_spec.host_name, 209 dut_locker.dut_spec.board_name) 210 if dut_locker.retries == 0: 211 raise error.TestError("No more retries left to lock a " 212 "suitable device.") 213 return host_name 214 215 def get_dut_pool(self): 216 """Allocates a batch of locked DUTs for the test. 217 218 @return a list of DUTObject, locked on AFE. 219 """ 220 # We need this while loop to continuously loop over the for loop. 221 # To exit the while loop, we either: 222 # - locked batch_size number of duts and return them 223 # - exhausted all retries on a dut in duts_to_lock 224 225 # It is important to preserve the order of DUT sets, but the order of 226 # DUT's within the set is not important as all the DUT's within a set 227 # have to perform the same role. 228 dut_pool = [] 229 for dut_set in self.duts_to_lock: 230 dut_pool.append([]) 231 232 num_duts_to_lock = sum(map(len, self.duts_to_lock)) 233 while num_duts_to_lock: 234 set_num = 0 235 for dut_locker_set in self.duts_to_lock: 236 for dut_locker in dut_locker_set: 237 if dut_locker.to_be_locked: 238 host_name = self._lock_dut_in_afe(dut_locker) 239 if host_name: 240 dut_object = self._create_dut_object(host_name) 241 self.locked_duts.append(dut_object) 242 dut_pool[set_num].append(dut_object) 243 num_duts_to_lock -= 1 244 set_num += 1 245 246 logging.info('Remaining DUTs to lock: %d', num_duts_to_lock) 247 248 if num_duts_to_lock: 249 seconds_to_sleep = random.randint(self.MIN_SECONDS_TO_SLEEP, 250 self.MAX_SECONDS_TO_SLEEP) 251 logging.debug('Sleep %d sec before retry', seconds_to_sleep) 252 sleep(seconds_to_sleep) 253 return dut_pool 254 255 def _unlock_one_dut(self, dut): 256 """Unlock one DUT after we're done. 257 258 @param dut: a DUTObject corresponding to the DUT. 259 """ 260 host_name = dut.host.host_name 261 if self.manager.unlock(hosts=[host_name]): 262 self._locked_duts.remove(dut) 263 else: 264 logging.error('Tried to unlock a host we have not locked (%s)?', 265 host_name) 266 267 def unlock_duts(self): 268 """Unlock DUTs after we're done.""" 269 for dut in self.locked_duts: 270 self._unlock_one_dut(dut) 271 272 def unlock_and_close_duts(self): 273 """Unlock DUTs after we're done and close the associated WifiClient.""" 274 for dut in self.locked_duts: 275 dut.wifi_client.close() 276 self._unlock_one_dut(dut) 277
