Synced repos to: 64740

1 files changed, 132 insertions, 34 deletions

diff --git a/lock_machine.py b/lock_machine.py
index c5f98092..0f948c3d 100755
--- a/lock_machine.py
+++ b/lock_machine.py
@@ -1,10 +1,8 @@
-#!/usr/bin/python2.6
+#!/usr/bin/python
 #
 # Copyright 2010 Google Inc. All Rights Reserved.
 
-"""Script to lock/unlock machines.
-
-"""
+"""Script to lock/unlock machines."""
 
 __author__ = "asharif@google.com (Ahmad Sharif)"
 
@@ -12,14 +10,31 @@ import datetime
 import fcntl
 import getpass
 import glob
+import json
 import optparse
 import os
-import pickle
 import socket
 import sys
 import time
+
 from utils import logger
 
+LOCK_SUFFIX = "_check_lock_liveness"
+
+
+def FileCheckName(name):
+  return name + LOCK_SUFFIX
+
+
+def OpenLiveCheck(file_name):
+  with FileCreationMask(0000):
+    fd = open(file_name, "a+w")
+  try:
+    fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+  except IOError:
+    raise
+  return fd
+
 
 class FileCreationMask(object):
   def __init__(self, mask):
@@ -33,12 +48,23 @@ class FileCreationMask(object):
 
 
 class LockDescription(object):
-  def __init__(self):
-    self.owner = ""
-    self.exclusive = False
-    self.counter = 0
-    self.time = 0
-    self.reason = ""
+  """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
@@ -48,23 +74,26 @@ class LockDescription(object):
                      "Exclusive: %s" % self.exclusive,
                      "Counter: %s" % self.counter,
                      "Time: %s" % self.time,
-                     "Reason: %s" % self.reason])
+                     "Reason: %s" % self.reason,
+                     "Auto: %s" % self.auto])
 
 
 class FileLock(object):
-  LOCKS_DIR = "/home/mobiletc-prebuild/locks"
+  """File lock operation class."""
+  FILE_OPS = []
 
   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._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
 
   @classmethod
   def AsString(cls, file_locks):
-    stringify_fmt = "%-30s %-15s %-4s %-4s %-15s %-40s"
+    stringify_fmt = "%-30s %-15s %-4s %-4s %-15s %-40s %-4s"
     header = stringify_fmt % ("machine", "owner", "excl", "ctr",
-                              "elapsed", "reason")
+                              "elapsed", "reason", "auto")
     lock_strings = []
     for file_lock in file_locks:
 
@@ -77,15 +106,20 @@ class FileLock(object):
                            file_lock._description.exclusive,
                            file_lock._description.counter,
                            elapsed_time,
-                           file_lock._description.reason))
+                           file_lock._description.reason,
+                           file_lock._description.auto))
     table = "\n".join(lock_strings)
     return "\n".join([header, table])
 
   @classmethod
-  def ListLock(cls, pattern):
-    full_pattern = os.path.join(cls.LOCKS_DIR, pattern)
+  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():
@@ -102,9 +136,26 @@ class FileLock(object):
           raise IOError("flock(%s, LOCK_EX) failed!" % self._filepath)
 
         try:
-          self._description = pickle.load(self._file)
-        except (EOFError, pickle.PickleError):
-          self._description = LockDescription()
+          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:
@@ -113,7 +164,7 @@ class FileLock(object):
 
   def __exit__(self, type, value, traceback):
     self._file.truncate(0)
-    self._file.write(pickle.dumps(self._description))
+    self._file.write(json.dumps(self._description.__dict__, skipkeys=True))
     self._file.close()
 
   def __str__(self):
@@ -121,12 +172,14 @@ class FileLock(object):
 
 
 class Lock(object):
-  def __init__(self, to_lock):
-    self._to_lock = to_lock
+  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._to_lock) as lock:
+    with FileLock(self._lock_file) as lock:
       if lock.exclusive:
         self._logger.LogError(
             "Exclusive lock already acquired by %s. Reason: %s" %
@@ -137,17 +190,22 @@ class Lock(object):
         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._to_lock) as lock:
+    with FileLock(self._lock_file) as lock:
       if not lock.IsLocked():
         self._logger.LogError("Can't unlock unlocked machine!")
         return False
@@ -161,28 +219,61 @@ class Lock(object):
           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):
-  def __init__(self, name):
+  LOCKS_DIR = "/home/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)
+    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 not 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)
+    lock = Lock(self._full_name, self._auto)
     return lock.Unlock(exclusive, ignore_ownership)
 
 
@@ -218,9 +309,16 @@ def Main(argv):
                     action="store_true",
                     default=False,
                     help="Use this for a shared (non-exclusive) lock.")
+  parser.add_option("-d",
+                    "--dir",
+                    dest="locks_dir",
+                    action="store",
+                    default=Machine.LOCKS_DIR,
+                    help="Use this to set different locks_dir")
 
   options, args = parser.parse_args(argv)
 
+  options.locks_dir = os.path.abspath(options.locks_dir)
   exclusive = not options.shared
 
   if not options.list_locks and len(args) != 2:
@@ -229,12 +327,12 @@ def Main(argv):
     return 1
 
   if len(args) > 1:
-    machine = Machine(args[1])
+    machine = Machine(args[1], options.locks_dir, auto=False)
   else:
     machine = None
 
   if options.list_locks:
-    FileLock.ListLock("*")
+    FileLock.ListLock("*", options.locks_dir)
     retval = True
   elif options.unlock:
     retval = machine.Unlock(exclusive, options.ignore_ownership)