1# Copyright (c) 2012 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 5import logging 6import signal 7import common 8 9from autotest_lib.server import site_utils 10from autotest_lib.server.cros.chaos_lib import chaos_datastore_utils 11"""HostLockManager class, for the dynamic_suite module. 12 13A HostLockManager instance manages locking and unlocking a set of autotest DUTs. 14A caller can lock or unlock one or more DUTs. If the caller fails to unlock() 15locked hosts before the instance is destroyed, it will attempt to unlock() the 16hosts automatically, but this is to be avoided. 17 18Sample usage: 19 manager = host_lock_manager.HostLockManager() 20 try: 21 manager.lock(['host1']) 22 # do things 23 finally: 24 manager.unlock() 25""" 26 27class HostLockManager(object): 28 """ 29 @attribute _afe: an instance of AFE as defined in server/frontend.py. 30 @attribute _locked_hosts: a set of DUT hostnames. 31 @attribute LOCK: a string. 32 @attribute UNLOCK: a string. 33 """ 34 35 LOCK = 'lock' 36 UNLOCK = 'unlock' 37 38 39 @property 40 def locked_hosts(self): 41 """@returns set of locked hosts.""" 42 return self._locked_hosts 43 44 45 @locked_hosts.setter 46 def locked_hosts(self, hosts): 47 """Sets value of locked_hosts. 48 49 @param hosts: a set of strings. 50 """ 51 self._locked_hosts = hosts 52 53 54 def __init__(self, afe=None): 55 """ 56 Constructor 57 """ 58 self.dutils = chaos_datastore_utils.ChaosDataStoreUtils() 59 # Keep track of hosts locked by this instance. 60 self._locked_hosts = set() 61 62 63 def __del__(self): 64 if self._locked_hosts: 65 logging.warning('Caller failed to unlock %r! Forcing unlock now.', 66 self._locked_hosts) 67 self.unlock() 68 69 70 def _check_host(self, host, operation): 71 """Checks host for desired operation. 72 73 @param host: a string, hostname. 74 @param operation: a string, LOCK or UNLOCK. 75 @returns a string: host name, if desired operation can be performed on 76 host or None otherwise. 77 """ 78 host_checked = host 79 # Get host details from DataStore 80 host_info = self.dutils.show_device(host) 81 82 if not host_info: 83 logging.warning('Host (AP) details not found in DataStore') 84 return None 85 86 if operation == self.LOCK and host_info['lock_status']: 87 err = ('Contention detected: %s is locked by %s at %s.' % 88 (host, host_info['locked_by'], 89 host_info['lock_status_updated'])) 90 logging.error(err) 91 return None 92 93 elif operation == self.UNLOCK and not host_info['lock_status']: 94 logging.info('%s not locked.', host) 95 return None 96 97 return host_checked 98 99 100 def lock(self, hosts, lock_reason='Locked by HostLockManager'): 101 """Lock hosts in datastore. 102 103 @param hosts: a list of strings, host names. 104 @param lock_reason: a string, a reason for locking the hosts. 105 106 @returns a boolean, True == at least one host from hosts is locked. 107 """ 108 # Filter out hosts that we may have already locked 109 new_hosts = set(hosts).difference(self._locked_hosts) 110 logging.info('Attempt to lock %s', new_hosts) 111 if not new_hosts: 112 return False 113 114 return self._host_modifier(new_hosts, self.LOCK, lock_reason=lock_reason) 115 116 117 def unlock(self, hosts=None): 118 """Unlock hosts in datastore after use. 119 120 @param hosts: a list of strings, host names. 121 @returns a boolean, True == at least one host from self._locked_hosts is 122 unlocked. 123 """ 124 # Filter out hosts that we did not lock 125 updated_hosts = self._locked_hosts 126 if hosts: 127 unknown_hosts = set(hosts).difference(self._locked_hosts) 128 logging.warning('Skip unknown hosts: %s', unknown_hosts) 129 updated_hosts = set(hosts) - unknown_hosts 130 logging.info('Valid hosts: %s', updated_hosts) 131 updated_hosts = updated_hosts.intersection(self._locked_hosts) 132 133 if not updated_hosts: 134 return False 135 136 logging.info('Unlocking hosts (APs / PCAPs): %s', updated_hosts) 137 return self._host_modifier(updated_hosts, self.UNLOCK) 138 139 140 def _host_modifier(self, hosts, operation, lock_reason=None): 141 """Helper that locks hosts in DataStore. 142 143 @param: hosts, a set of strings, host names. 144 @param operation: a string, LOCK or UNLOCK. 145 @param lock_reason: a string, a reason must be provided when locking. 146 147 @returns a boolean, if operation succeeded on at least one host in 148 hosts. 149 """ 150 updated_hosts = set() 151 for host in hosts: 152 verified_host = self._check_host(host, operation) 153 if verified_host is not None: 154 updated_hosts.add(verified_host) 155 156 logging.info('host_modifier: updated_hosts = %s', updated_hosts) 157 if not updated_hosts: 158 logging.info('host_modifier: no host to update') 159 return False 160 161 for host in updated_hosts: 162 if operation == self.LOCK: 163 if self.dutils.lock_device(host, lock_reason): 164 logging.info('Locked host in datastore: %s', host) 165 self._locked_hosts = self._locked_hosts.union([host]) 166 else: 167 logging.error('Unable to lock host: ', host) 168 169 if operation == self.UNLOCK: 170 if self.dutils.unlock_device(host): 171 logging.info('Unlocked host in datastore: %s', host) 172 self._locked_hosts = self._locked_hosts.difference([host]) 173 else: 174 logging.error('Unable to un-lock host: %s', host) 175 176 return True 177 178 179class HostsLockedBy(object): 180 """Context manager to make sure that a HostLockManager will always unlock 181 its machines. This protects against both exceptions and SIGTERM.""" 182 183 def _make_handler(self): 184 def _chaining_signal_handler(signal_number, frame): 185 self._manager.unlock() 186 # self._old_handler can also be signal.SIG_{IGN,DFL} which are ints. 187 if callable(self._old_handler): 188 self._old_handler(signal_number, frame) 189 return _chaining_signal_handler 190 191 192 def __init__(self, manager): 193 """ 194 @param manager: The HostLockManager used to lock the hosts. 195 """ 196 self._manager = manager 197 self._old_handler = signal.SIG_DFL 198 199 200 def __enter__(self): 201 self._old_handler = signal.signal(signal.SIGTERM, self._make_handler()) 202 203 204 def __exit__(self, exntype, exnvalue, backtrace): 205 signal.signal(signal.SIGTERM, self._old_handler) 206 self._manager.unlock() 207