Home | History | Annotate | Download | only in py_utils 
      1 # Copyright 2016 The Chromium 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 import contextlib
      6 import os
      7 
      8 LOCK_EX = None  # Exclusive lock
      9 LOCK_SH = None  # Shared lock
     10 LOCK_NB = None  # Non-blocking (LockException is raised if resource is locked)
     11 
     12 
     13 class LockException(Exception):
     14   pass
     15 
     16 
     17 if os.name == 'nt':
     18   import win32con    # pylint: disable=import-error
     19   import win32file   # pylint: disable=import-error
     20   import pywintypes  # pylint: disable=import-error
     21   LOCK_EX = win32con.LOCKFILE_EXCLUSIVE_LOCK
     22   LOCK_SH = 0  # the default
     23   LOCK_NB = win32con.LOCKFILE_FAIL_IMMEDIATELY
     24   _OVERLAPPED = pywintypes.OVERLAPPED()
     25 elif os.name == 'posix':
     26   import fcntl       # pylint: disable=import-error
     27   LOCK_EX = fcntl.LOCK_EX
     28   LOCK_SH = fcntl.LOCK_SH
     29   LOCK_NB = fcntl.LOCK_NB
     30 
     31 
     32 @contextlib.contextmanager
     33 def FileLock(target_file, flags):
     34   """ Lock the target file. Similar to AcquireFileLock but allow user to write:
     35         with FileLock(f, LOCK_EX):
     36            ...do stuff on file f without worrying about race condition
     37     Args: see AcquireFileLock's documentation.
     38   """
     39   AcquireFileLock(target_file, flags)
     40   try:
     41     yield
     42   finally:
     43     ReleaseFileLock(target_file)
     44 
     45 
     46 def AcquireFileLock(target_file, flags):
     47   """ Lock the target file. Note that if |target_file| is closed, the lock is
     48     automatically released.
     49   Args:
     50     target_file: file handle of the file to acquire lock.
     51     flags: can be any of the type LOCK_EX, LOCK_SH, LOCK_NB, or a bitwise
     52       OR combination of flags.
     53   """
     54   assert flags in (
     55       LOCK_EX, LOCK_SH, LOCK_NB, LOCK_EX | LOCK_NB, LOCK_SH | LOCK_NB)
     56   if os.name == 'nt':
     57     _LockImplWin(target_file, flags)
     58   elif os.name == 'posix':
     59     _LockImplPosix(target_file, flags)
     60   else:
     61     raise NotImplementedError('%s is not supported' % os.name)
     62 
     63 
     64 def ReleaseFileLock(target_file):
     65   """ Unlock the target file.
     66   Args:
     67     target_file: file handle of the file to release the lock.
     68   """
     69   if os.name == 'nt':
     70     _UnlockImplWin(target_file)
     71   elif os.name == 'posix':
     72     _UnlockImplPosix(target_file)
     73   else:
     74     raise NotImplementedError('%s is not supported' % os.name)
     75 
     76 # These implementations are based on
     77 # http://code.activestate.com/recipes/65203/
     78 
     79 def _LockImplWin(target_file, flags):
     80   hfile = win32file._get_osfhandle(target_file.fileno())
     81   try:
     82     win32file.LockFileEx(hfile, flags, 0, -0x10000, _OVERLAPPED)
     83   except pywintypes.error, exc_value:
     84     if exc_value[0] == 33:
     85       raise LockException('Error trying acquiring lock of %s: %s' %
     86                           (target_file.name, exc_value[2]))
     87     else:
     88       raise
     89 
     90 
     91 def _UnlockImplWin(target_file):
     92   hfile = win32file._get_osfhandle(target_file.fileno())
     93   try:
     94     win32file.UnlockFileEx(hfile, 0, -0x10000, _OVERLAPPED)
     95   except pywintypes.error, exc_value:
     96     if exc_value[0] == 158:
     97       # error: (158, 'UnlockFileEx', 'The segment is already unlocked.')
     98       # To match the 'posix' implementation, silently ignore this error
     99       pass
    100     else:
    101       # Q:  Are there exceptions/codes we should be dealing with here?
    102       raise
    103 
    104 
    105 def _LockImplPosix(target_file, flags):
    106   try:
    107     fcntl.flock(target_file.fileno(), flags)
    108   except IOError, exc_value:
    109     if exc_value[0] == 11 or exc_value[0] == 35:
    110       raise LockException('Error trying acquiring lock of %s: %s' %
    111                           (target_file.name, exc_value[1]))
    112     else:
    113       raise
    114 
    115 
    116 def _UnlockImplPosix(target_file):
    117   fcntl.flock(target_file.fileno(), fcntl.LOCK_UN)
    118