resource.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238
  1. """Generic resource pool implementation."""
  2. from __future__ import absolute_import, unicode_literals
  3. import os
  4. from collections import deque
  5. from . import exceptions
  6. from .five import Empty, LifoQueue as _LifoQueue
  7. from .utils.compat import register_after_fork
  8. from .utils.functional import lazy
  9. def _after_fork_cleanup_resource(resource):
  10. try:
  11. resource.force_close_all()
  12. except Exception:
  13. pass
  14. class LifoQueue(_LifoQueue):
  15. """Last in first out version of Queue."""
  16. def _init(self, maxsize):
  17. self.queue = deque()
  18. class Resource(object):
  19. """Pool of resources."""
  20. LimitExceeded = exceptions.LimitExceeded
  21. close_after_fork = False
  22. def __init__(self, limit=None, preload=None, close_after_fork=None):
  23. self._limit = limit
  24. self.preload = preload or 0
  25. self._closed = False
  26. self.close_after_fork = (
  27. close_after_fork
  28. if close_after_fork is not None else self.close_after_fork
  29. )
  30. self._resource = LifoQueue()
  31. self._dirty = set()
  32. if self.close_after_fork and register_after_fork is not None:
  33. register_after_fork(self, _after_fork_cleanup_resource)
  34. self.setup()
  35. def setup(self):
  36. raise NotImplementedError('subclass responsibility')
  37. def _add_when_empty(self):
  38. if self.limit and len(self._dirty) >= self.limit:
  39. raise self.LimitExceeded(self.limit)
  40. # All taken, put new on the queue and
  41. # try get again, this way the first in line
  42. # will get the resource.
  43. self._resource.put_nowait(self.new())
  44. def acquire(self, block=False, timeout=None):
  45. """Acquire resource.
  46. Arguments:
  47. block (bool): If the limit is exceeded,
  48. then block until there is an available item.
  49. timeout (float): Timeout to wait
  50. if ``block`` is true. Default is :const:`None` (forever).
  51. Raises:
  52. LimitExceeded: if block is false and the limit has been exceeded.
  53. """
  54. if self._closed:
  55. raise RuntimeError('Acquire on closed pool')
  56. if self.limit:
  57. while 1:
  58. try:
  59. R = self._resource.get(block=block, timeout=timeout)
  60. except Empty:
  61. self._add_when_empty()
  62. else:
  63. try:
  64. R = self.prepare(R)
  65. except BaseException:
  66. if isinstance(R, lazy):
  67. # not evaluated yet, just put it back
  68. self._resource.put_nowait(R)
  69. else:
  70. # evaluted so must try to release/close first.
  71. self.release(R)
  72. raise
  73. self._dirty.add(R)
  74. break
  75. else:
  76. R = self.prepare(self.new())
  77. def release():
  78. """Release resource so it can be used by another thread.
  79. Warnings:
  80. The caller is responsible for discarding the object,
  81. and to never use the resource again. A new resource must
  82. be acquired if so needed.
  83. """
  84. self.release(R)
  85. R.release = release
  86. return R
  87. def prepare(self, resource):
  88. return resource
  89. def close_resource(self, resource):
  90. resource.close()
  91. def release_resource(self, resource):
  92. pass
  93. def replace(self, resource):
  94. """Replace existing resource with a new instance.
  95. This can be used in case of defective resources.
  96. """
  97. if self.limit:
  98. self._dirty.discard(resource)
  99. self.close_resource(resource)
  100. def release(self, resource):
  101. if self.limit:
  102. self._dirty.discard(resource)
  103. self._resource.put_nowait(resource)
  104. self.release_resource(resource)
  105. else:
  106. self.close_resource(resource)
  107. def collect_resource(self, resource):
  108. pass
  109. def force_close_all(self):
  110. """Close and remove all resources in the pool (also those in use).
  111. Used to close resources from parent processes after fork
  112. (e.g. sockets/connections).
  113. """
  114. if self._closed:
  115. return
  116. self._closed = True
  117. dirty = self._dirty
  118. resource = self._resource
  119. while 1: # - acquired
  120. try:
  121. dres = dirty.pop()
  122. except KeyError:
  123. break
  124. try:
  125. self.collect_resource(dres)
  126. except AttributeError: # Issue #78
  127. pass
  128. while 1: # - available
  129. # deque supports '.clear', but lists do not, so for that
  130. # reason we use pop here, so that the underlying object can
  131. # be any object supporting '.pop' and '.append'.
  132. try:
  133. res = resource.queue.pop()
  134. except IndexError:
  135. break
  136. try:
  137. self.collect_resource(res)
  138. except AttributeError:
  139. pass # Issue #78
  140. def resize(self, limit, force=False, ignore_errors=False, reset=False):
  141. prev_limit = self._limit
  142. if (self._dirty and 0 < limit < self._limit) and not ignore_errors:
  143. if not force:
  144. raise RuntimeError(
  145. "Can't shrink pool when in use: was={0} now={1}".format(
  146. self._limit, limit))
  147. reset = True
  148. self._limit = limit
  149. if reset:
  150. try:
  151. self.force_close_all()
  152. except Exception:
  153. pass
  154. self.setup()
  155. if limit < prev_limit:
  156. self._shrink_down(collect=limit > 0)
  157. def _shrink_down(self, collect=True):
  158. class Noop:
  159. def __enter__(self):
  160. pass
  161. def __exit__(self, type, value, traceback):
  162. pass
  163. resource = self._resource
  164. # Items to the left are last recently used, so we remove those first.
  165. with getattr(resource, 'mutex', Noop()):
  166. while len(resource.queue) > self.limit:
  167. R = resource.queue.popleft()
  168. if collect:
  169. self.collect_resource(R)
  170. @property
  171. def limit(self):
  172. return self._limit
  173. @limit.setter
  174. def limit(self, limit):
  175. self.resize(limit)
  176. if os.environ.get('KOMBU_DEBUG_POOL'): # pragma: no cover
  177. _orig_acquire = acquire
  178. _orig_release = release
  179. _next_resource_id = 0
  180. def acquire(self, *args, **kwargs): # noqa
  181. import traceback
  182. id = self._next_resource_id = self._next_resource_id + 1
  183. print('+{0} ACQUIRE {1}'.format(id, self.__class__.__name__))
  184. r = self._orig_acquire(*args, **kwargs)
  185. r._resource_id = id
  186. print('-{0} ACQUIRE {1}'.format(id, self.__class__.__name__))
  187. if not hasattr(r, 'acquired_by'):
  188. r.acquired_by = []
  189. r.acquired_by.append(traceback.format_stack())
  190. return r
  191. def release(self, resource): # noqa
  192. id = resource._resource_id
  193. print('+{0} RELEASE {1}'.format(id, self.__class__.__name__))
  194. r = self._orig_release(resource)
  195. print('-{0} RELEASE {1}'.format(id, self.__class__.__name__))
  196. self._next_resource_id -= 1
  197. return r