src/third_party/chromite/lib/paygen/flock.py

   1 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """A library for managing file locks."""
   6
   7 from __future__ import print_function
   8
   9 import errno
  10 import fcntl
  11 import os
  12
  13
  14 LOCK_DIR = '/tmp/run_once'
  15
  16
  17 class LockNotAcquired(Exception):
  18   """Raised when the run_once lock is already held by another lock object.
  19
  20   Note that this can happen even within the same process.  A second lock
  21   object targeting the same lock file will fail to acquire the lock regardless
  22   of process.
  23
  24   self.lock_file_path has path to lock file involved.
  25   self.owner_pid has pid of process that currently has lock.
  26   """
  27
  28   def __init__(self, lock_file_path, owner_pid, *args, **kwargs):
  29     Exception.__init__(self, *args, **kwargs)
  30     self.lock_file_path = lock_file_path
  31     self.owner_pid = owner_pid
  32
  33   def __str__(self):
  34     return "Lock (%s) held by pid %s" % (self.lock_file_path, self.owner_pid)
  35
  36
  37 class Lock(object):
  38   """This class grabs an exclusive flock on a file in a specified directory.
  39
  40   This class can be used in combination with the "with" statement.
  41
  42   Because the lock is associated with a file descriptor, the lock will
  43   continue to be held for as long as the file descriptor is open (even
  44   in subprocesses, exec'd executables, etc).
  45
  46   For informational purposes only, the pid of the current process is
  47   written into the lock file when it is held, but it's never removed.
  48   """
  49
  50   def __init__(self, lock_name, lock_dir=None, blocking=False, shared=False):
  51     """Setup our lock class, but don't do any work yet (or grab the lock).
  52
  53     Args:
  54       lock_name: The file name of the lock to file. If it's a relative name,
  55                  it will be expanded based on lock_dir.
  56       lock_dir: is the directory in which to create lock files, it defaults
  57                 to LOCK_DIR.
  58       blocking: When trying to acquire a lock, do we block until it's
  59                 available, or raise "LockNotAcquired"? If we block,
  60                 there is no timeout.
  61       shared: A value of False means get an exclusive lock, and True
  62               means to get a shared lock.
  63     """
  64     if lock_dir == None:
  65       lock_dir = LOCK_DIR
  66
  67     # os.path.join will ignore lock_dir, if lock_name is an absolute path.
  68     self.lock_file = os.path.join(lock_dir, lock_name)
  69
  70     self._blocking = blocking
  71     self._shared = shared
  72     self._fd = None
  73
  74   def Acquire(self):
  75     """Acquire the flock.
  76
  77     It's safe to call this multiple times, though the first Unlock will
  78     release the lock.
  79     """
  80     # Create the directory for our lock files if it doesn't already exist
  81     try:
  82       os.makedirs(os.path.dirname(self.lock_file))
  83     except OSError as e:
  84       if e.errno is not errno.EEXIST:
  85         raise
  86
  87     if not self._fd:
  88       self._fd = open(self.lock_file, 'a')
  89
  90     try:
  91       if self._shared:
  92         flags = fcntl.LOCK_SH
  93       else:
  94         flags = fcntl.LOCK_EX
  95
  96       if not self._blocking:
  97         flags |= fcntl.LOCK_NB
  98
  99       fcntl.flock(self._fd, flags)
 100
 101       # We have the lock, write our pid into it for informational purposes.
 102       self._fd.truncate(0)
 103       self._fd.write(str(os.getpid())+'\n')
 104       self._fd.flush()
 105
 106     except IOError as e:
 107       self._fd.close()
 108       self._fd = None
 109
 110       # We got the error that someone else already held the locked.
 111       # Can only happen if we are blocking == False.
 112       if e.errno == errno.EAGAIN:
 113         # To be helpful, grab pid of owner process from file.
 114         owner_pid = None
 115         with open(self.lock_file, 'r') as f:
 116           owner_pid = f.read().strip()
 117
 118         raise LockNotAcquired(self.lock_file, owner_pid)
 119
 120       # Pass along any other error for debugging
 121       raise
 122
 123   def Release(self):
 124     """Release the flock."""
 125
 126     if self._fd:
 127       fcntl.flock(self._fd, fcntl.LOCK_UN)
 128       self._fd.close()
 129       self._fd = None
 130
 131   def IsLocked(self):
 132     """Return True if lock is currently acquired."""
 133     return bool(self._fd)
 134
 135   # Lock objects can be used with "with" statements.
 136   def __enter__(self):
 137     self.Acquire()
 138     return self
 139
 140   def __exit__(self, _type, _value, _traceback):
 141     self.Release()
 142
 143
 144 def ExecWithLock(cmd, lock_name=None, lock_dir=None, blocking=False):
 145   """Helper method that execs another program with an flock.
 146
 147   Args:
 148     cmd: The command to run through flock.
 149     lock_name: defaults to the name of the command.
 150     lock_dir: defaults to LOCK_DIR.
 151     blocking: Whether to take a blocking lock.
 152
 153   Raises:
 154     LockNotAcquired: If the lock wasn't acquired
 155   """
 156   if not lock_name:
 157     lock_name = os.path.basename(cmd[0])
 158
 159   with Lock(lock_name, lock_dir=lock_dir, blocking=blocking):
 160     # Our lock file is locked, exec our subprocess. It has an extra fd
 161     # in it's environment, and the lock on that fd will be held until
 162     # that fd is closed on exit.
 163     os.execvp(cmd[0], cmd)
 164
 165     # Note that the above new process will not return here, which has
 166     # the effect of never exiting this 'with' context, which means
 167     # Lock.Unlock() is never called. The lock is released all the same.