diff options
Diffstat (limited to 'file_lock_machine.py')
| -rwxr-xr-x | file_lock_machine.py | 690 |
1 files changed, 362 insertions, 328 deletions
diff --git a/file_lock_machine.py b/file_lock_machine.py index 5bba4430..5268398c 100755 --- a/file_lock_machine.py +++ b/file_lock_machine.py @@ -1,15 +1,13 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -# Copyright 2019 The Chromium OS Authors. All rights reserved. +# Copyright 2019 The ChromiumOS Authors # Use of this source code is governed by a BSD-style license that can be # found in the LICENSE file. """Script to lock/unlock machines.""" -from __future__ import division -from __future__ import print_function -__author__ = 'asharif@google.com (Ahmad Sharif)' +__author__ = "asharif@google.com (Ahmad Sharif)" import argparse import datetime @@ -24,7 +22,8 @@ import time from cros_utils import logger -LOCK_SUFFIX = '_check_lock_liveness' + +LOCK_SUFFIX = "_check_lock_liveness" # The locks file directory REQUIRES that 'group' only has read/write # privileges and 'world' has no privileges. So the mask must be @@ -33,357 +32,392 @@ LOCK_MASK = 0o27 def FileCheckName(name): - return name + LOCK_SUFFIX + return name + LOCK_SUFFIX def OpenLiveCheck(file_name): - with FileCreationMask(LOCK_MASK): - fd = open(file_name, 'a') - try: - fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - except IOError as e: - logger.GetLogger().LogError(e) - raise - return fd + with FileCreationMask(LOCK_MASK): + fd = open(file_name, "a") + try: + fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except IOError as e: + logger.GetLogger().LogError(e) + raise + return fd class FileCreationMask(object): - """Class for the file creation mask.""" + """Class for the file creation mask.""" - def __init__(self, mask): - self._mask = mask - self._old_mask = None + def __init__(self, mask): + self._mask = mask + self._old_mask = None - def __enter__(self): - self._old_mask = os.umask(self._mask) + def __enter__(self): + self._old_mask = os.umask(self._mask) - def __exit__(self, typ, value, traceback): - os.umask(self._old_mask) + def __exit__(self, typ, value, traceback): + os.umask(self._old_mask) class LockDescription(object): - """The description of the lock.""" + """The description of the lock.""" - def __init__(self, desc=None): - try: - self.owner = desc['owner'] - self.exclusive = desc['exclusive'] - self.counter = desc['counter'] - self.time = desc['time'] - self.reason = desc['reason'] - self.auto = desc['auto'] - except (KeyError, TypeError): - self.owner = '' - self.exclusive = False - self.counter = 0 - self.time = 0 - self.reason = '' - self.auto = False - - def IsLocked(self): - return self.counter or self.exclusive - - def __str__(self): - return ' '.join([ - 'Owner: %s' % self.owner, - 'Exclusive: %s' % self.exclusive, - 'Counter: %s' % self.counter, - 'Time: %s' % self.time, - 'Reason: %s' % self.reason, - 'Auto: %s' % self.auto, - ]) + def __init__(self, desc=None): + try: + self.owner = desc["owner"] + self.exclusive = desc["exclusive"] + self.counter = desc["counter"] + self.time = desc["time"] + self.reason = desc["reason"] + self.auto = desc["auto"] + except (KeyError, TypeError): + self.owner = "" + self.exclusive = False + self.counter = 0 + self.time = 0 + self.reason = "" + self.auto = False + + def IsLocked(self): + return self.counter or self.exclusive + + def __str__(self): + return " ".join( + [ + "Owner: %s" % self.owner, + "Exclusive: %s" % self.exclusive, + "Counter: %s" % self.counter, + "Time: %s" % self.time, + "Reason: %s" % self.reason, + "Auto: %s" % self.auto, + ] + ) class FileLock(object): - """File lock operation class.""" - FILE_OPS = [] - - def __init__(self, lock_filename): - self._filepath = lock_filename - lock_dir = os.path.dirname(lock_filename) - assert os.path.isdir(lock_dir), ("Locks dir: %s doesn't exist!" % lock_dir) - self._file = None - self._description = None - - self.exclusive = None - self.auto = None - self.reason = None - self.time = None - self.owner = None - - def getDescription(self): - return self._description - - def getFilePath(self): - return self._filepath - - def setDescription(self, desc): - self._description = desc - - @classmethod - def AsString(cls, file_locks): - stringify_fmt = '%-30s %-15s %-4s %-4s %-15s %-40s %-4s' - header = stringify_fmt % ('machine', 'owner', 'excl', 'ctr', 'elapsed', - 'reason', 'auto') - lock_strings = [] - for file_lock in file_locks: - - elapsed_time = datetime.timedelta( - seconds=int(time.time() - file_lock.getDescription().time)) - elapsed_time = '%s ago' % elapsed_time - lock_strings.append( - stringify_fmt % (os.path.basename(file_lock.getFilePath), - file_lock.getDescription().owner, - file_lock.getDescription().exclusive, - file_lock.getDescription().counter, elapsed_time, - file_lock.getDescription().reason, - file_lock.getDescription().auto)) - table = '\n'.join(lock_strings) - return '\n'.join([header, table]) - - @classmethod - def ListLock(cls, pattern, locks_dir): - if not locks_dir: - locks_dir = Machine.LOCKS_DIR - full_pattern = os.path.join(locks_dir, pattern) - file_locks = [] - for lock_filename in glob.glob(full_pattern): - if LOCK_SUFFIX in lock_filename: - continue - file_lock = FileLock(lock_filename) - with file_lock as lock: - if lock.IsLocked(): - file_locks.append(file_lock) - logger.GetLogger().LogOutput('\n%s' % cls.AsString(file_locks)) - - def __enter__(self): - with FileCreationMask(LOCK_MASK): - try: - self._file = open(self._filepath, 'a+') - self._file.seek(0, os.SEEK_SET) - - if fcntl.flock(self._file.fileno(), fcntl.LOCK_EX) == -1: - raise IOError('flock(%s, LOCK_EX) failed!' % self._filepath) - - try: - desc = json.load(self._file) - except (EOFError, ValueError): - desc = None - self._description = LockDescription(desc) - - if self._description.exclusive and self._description.auto: - locked_byself = False - for fd in self.FILE_OPS: - if fd.name == FileCheckName(self._filepath): - locked_byself = True - break - if not locked_byself: - try: - fp = OpenLiveCheck(FileCheckName(self._filepath)) - except IOError: - pass - else: - self._description = LockDescription() - fcntl.lockf(fp, fcntl.LOCK_UN) - fp.close() + """File lock operation class.""" + + FILE_OPS = [] + + def __init__(self, lock_filename): + self._filepath = lock_filename + lock_dir = os.path.dirname(lock_filename) + assert os.path.isdir(lock_dir), ( + "Locks dir: %s doesn't exist!" % lock_dir + ) + self._file = None + self._description = None + + self.exclusive = None + self.auto = None + self.reason = None + self.time = None + self.owner = None + + def getDescription(self): return self._description - # Check this differently? - except IOError as ex: - logger.GetLogger().LogError(ex) - return None - - def __exit__(self, typ, value, traceback): - self._file.truncate(0) - self._file.write(json.dumps(self._description.__dict__, skipkeys=True)) - self._file.close() - def __str__(self): - return self.AsString([self]) + def getFilePath(self): + return self._filepath + + def setDescription(self, desc): + self._description = desc + + @classmethod + def AsString(cls, file_locks): + stringify_fmt = "%-30s %-15s %-4s %-4s %-15s %-40s %-4s" + header = stringify_fmt % ( + "machine", + "owner", + "excl", + "ctr", + "elapsed", + "reason", + "auto", + ) + lock_strings = [] + for file_lock in file_locks: + + elapsed_time = datetime.timedelta( + seconds=int(time.time() - file_lock.getDescription().time) + ) + elapsed_time = "%s ago" % elapsed_time + lock_strings.append( + stringify_fmt + % ( + os.path.basename(file_lock.getFilePath), + file_lock.getDescription().owner, + file_lock.getDescription().exclusive, + file_lock.getDescription().counter, + elapsed_time, + file_lock.getDescription().reason, + file_lock.getDescription().auto, + ) + ) + table = "\n".join(lock_strings) + return "\n".join([header, table]) + + @classmethod + def ListLock(cls, pattern, locks_dir): + if not locks_dir: + locks_dir = Machine.LOCKS_DIR + full_pattern = os.path.join(locks_dir, pattern) + file_locks = [] + for lock_filename in glob.glob(full_pattern): + if LOCK_SUFFIX in lock_filename: + continue + file_lock = FileLock(lock_filename) + with file_lock as lock: + if lock.IsLocked(): + file_locks.append(file_lock) + logger.GetLogger().LogOutput("\n%s" % cls.AsString(file_locks)) + + def __enter__(self): + with FileCreationMask(LOCK_MASK): + try: + self._file = open(self._filepath, "a+") + self._file.seek(0, os.SEEK_SET) + + if fcntl.flock(self._file.fileno(), fcntl.LOCK_EX) == -1: + raise IOError("flock(%s, LOCK_EX) failed!" % self._filepath) + + try: + desc = json.load(self._file) + except (EOFError, ValueError): + desc = None + self._description = LockDescription(desc) + + if self._description.exclusive and self._description.auto: + locked_byself = False + for fd in self.FILE_OPS: + if fd.name == FileCheckName(self._filepath): + locked_byself = True + break + if not locked_byself: + try: + fp = OpenLiveCheck(FileCheckName(self._filepath)) + except IOError: + pass + else: + self._description = LockDescription() + fcntl.lockf(fp, fcntl.LOCK_UN) + fp.close() + return self._description + # Check this differently? + except IOError as ex: + logger.GetLogger().LogError(ex) + return None + + def __exit__(self, typ, value, traceback): + self._file.truncate(0) + self._file.write(json.dumps(self._description.__dict__, skipkeys=True)) + self._file.close() + + def __str__(self): + return self.AsString([self]) class Lock(object): - """Lock class""" - - def __init__(self, lock_file, auto=True): - self._to_lock = os.path.basename(lock_file) - self._lock_file = lock_file - self._logger = logger.GetLogger() - self._auto = auto - - def NonBlockingLock(self, exclusive, reason=''): - with FileLock(self._lock_file) as lock: - if lock.exclusive: - self._logger.LogError( - 'Exclusive lock already acquired by %s. Reason: %s' % (lock.owner, - lock.reason)) - return False - - if exclusive: - if lock.counter: - self._logger.LogError('Shared lock already acquired') - return False - lock_file_check = FileCheckName(self._lock_file) - fd = OpenLiveCheck(lock_file_check) - FileLock.FILE_OPS.append(fd) - - lock.exclusive = True - lock.reason = reason - lock.owner = getpass.getuser() - lock.time = time.time() - lock.auto = self._auto - else: - lock.counter += 1 - self._logger.LogOutput('Successfully locked: %s' % self._to_lock) - return True - - def Unlock(self, exclusive, force=False): - with FileLock(self._lock_file) as lock: - if not lock.IsLocked(): - self._logger.LogWarning("Can't unlock unlocked machine!") + """Lock class""" + + def __init__(self, lock_file, auto=True): + self._to_lock = os.path.basename(lock_file) + self._lock_file = lock_file + self._logger = logger.GetLogger() + self._auto = auto + + def NonBlockingLock(self, exclusive, reason=""): + with FileLock(self._lock_file) as lock: + if lock.exclusive: + self._logger.LogError( + "Exclusive lock already acquired by %s. Reason: %s" + % (lock.owner, lock.reason) + ) + return False + + if exclusive: + if lock.counter: + self._logger.LogError("Shared lock already acquired") + return False + lock_file_check = FileCheckName(self._lock_file) + fd = OpenLiveCheck(lock_file_check) + FileLock.FILE_OPS.append(fd) + + lock.exclusive = True + lock.reason = reason + lock.owner = getpass.getuser() + lock.time = time.time() + lock.auto = self._auto + else: + lock.counter += 1 + self._logger.LogOutput("Successfully locked: %s" % self._to_lock) return True - if lock.exclusive != exclusive: - self._logger.LogError('shared locks must be unlocked with --shared') - return False - - if lock.exclusive: - if lock.owner != getpass.getuser() and not force: - self._logger.LogError("%s can't unlock lock owned by: %s" % - (getpass.getuser(), lock.owner)) - return False - if lock.auto != self._auto: - self._logger.LogError("Can't unlock lock with different -a" - ' parameter.') - return False - lock.exclusive = False - lock.reason = '' - lock.owner = '' - - if self._auto: - del_list = [ - i for i in FileLock.FILE_OPS - if i.name == FileCheckName(self._lock_file) - ] - for i in del_list: - FileLock.FILE_OPS.remove(i) - for f in del_list: - fcntl.lockf(f, fcntl.LOCK_UN) - f.close() - del del_list - os.remove(FileCheckName(self._lock_file)) - - else: - lock.counter -= 1 - return True + def Unlock(self, exclusive, force=False): + with FileLock(self._lock_file) as lock: + if not lock.IsLocked(): + self._logger.LogWarning("Can't unlock unlocked machine!") + return True + + if lock.exclusive != exclusive: + self._logger.LogError( + "shared locks must be unlocked with --shared" + ) + return False + + if lock.exclusive: + if lock.owner != getpass.getuser() and not force: + self._logger.LogError( + "%s can't unlock lock owned by: %s" + % (getpass.getuser(), lock.owner) + ) + return False + if lock.auto != self._auto: + self._logger.LogError( + "Can't unlock lock with different -a" " parameter." + ) + return False + lock.exclusive = False + lock.reason = "" + lock.owner = "" + + if self._auto: + del_list = [ + i + for i in FileLock.FILE_OPS + if i.name == FileCheckName(self._lock_file) + ] + for i in del_list: + FileLock.FILE_OPS.remove(i) + for f in del_list: + fcntl.lockf(f, fcntl.LOCK_UN) + f.close() + del del_list + os.remove(FileCheckName(self._lock_file)) + + else: + lock.counter -= 1 + return True class Machine(object): - """Machine class""" + """Machine class""" - LOCKS_DIR = '/google/data/rw/users/mo/mobiletc-prebuild/locks' + LOCKS_DIR = "/google/data/rw/users/mo/mobiletc-prebuild/locks" - def __init__(self, name, locks_dir=LOCKS_DIR, auto=True): - self._name = name - self._auto = auto - try: - self._full_name = socket.gethostbyaddr(name)[0] - except socket.error: - self._full_name = self._name - self._full_name = os.path.join(locks_dir, self._full_name) - - def Lock(self, exclusive=False, reason=''): - lock = Lock(self._full_name, self._auto) - return lock.NonBlockingLock(exclusive, reason) - - def TryLock(self, timeout=300, exclusive=False, reason=''): - locked = False - sleep = timeout / 10 - while True: - locked = self.Lock(exclusive, reason) - if locked or timeout < 0: - break - print('Lock not acquired for {0}, wait {1} seconds ...'.format( - self._name, sleep)) - time.sleep(sleep) - timeout -= sleep - return locked - - def Unlock(self, exclusive=False, ignore_ownership=False): - lock = Lock(self._full_name, self._auto) - return lock.Unlock(exclusive, ignore_ownership) + def __init__(self, name, locks_dir=LOCKS_DIR, auto=True): + self._name = name + self._auto = auto + try: + self._full_name = socket.gethostbyaddr(name)[0] + except socket.error: + self._full_name = self._name + self._full_name = os.path.join(locks_dir, self._full_name) + + def Lock(self, exclusive=False, reason=""): + lock = Lock(self._full_name, self._auto) + return lock.NonBlockingLock(exclusive, reason) + + def TryLock(self, timeout=300, exclusive=False, reason=""): + locked = False + sleep = timeout / 10 + while True: + locked = self.Lock(exclusive, reason) + if locked or timeout < 0: + break + print( + "Lock not acquired for {0}, wait {1} seconds ...".format( + self._name, sleep + ) + ) + time.sleep(sleep) + timeout -= sleep + return locked + + def Unlock(self, exclusive=False, ignore_ownership=False): + lock = Lock(self._full_name, self._auto) + return lock.Unlock(exclusive, ignore_ownership) def Main(argv): - """The main function.""" - - parser = argparse.ArgumentParser() - parser.add_argument( - '-r', '--reason', dest='reason', default='', help='The lock reason.') - parser.add_argument( - '-u', - '--unlock', - dest='unlock', - action='store_true', - default=False, - help='Use this to unlock.') - parser.add_argument( - '-l', - '--list_locks', - dest='list_locks', - action='store_true', - default=False, - help='Use this to list locks.') - parser.add_argument( - '-f', - '--ignore_ownership', - dest='ignore_ownership', - action='store_true', - default=False, - help="Use this to force unlock on a lock you don't own.") - parser.add_argument( - '-s', - '--shared', - dest='shared', - action='store_true', - default=False, - help='Use this for a shared (non-exclusive) lock.') - parser.add_argument( - '-d', - '--dir', - dest='locks_dir', - action='store', - default=Machine.LOCKS_DIR, - help='Use this to set different locks_dir') - parser.add_argument('args', nargs='*', help='Machine arg.') - - options = parser.parse_args(argv) - - options.locks_dir = os.path.abspath(options.locks_dir) - exclusive = not options.shared - - if not options.list_locks and len(options.args) != 2: - logger.GetLogger().LogError( - 'Either --list_locks or a machine arg is needed.') - return 1 - - if len(options.args) > 1: - machine = Machine(options.args[1], options.locks_dir, auto=False) - else: - machine = None - - if options.list_locks: - FileLock.ListLock('*', options.locks_dir) - retval = True - elif options.unlock: - retval = machine.Unlock(exclusive, options.ignore_ownership) - else: - retval = machine.Lock(exclusive, options.reason) - - if retval: - return 0 - else: - return 1 - - -if __name__ == '__main__': - sys.exit(Main(sys.argv[1:])) + """The main function.""" + + parser = argparse.ArgumentParser() + parser.add_argument( + "-r", "--reason", dest="reason", default="", help="The lock reason." + ) + parser.add_argument( + "-u", + "--unlock", + dest="unlock", + action="store_true", + default=False, + help="Use this to unlock.", + ) + parser.add_argument( + "-l", + "--list_locks", + dest="list_locks", + action="store_true", + default=False, + help="Use this to list locks.", + ) + parser.add_argument( + "-f", + "--ignore_ownership", + dest="ignore_ownership", + action="store_true", + default=False, + help="Use this to force unlock on a lock you don't own.", + ) + parser.add_argument( + "-s", + "--shared", + dest="shared", + action="store_true", + default=False, + help="Use this for a shared (non-exclusive) lock.", + ) + parser.add_argument( + "-d", + "--dir", + dest="locks_dir", + action="store", + default=Machine.LOCKS_DIR, + help="Use this to set different locks_dir", + ) + parser.add_argument("args", nargs="*", help="Machine arg.") + + options = parser.parse_args(argv) + + options.locks_dir = os.path.abspath(options.locks_dir) + exclusive = not options.shared + + if not options.list_locks and len(options.args) != 2: + logger.GetLogger().LogError( + "Either --list_locks or a machine arg is needed." + ) + return 1 + + if len(options.args) > 1: + machine = Machine(options.args[1], options.locks_dir, auto=False) + else: + machine = None + + if options.list_locks: + FileLock.ListLock("*", options.locks_dir) + retval = True + elif options.unlock: + retval = machine.Unlock(exclusive, options.ignore_ownership) + else: + retval = machine.Lock(exclusive, options.reason) + + if retval: + return 0 + else: + return 1 + + +if __name__ == "__main__": + sys.exit(Main(sys.argv[1:])) |