1 #!/usr/bin/env python2 2 # 3 # Copyright 2010 Google Inc. All Rights Reserved. 4 """Script to lock/unlock machines.""" 5 6 from __future__ import print_function 7 8 __author__ = 'asharif (at] google.com (Ahmad Sharif)' 9 10 import argparse 11 import datetime 12 import fcntl 13 import getpass 14 import glob 15 import json 16 import os 17 import socket 18 import sys 19 import time 20 21 from cros_utils import logger 22 23 LOCK_SUFFIX = '_check_lock_liveness' 24 25 # The locks file directory REQUIRES that 'group' only has read/write 26 # privileges and 'world' has no privileges. So the mask must be 27 # '0027': 0777 - 0027 = 0750. 28 LOCK_MASK = 0027 29 30 31 def FileCheckName(name): 32 return name + LOCK_SUFFIX 33 34 35 def OpenLiveCheck(file_name): 36 with FileCreationMask(LOCK_MASK): 37 fd = open(file_name, 'a') 38 try: 39 fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) 40 except IOError: 41 raise 42 return fd 43 44 45 class FileCreationMask(object): 46 """Class for the file creation mask.""" 47 48 def __init__(self, mask): 49 self._mask = mask 50 self._old_mask = None 51 52 def __enter__(self): 53 self._old_mask = os.umask(self._mask) 54 55 def __exit__(self, typ, value, traceback): 56 os.umask(self._old_mask) 57 58 59 class LockDescription(object): 60 """The description of the lock.""" 61 62 def __init__(self, desc=None): 63 try: 64 self.owner = desc['owner'] 65 self.exclusive = desc['exclusive'] 66 self.counter = desc['counter'] 67 self.time = desc['time'] 68 self.reason = desc['reason'] 69 self.auto = desc['auto'] 70 except (KeyError, TypeError): 71 self.owner = '' 72 self.exclusive = False 73 self.counter = 0 74 self.time = 0 75 self.reason = '' 76 self.auto = False 77 78 def IsLocked(self): 79 return self.counter or self.exclusive 80 81 def __str__(self): 82 return ' '.join([ 83 'Owner: %s' % self.owner, 'Exclusive: %s' % self.exclusive, 84 'Counter: %s' % self.counter, 'Time: %s' % self.time, 85 'Reason: %s' % self.reason, 'Auto: %s' % self.auto 86 ]) 87 88 89 class FileLock(object): 90 """File lock operation class.""" 91 FILE_OPS = [] 92 93 def __init__(self, lock_filename): 94 self._filepath = lock_filename 95 lock_dir = os.path.dirname(lock_filename) 96 assert os.path.isdir(lock_dir), ("Locks dir: %s doesn't exist!" % lock_dir) 97 self._file = None 98 self._description = None 99 100 def getDescription(self): 101 return self._description 102 103 def getFilePath(self): 104 return self._filepath 105 106 def setDescription(self, desc): 107 self._description = desc 108 109 @classmethod 110 def AsString(cls, file_locks): 111 stringify_fmt = '%-30s %-15s %-4s %-4s %-15s %-40s %-4s' 112 header = stringify_fmt % ('machine', 'owner', 'excl', 'ctr', 'elapsed', 113 'reason', 'auto') 114 lock_strings = [] 115 for file_lock in file_locks: 116 117 elapsed_time = datetime.timedelta( 118 seconds=int(time.time() - file_lock.getDescription().time)) 119 elapsed_time = '%s ago' % elapsed_time 120 lock_strings.append( 121 stringify_fmt % 122 (os.path.basename(file_lock.getFilePath), 123 file_lock.getDescription().owner, 124 file_lock.getDescription().exclusive, 125 file_lock.getDescription().counter, elapsed_time, 126 file_lock.getDescription().reason, file_lock.getDescription().auto)) 127 table = '\n'.join(lock_strings) 128 return '\n'.join([header, table]) 129 130 @classmethod 131 def ListLock(cls, pattern, locks_dir): 132 if not locks_dir: 133 locks_dir = Machine.LOCKS_DIR 134 full_pattern = os.path.join(locks_dir, pattern) 135 file_locks = [] 136 for lock_filename in glob.glob(full_pattern): 137 if LOCK_SUFFIX in lock_filename: 138 continue 139 file_lock = FileLock(lock_filename) 140 with file_lock as lock: 141 if lock.IsLocked(): 142 file_locks.append(file_lock) 143 logger.GetLogger().LogOutput('\n%s' % cls.AsString(file_locks)) 144 145 def __enter__(self): 146 with FileCreationMask(LOCK_MASK): 147 try: 148 self._file = open(self._filepath, 'a+') 149 self._file.seek(0, os.SEEK_SET) 150 151 if fcntl.flock(self._file.fileno(), fcntl.LOCK_EX) == -1: 152 raise IOError('flock(%s, LOCK_EX) failed!' % self._filepath) 153 154 try: 155 desc = json.load(self._file) 156 except (EOFError, ValueError): 157 desc = None 158 self._description = LockDescription(desc) 159 160 if self._description.exclusive and self._description.auto: 161 locked_byself = False 162 for fd in self.FILE_OPS: 163 if fd.name == FileCheckName(self._filepath): 164 locked_byself = True 165 break 166 if not locked_byself: 167 try: 168 fp = OpenLiveCheck(FileCheckName(self._filepath)) 169 except IOError: 170 pass 171 else: 172 self._description = LockDescription() 173 fcntl.lockf(fp, fcntl.LOCK_UN) 174 fp.close() 175 return self._description 176 # Check this differently? 177 except IOError as ex: 178 logger.GetLogger().LogError(ex) 179 return None 180 181 def __exit__(self, typ, value, traceback): 182 self._file.truncate(0) 183 self._file.write(json.dumps(self._description.__dict__, skipkeys=True)) 184 self._file.close() 185 186 def __str__(self): 187 return self.AsString([self]) 188 189 190 class Lock(object): 191 """Lock class""" 192 193 def __init__(self, lock_file, auto=True): 194 self._to_lock = os.path.basename(lock_file) 195 self._lock_file = lock_file 196 self._logger = logger.GetLogger() 197 self._auto = auto 198 199 def NonBlockingLock(self, exclusive, reason=''): 200 with FileLock(self._lock_file) as lock: 201 if lock.exclusive: 202 self._logger.LogError( 203 'Exclusive lock already acquired by %s. Reason: %s' % (lock.owner, 204 lock.reason)) 205 return False 206 207 if exclusive: 208 if lock.counter: 209 self._logger.LogError('Shared lock already acquired') 210 return False 211 lock_file_check = FileCheckName(self._lock_file) 212 fd = OpenLiveCheck(lock_file_check) 213 FileLock.FILE_OPS.append(fd) 214 215 lock.exclusive = True 216 lock.reason = reason 217 lock.owner = getpass.getuser() 218 lock.time = time.time() 219 lock.auto = self._auto 220 else: 221 lock.counter += 1 222 self._logger.LogOutput('Successfully locked: %s' % self._to_lock) 223 return True 224 225 def Unlock(self, exclusive, force=False): 226 with FileLock(self._lock_file) as lock: 227 if not lock.IsLocked(): 228 self._logger.LogWarning("Can't unlock unlocked machine!") 229 return True 230 231 if lock.exclusive != exclusive: 232 self._logger.LogError('shared locks must be unlocked with --shared') 233 return False 234 235 if lock.exclusive: 236 if lock.owner != getpass.getuser() and not force: 237 self._logger.LogError("%s can't unlock lock owned by: %s" % 238 (getpass.getuser(), lock.owner)) 239 return False 240 if lock.auto != self._auto: 241 self._logger.LogError("Can't unlock lock with different -a" 242 ' parameter.') 243 return False 244 lock.exclusive = False 245 lock.reason = '' 246 lock.owner = '' 247 248 if self._auto: 249 del_list = [ 250 i for i in FileLock.FILE_OPS 251 if i.name == FileCheckName(self._lock_file) 252 ] 253 for i in del_list: 254 FileLock.FILE_OPS.remove(i) 255 for f in del_list: 256 fcntl.lockf(f, fcntl.LOCK_UN) 257 f.close() 258 del del_list 259 os.remove(FileCheckName(self._lock_file)) 260 261 else: 262 lock.counter -= 1 263 return True 264 265 266 class Machine(object): 267 """Machine class""" 268 269 LOCKS_DIR = '/google/data/rw/users/mo/mobiletc-prebuild/locks' 270 271 def __init__(self, name, locks_dir=LOCKS_DIR, auto=True): 272 self._name = name 273 self._auto = auto 274 try: 275 self._full_name = socket.gethostbyaddr(name)[0] 276 except socket.error: 277 self._full_name = self._name 278 self._full_name = os.path.join(locks_dir, self._full_name) 279 280 def Lock(self, exclusive=False, reason=''): 281 lock = Lock(self._full_name, self._auto) 282 return lock.NonBlockingLock(exclusive, reason) 283 284 def TryLock(self, timeout=300, exclusive=False, reason=''): 285 locked = False 286 sleep = timeout / 10 287 while True: 288 locked = self.Lock(exclusive, reason) 289 if locked or not timeout >= 0: 290 break 291 print('Lock not acquired for {0}, wait {1} seconds ...'.format( 292 self._name, sleep)) 293 time.sleep(sleep) 294 timeout -= sleep 295 return locked 296 297 def Unlock(self, exclusive=False, ignore_ownership=False): 298 lock = Lock(self._full_name, self._auto) 299 return lock.Unlock(exclusive, ignore_ownership) 300 301 302 def Main(argv): 303 """The main function.""" 304 305 parser = argparse.ArgumentParser() 306 parser.add_argument( 307 '-r', '--reason', dest='reason', default='', help='The lock reason.') 308 parser.add_argument( 309 '-u', 310 '--unlock', 311 dest='unlock', 312 action='store_true', 313 default=False, 314 help='Use this to unlock.') 315 parser.add_argument( 316 '-l', 317 '--list_locks', 318 dest='list_locks', 319 action='store_true', 320 default=False, 321 help='Use this to list locks.') 322 parser.add_argument( 323 '-f', 324 '--ignore_ownership', 325 dest='ignore_ownership', 326 action='store_true', 327 default=False, 328 help="Use this to force unlock on a lock you don't own.") 329 parser.add_argument( 330 '-s', 331 '--shared', 332 dest='shared', 333 action='store_true', 334 default=False, 335 help='Use this for a shared (non-exclusive) lock.') 336 parser.add_argument( 337 '-d', 338 '--dir', 339 dest='locks_dir', 340 action='store', 341 default=Machine.LOCKS_DIR, 342 help='Use this to set different locks_dir') 343 parser.add_argument('args', nargs='*', help='Machine arg.') 344 345 options = parser.parse_args(argv) 346 347 options.locks_dir = os.path.abspath(options.locks_dir) 348 exclusive = not options.shared 349 350 if not options.list_locks and len(options.args) != 2: 351 logger.GetLogger().LogError( 352 'Either --list_locks or a machine arg is needed.') 353 return 1 354 355 if len(options.args) > 1: 356 machine = Machine(options.args[1], options.locks_dir, auto=False) 357 else: 358 machine = None 359 360 if options.list_locks: 361 FileLock.ListLock('*', options.locks_dir) 362 retval = True 363 elif options.unlock: 364 retval = machine.Unlock(exclusive, options.ignore_ownership) 365 else: 366 retval = machine.Lock(exclusive, options.reason) 367 368 if retval: 369 return 0 370 else: 371 return 1 372 373 374 if __name__ == '__main__': 375 sys.exit(Main(sys.argv[1:])) 376
