Split LockMachine() into 2 functions: one for lock and the other for unlock.
Also added a check before the unlock takes place to make sure the user
unlocking is the same as the user who locked the machine.
PRESUBMIT=passed
R=raymes,bjanakiraman,kbaclawski
DELTA=278 (188 added, 17 deleted, 73 changed)
OCL=55045-p2
RCL=55189-p2
RDATE=2011/08/23 10:37:54
P4 change: 42652763
diff --git a/v14/lock_machine.py b/v14/lock_machine.py
index 9f01723..a0079a8 100755
--- a/v14/lock_machine.py
+++ b/v14/lock_machine.py
@@ -9,94 +9,189 @@
__author__ = "asharif@google.com (Ahmad Sharif)"
import datetime
-import getpass
+import fcntl
import glob
import optparse
import os
-import sys
-import tc_enter_chroot
-import build_chromeos
-import setup_chromeos
+import pickle
import socket
-from utils import command_executer
-from utils import utils
+import sys
+import time
from utils import logger
-LOCK_DIR = "locks"
-LOCK_USERNAME = "mobiletc-prebuild"
-REASON_FILE = "reason.txt"
-UMASK_COMMAND = "umask a+rwx"
+class FileCreationMask(object):
+ def __init__(self, mask):
+ self._mask = mask
+
+ def __enter__(self):
+ self._old_mask = os.umask(self._mask)
+
+ def __exit__(self, type, value, traceback):
+ os.umask(self._old_mask)
-# TODO(asharif): Use duration?
-def LockMachine(machine, unlock=False, duration=None, reason=None):
- ce = command_executer.GetCommandExecuter()
- l = logger.GetLogger()
- locks_dir = os.path.join("/home", LOCK_USERNAME, LOCK_DIR)
+class LockDescription(object):
+ def __init__(self):
+ self.owner = ""
+ self.exclusive = False
+ self.counter = 0
+ self.time = 0
+ self.reason = ""
- if not os.path.exists(locks_dir):
- l.LogError("Locks dir: %s must exist" % locks_dir)
- return 1
+ def IsLocked(self):
+ return self.counter or self.exclusive
- machine_lock_dir = os.path.join(locks_dir, machine)
-
- if unlock:
- lock_program = "rm -r"
- else:
- lock_program = "%s && mkdir" % UMASK_COMMAND
- command = ("%s %s" %
- (lock_program,
- machine_lock_dir))
- retval = ce.RunCommand(command)
- if retval: return retval
-
- reason_file = os.path.join(machine_lock_dir, "reason.txt")
- if not unlock:
- if not reason:
- reason = ""
- full_reason = ("Locked by: %s@%s on %s: %s" %
- (getpass.getuser(),
- os.uname()[1],
- str(datetime.datetime.now()),
- reason))
- command = ("%s && echo \"%s\" > %s" %
- (UMASK_COMMAND,
- full_reason,
- reason_file))
- retval = ce.RunCommand(command)
-
- return 0
+ 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])
-def ListLocks(machine=None):
- if not machine:
- machine = "*"
- locks_dir = os.path.join("/home", LOCK_USERNAME, LOCK_DIR)
- print "Machine: Reason"
- print "---------------"
- for current_dir in glob.glob(os.path.join(locks_dir, machine)):
- f = open(os.path.join(current_dir, REASON_FILE))
- reason = f.read()
- reason = reason.strip()
- print "%s: %s" % (os.path.basename(current_dir), reason)
- return 0
+class FileLock(object):
+ LOCKS_DIR = "/home/mobiletc-prebuild/locks"
+
+ def __init__(self, lock_filename):
+ assert os.path.isdir(self.LOCKS_DIR), (
+ "Locks dir: %s doesn't exist!" % self.LOCKS_DIR)
+ self._filepath = os.path.join(self.LOCKS_DIR, lock_filename)
+ self._file = None
+
+ @classmethod
+ def AsString(cls, file_locks):
+ stringify_fmt = "%-30s %-15s %-4s %-4s %-15s %-40s"
+ header = stringify_fmt % ("machine", "owner", "excl", "ctr",
+ "elapsed", "reason")
+ lock_strings = []
+ for file_lock in file_locks:
+
+ elapsed_time = datetime.timedelta(
+ seconds=int(time.time() - file_lock._description.time))
+ elapsed_time = "%s ago" % elapsed_time
+ lock_strings.append(stringify_fmt %
+ (os.path.basename(file_lock._filepath),
+ file_lock._description.owner,
+ file_lock._description.exclusive,
+ file_lock._description.counter,
+ elapsed_time,
+ file_lock._description.reason))
+ table = "\n".join(lock_strings)
+ return "\n".join([header, table])
+
+ @classmethod
+ def ListLock(cls, pattern):
+ full_pattern = os.path.join(cls.LOCKS_DIR, pattern)
+ file_locks = []
+ for lock_filename in glob.glob(full_pattern):
+ 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(0000):
+ 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:
+ self._description = pickle.load(self._file)
+ except (EOFError, pickle.PickleError):
+ self._description = LockDescription()
+ return self._description
+ # Check this differently?
+ except IOError as ex:
+ logger.GetLogger().LogError(ex)
+ return None
+
+ def __exit__(self, type, value, traceback):
+ self._file.truncate(0)
+ self._file.write(pickle.dumps(self._description))
+ self._file.close()
+
+ def __str__(self):
+ return self.AsString([self])
+
+
+class Lock(object):
+ def __init__(self, to_lock):
+ self._to_lock = to_lock
+ self._logger = logger.GetLogger()
+
+ def NonBlockingLock(self, exclusive, reason=""):
+ with FileLock(self._to_lock) 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.exclusive = True
+ lock.reason = reason
+ lock.owner = os.getlogin()
+ lock.time = time.time()
+ else:
+ lock.counter += 1
+ self._logger.LogOutput("Successfully locked: %s" % self._to_lock)
+ return True
+
+ def Unlock(self, exclusive, force=False):
+ with FileLock(self._to_lock) as lock:
+ if not lock.IsLocked():
+ self._logger.LogError("Can't unlock unlocked machine!")
+ return False
+
+ if lock.exclusive != exclusive:
+ self._logger.LogError("shared locks must be unlocked with --shared")
+ return False
+
+ if lock.exclusive:
+ if lock.owner != os.getlogin() and not force:
+ self._logger.LogError("%s can't unlock lock owned by: %s" %
+ (os.getlogin(), lock.owner))
+ return False
+ lock.exclusive = False
+ lock.reason = ""
+ lock.owner = ""
+ else:
+ lock.counter -= 1
+ return True
+
+
+class Machine(object):
+ def __init__(self, name):
+ self._name = name
+ try:
+ self._full_name = socket.gethostbyaddr(name)[0]
+ except socket.error:
+ self._full_name = self._name
+
+ def Lock(self, exclusive=False, reason=""):
+ lock = Lock(self._full_name)
+ return lock.NonBlockingLock(exclusive, reason)
+
+ def Unlock(self, exclusive=False, ignore_ownership=False):
+ lock = Lock(self._full_name)
+ return lock.Unlock(exclusive, ignore_ownership)
def Main(argv):
"""The main function."""
- # Common initializations
- ce = command_executer.GetCommandExecuter()
- l = logger.GetLogger()
-
parser = optparse.OptionParser()
- parser.add_option("-m",
- "--machine",
- dest="machine",
- help="The machine to be locked.")
parser.add_option("-r",
"--reason",
dest="reason",
+ default="",
help="The lock reason.")
parser.add_option("-u",
"--unlock",
@@ -110,27 +205,45 @@
action="store_true",
default=False,
help="Use this to list locks.")
+ parser.add_option("-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_option("-s",
+ "--shared",
+ dest="shared",
+ action="store_true",
+ default=False,
+ help="Use this for a shared (non-exclusive) lock.")
- options = parser.parse_args(argv)[0]
+ options, args = parser.parse_args(argv)
- if not options.list_locks and not options.machine:
- l.LogError("Either --list_locks or --machine option is needed.")
+ exclusive = not options.shared
+
+ if not options.list_locks and len(args) != 2:
+ logger.GetLogger().LogError(
+ "Either --list_locks or a machine arg is needed.")
return 1
- machine = options.machine
- unlock = options.unlock
- reason = options.reason
-
- # Canonicalize machine name
- if machine:
- machine = socket.gethostbyaddr(machine)[0]
+ if len(args) > 1:
+ machine = Machine(args[1])
+ else:
+ machine = None
if options.list_locks:
- retval = ListLocks(machine)
+ FileLock.ListLock("*")
+ retval = True
+ elif options.unlock:
+ retval = machine.Unlock(exclusive, options.ignore_ownership)
else:
- retval = LockMachine(machine, unlock=unlock, reason=reason)
- return retval
+ retval = machine.Lock(exclusive, options.reason)
+
+ if retval:
+ return 0
+ else:
+ return 1
if __name__ == "__main__":
- retval = Main(sys.argv)
- sys.exit(retval)
+ sys.exit(Main(sys.argv))
diff --git a/v14/lock_machine_test.py b/v14/lock_machine_test.py
new file mode 100644
index 0000000..8c00c40
--- /dev/null
+++ b/v14/lock_machine_test.py
@@ -0,0 +1,58 @@
+#!/usr/bin/python2.6
+#
+# Copyright 2010 Google Inc. All Rights Reserved.
+
+"""lock_machine.py related unit-tests.
+
+MachineManagerTest tests MachineManager.
+"""
+
+__author__ = "asharif@google.com (Ahmad Sharif)"
+
+
+import lock_machine
+import unittest
+
+
+class MachineTest(unittest.TestCase):
+ def setUp(self):
+ pass
+
+
+ def testRepeatedUnlock(self):
+ mach = lock_machine.Machine("qqqraymes.mtv")
+ for i in range(10):
+ self.assertFalse(mach.Unlock())
+
+ def testLockUnlock(self):
+ mach = lock_machine.Machine("otter.mtv")
+ for i in range(10):
+ self.assertTrue(mach.Lock(exclusive=True))
+ self.assertTrue(mach.Unlock(exclusive=True))
+
+ def testSharedLock(self):
+ mach = lock_machine.Machine("chrotomation.mtv")
+ for i in range(10):
+ self.assertTrue(mach.Lock(exclusive=False))
+ for i in range(10):
+ self.assertTrue(mach.Unlock(exclusive=False))
+ self.assertTrue(mach.Lock(exclusive=True))
+ self.assertTrue(mach.Unlock(exclusive=True))
+
+ def testExclusiveLock(self):
+ mach = lock_machine.Machine("atree.mtv")
+ self.assertTrue(mach.Lock(exclusive=True))
+ for i in range(10):
+ self.assertFalse(mach.Lock(exclusive=True))
+ self.assertFalse(mach.Lock(exclusive=False))
+ self.assertTrue(mach.Unlock(exclusive=True))
+
+ def testExclusiveState(self):
+ mach = lock_machine.Machine("testExclusiveState")
+ self.assertTrue(mach.Lock(exclusive=True))
+ for i in range(10):
+ self.assertFalse(mach.Lock(exclusive=False))
+ self.assertTrue(mach.Unlock(exclusive=True))
+
+if __name__ == "__main__":
+ unittest.main()