Implement new global locking scheme for machines.

This CL implements a new machine locking mechanism using
Autotest Front End servers.  When locking/unlocking a lab machine,
it uses the ChromeOS HW lab server; when locking/unlocking a local
machine, it uses a local AFE server on chrotomation2.

BUG=None
TEST=Tested the script manually to lock/unlock machines and query status.
Also tested with image_chromeos and with crosperf scripts.

Change-Id: I2793bc1f7dc056e725694e81ded656d9f49d227b
Reviewed-on: https://chrome-internal-review.googlesource.com/217560
Reviewed-by: Luis Lozano <llozano@chromium.org>
Reviewed-by: David Sharp <dhsharp@google.com>
Tested-by: Caroline Tice <cmtice@google.com>
Commit-Queue: Caroline Tice <cmtice@google.com>

1 files changed, 348 insertions, 0 deletions

diff --git a/file_lock_machine.py b/file_lock_machine.py
new file mode 100755
index 00000000..1496a2c7
--- /dev/null
+++ b/file_lock_machine.py
@@ -0,0 +1,348 @@
+#!/usr/bin/python
+#
+# Copyright 2010 Google Inc. All Rights Reserved.
+
+"""Script to lock/unlock machines."""
+
+__author__ = "asharif@google.com (Ahmad Sharif)"
+
+import datetime
+import fcntl
+import getpass
+import glob
+import json
+import optparse
+import os
+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):
+    self._mask = mask
+
+  def __enter__(self):
+    self._old_mask = os.umask(self._mask)
+
+  def __exit__(self, type, value, traceback):
+    os.umask(self._old_mask)
+
+
+class LockDescription(object):
+  """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])
+
+
+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
+
+  @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._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,
+                           file_lock._description.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(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:
+          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, type, 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):
+  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!")
+        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):
+  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, 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, self._auto)
+    return lock.Unlock(exclusive, ignore_ownership)
+
+
+def Main(argv):
+  """The main function."""
+  parser = optparse.OptionParser()
+  parser.add_option("-r",
+                    "--reason",
+                    dest="reason",
+                    default="",
+                    help="The lock reason.")
+  parser.add_option("-u",
+                    "--unlock",
+                    dest="unlock",
+                    action="store_true",
+                    default=False,
+                    help="Use this to unlock.")
+  parser.add_option("-l",
+                    "--list_locks",
+                    dest="list_locks",
+                    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.")
+  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:
+    logger.GetLogger().LogError(
+        "Either --list_locks or a machine arg is needed.")
+    return 1
+
+  if len(args) > 1:
+    machine = Machine(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))