lock.py 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260
  1. # Copyright (c) 2009-2012 Denis Bilenko. See LICENSE for details.
  2. """Locking primitives"""
  3. from __future__ import absolute_import
  4. from gevent.hub import getcurrent
  5. from gevent._compat import PYPY
  6. from gevent._semaphore import Semaphore, BoundedSemaphore # pylint:disable=no-name-in-module,import-error
  7. __all__ = [
  8. 'Semaphore',
  9. 'DummySemaphore',
  10. 'BoundedSemaphore',
  11. 'RLock',
  12. ]
  13. # On PyPy, we don't compile the Semaphore class with Cython. Under
  14. # Cython, each individual method holds the GIL for its entire
  15. # duration, ensuring that no other thread can interrupt us in an
  16. # unsafe state (only when we _do_wait do we call back into Python and
  17. # allow switching threads). Simulate that here through the use of a manual
  18. # lock. (We use a separate lock for each semaphore to allow sys.settrace functions
  19. # to use locks *other* than the one being traced.)
  20. if PYPY:
  21. # TODO: Need to use monkey.get_original?
  22. try:
  23. from _thread import allocate_lock as _allocate_lock # pylint:disable=import-error,useless-suppression
  24. from _thread import get_ident as _get_ident # pylint:disable=import-error,useless-suppression
  25. except ImportError:
  26. # Python 2
  27. from thread import allocate_lock as _allocate_lock # pylint:disable=import-error,useless-suppression
  28. from thread import get_ident as _get_ident # pylint:disable=import-error,useless-suppression
  29. _sem_lock = _allocate_lock()
  30. def untraceable(f):
  31. # Don't allow re-entry to these functions in a single thread, as can
  32. # happen if a sys.settrace is used
  33. def wrapper(self):
  34. me = _get_ident()
  35. try:
  36. count = self._locking[me]
  37. except KeyError:
  38. count = self._locking[me] = 1
  39. else:
  40. count = self._locking[me] = count + 1
  41. if count:
  42. return
  43. try:
  44. return f(self)
  45. finally:
  46. count = count - 1
  47. if not count:
  48. del self._locking[me]
  49. else:
  50. self._locking[me] = count
  51. return wrapper
  52. class _OwnedLock(object):
  53. def __init__(self):
  54. self._owner = None
  55. self._block = _allocate_lock()
  56. self._locking = {}
  57. self._count = 0
  58. @untraceable
  59. def acquire(self):
  60. me = _get_ident()
  61. if self._owner == me:
  62. self._count += 1
  63. return
  64. self._owner = me
  65. self._block.acquire()
  66. self._count = 1
  67. @untraceable
  68. def release(self):
  69. self._count = count = self._count - 1
  70. if not count:
  71. self._block.release()
  72. self._owner = None
  73. # acquire, wait, and release all acquire the lock on entry and release it
  74. # on exit. acquire and wait can call _do_wait, which must release it on entry
  75. # and re-acquire it for them on exit.
  76. class _around(object):
  77. __slots__ = ('before', 'after')
  78. def __init__(self, before, after):
  79. self.before = before
  80. self.after = after
  81. def __enter__(self):
  82. self.before()
  83. def __exit__(self, t, v, tb):
  84. self.after()
  85. def _decorate(func, cmname):
  86. # functools.wrap?
  87. def wrapped(self, *args, **kwargs):
  88. with getattr(self, cmname):
  89. return func(self, *args, **kwargs)
  90. return wrapped
  91. Semaphore._py3k_acquire = Semaphore.acquire = _decorate(Semaphore.acquire, '_lock_locked')
  92. Semaphore.release = _decorate(Semaphore.release, '_lock_locked')
  93. Semaphore.wait = _decorate(Semaphore.wait, '_lock_locked')
  94. Semaphore._do_wait = _decorate(Semaphore._do_wait, '_lock_unlocked')
  95. _Sem_init = Semaphore.__init__
  96. def __init__(self, *args, **kwargs):
  97. l = self._lock_lock = _OwnedLock()
  98. self._lock_locked = _around(l.acquire, l.release)
  99. self._lock_unlocked = _around(l.release, l.acquire)
  100. _Sem_init(self, *args, **kwargs)
  101. Semaphore.__init__ = __init__
  102. del _decorate
  103. del untraceable
  104. class DummySemaphore(object):
  105. """
  106. DummySemaphore(value=None) -> DummySemaphore
  107. A Semaphore initialized with "infinite" initial value. None of its
  108. methods ever block.
  109. This can be used to parameterize on whether or not to actually
  110. guard access to a potentially limited resource. If the resource is
  111. actually limited, such as a fixed-size thread pool, use a real
  112. :class:`Semaphore`, but if the resource is unbounded, use an
  113. instance of this class. In that way none of the supporting code
  114. needs to change.
  115. Similarly, it can be used to parameterize on whether or not to
  116. enforce mutual exclusion to some underlying object. If the
  117. underlying object is known to be thread-safe itself mutual
  118. exclusion is not needed and a ``DummySemaphore`` can be used, but
  119. if that's not true, use a real ``Semaphore``.
  120. """
  121. # Internally this is used for exactly the purpose described in the
  122. # documentation. gevent.pool.Pool uses it instead of a Semaphore
  123. # when the pool size is unlimited, and
  124. # gevent.fileobject.FileObjectThread takes a parameter that
  125. # determines whether it should lock around IO to the underlying
  126. # file object.
  127. def __init__(self, value=None):
  128. """
  129. .. versionchanged:: 1.1rc3
  130. Accept and ignore a *value* argument for compatibility with Semaphore.
  131. """
  132. pass
  133. def __str__(self):
  134. return '<%s>' % self.__class__.__name__
  135. def locked(self):
  136. """A DummySemaphore is never locked so this always returns False."""
  137. return False
  138. def release(self):
  139. """Releasing a dummy semaphore does nothing."""
  140. pass
  141. def rawlink(self, callback):
  142. # XXX should still work and notify?
  143. pass
  144. def unlink(self, callback):
  145. pass
  146. def wait(self, timeout=None):
  147. """Waiting for a DummySemaphore returns immediately."""
  148. pass
  149. def acquire(self, blocking=True, timeout=None):
  150. """
  151. A DummySemaphore can always be acquired immediately so this always
  152. returns True and ignores its arguments.
  153. .. versionchanged:: 1.1a1
  154. Always return *true*.
  155. """
  156. # pylint:disable=unused-argument
  157. return True
  158. def __enter__(self):
  159. pass
  160. def __exit__(self, typ, val, tb):
  161. pass
  162. class RLock(object):
  163. def __init__(self):
  164. self._block = Semaphore(1)
  165. self._owner = None
  166. self._count = 0
  167. def __repr__(self):
  168. return "<%s at 0x%x _block=%s _count=%r _owner=%r)>" % (
  169. self.__class__.__name__,
  170. id(self),
  171. self._block,
  172. self._count,
  173. self._owner)
  174. def acquire(self, blocking=1):
  175. me = getcurrent()
  176. if self._owner is me:
  177. self._count = self._count + 1
  178. return 1
  179. rc = self._block.acquire(blocking)
  180. if rc:
  181. self._owner = me
  182. self._count = 1
  183. return rc
  184. def __enter__(self):
  185. return self.acquire()
  186. def release(self):
  187. if self._owner is not getcurrent():
  188. raise RuntimeError("cannot release un-aquired lock")
  189. self._count = count = self._count - 1
  190. if not count:
  191. self._owner = None
  192. self._block.release()
  193. def __exit__(self, typ, value, tb):
  194. self.release()
  195. # Internal methods used by condition variables
  196. def _acquire_restore(self, count_owner):
  197. count, owner = count_owner
  198. self._block.acquire()
  199. self._count = count
  200. self._owner = owner
  201. def _release_save(self):
  202. count = self._count
  203. self._count = 0
  204. owner = self._owner
  205. self._owner = None
  206. self._block.release()
  207. return (count, owner)
  208. def _is_owned(self):
  209. return self._owner is getcurrent()