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