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
5import contextlib
6import os
7
8LOCK_EX = None  # Exclusive lock
9LOCK_SH = None  # Shared lock
10LOCK_NB = None  # Non-blocking (LockException is raised if resource is locked)
11
12
13class LockException(Exception):
14  pass
15
16
17if 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()
25elif 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
33def 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
46def 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
64def 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
79def _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
91def _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
105def _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
116def _UnlockImplPosix(target_file):
117  fcntl.flock(target_file.fileno(), fcntl.LOCK_UN)
118