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