pool.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. # Copyright 2015 Yahoo.com
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import collections
  15. import contextlib
  16. import sys
  17. import threading
  18. import time
  19. import six
  20. class ObjectPool(object):
  21. """A pool of objects that release/creates/destroys as needed."""
  22. def __init__(self, obj_creator,
  23. after_remove=None, max_size=None,
  24. idle_timeout=0,
  25. lock_generator=None):
  26. self._used_objs = collections.deque()
  27. self._free_objs = collections.deque()
  28. self._obj_creator = obj_creator
  29. if lock_generator is None:
  30. self._lock = threading.Lock()
  31. else:
  32. self._lock = lock_generator()
  33. self._after_remove = after_remove
  34. max_size = max_size or 2 ** 31
  35. if not isinstance(max_size, six.integer_types) or max_size < 0:
  36. raise ValueError('"max_size" must be a positive integer')
  37. self.max_size = max_size
  38. self.idle_timeout = idle_timeout
  39. self._idle_clock = time.time if idle_timeout else int
  40. @property
  41. def used(self):
  42. return tuple(self._used_objs)
  43. @property
  44. def free(self):
  45. return tuple(self._free_objs)
  46. @contextlib.contextmanager
  47. def get_and_release(self, destroy_on_fail=False):
  48. obj = self.get()
  49. try:
  50. yield obj
  51. except Exception:
  52. exc_info = sys.exc_info()
  53. if not destroy_on_fail:
  54. self.release(obj)
  55. else:
  56. self.destroy(obj)
  57. six.reraise(exc_info[0], exc_info[1], exc_info[2])
  58. self.release(obj)
  59. def get(self):
  60. with self._lock:
  61. # Find a free object, removing any that have idled for too long.
  62. now = self._idle_clock()
  63. while self._free_objs:
  64. obj = self._free_objs.popleft()
  65. if now - obj._last_used <= self.idle_timeout:
  66. break
  67. if self._after_remove is not None:
  68. self._after_remove(obj)
  69. else:
  70. # No free objects, create a new one.
  71. curr_count = len(self._used_objs)
  72. if curr_count >= self.max_size:
  73. raise RuntimeError("Too many objects,"
  74. " %s >= %s" % (curr_count,
  75. self.max_size))
  76. obj = self._obj_creator()
  77. self._used_objs.append(obj)
  78. obj._last_used = now
  79. return obj
  80. def destroy(self, obj, silent=True):
  81. was_dropped = False
  82. with self._lock:
  83. try:
  84. self._used_objs.remove(obj)
  85. was_dropped = True
  86. except ValueError:
  87. if not silent:
  88. raise
  89. if was_dropped and self._after_remove is not None:
  90. self._after_remove(obj)
  91. def release(self, obj, silent=True):
  92. with self._lock:
  93. try:
  94. self._used_objs.remove(obj)
  95. self._free_objs.append(obj)
  96. obj._last_used = self._idle_clock()
  97. except ValueError:
  98. if not silent:
  99. raise
  100. def clear(self):
  101. if self._after_remove is not None:
  102. needs_destroy = []
  103. with self._lock:
  104. needs_destroy.extend(self._used_objs)
  105. needs_destroy.extend(self._free_objs)
  106. self._free_objs.clear()
  107. self._used_objs.clear()
  108. for obj in needs_destroy:
  109. self._after_remove(obj)
  110. else:
  111. with self._lock:
  112. self._free_objs.clear()
  113. self._used_objs.clear()