pools.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. """Public resource pools."""
  2. from __future__ import absolute_import, unicode_literals
  3. import os
  4. from itertools import chain
  5. from .connection import Resource
  6. from .five import range, values
  7. from .messaging import Producer
  8. from .utils.collections import EqualityDict
  9. from .utils.compat import register_after_fork
  10. from .utils.functional import lazy
  11. __all__ = ('ProducerPool', 'PoolGroup', 'register_group',
  12. 'connections', 'producers', 'get_limit', 'set_limit', 'reset')
  13. _limit = [10]
  14. _groups = []
  15. use_global_limit = object()
  16. disable_limit_protection = os.environ.get('KOMBU_DISABLE_LIMIT_PROTECTION')
  17. def _after_fork_cleanup_group(group):
  18. group.clear()
  19. class ProducerPool(Resource):
  20. """Pool of :class:`kombu.Producer` instances."""
  21. Producer = Producer
  22. close_after_fork = True
  23. def __init__(self, connections, *args, **kwargs):
  24. self.connections = connections
  25. self.Producer = kwargs.pop('Producer', None) or self.Producer
  26. super(ProducerPool, self).__init__(*args, **kwargs)
  27. def _acquire_connection(self):
  28. return self.connections.acquire(block=True)
  29. def create_producer(self):
  30. conn = self._acquire_connection()
  31. try:
  32. return self.Producer(conn)
  33. except BaseException:
  34. conn.release()
  35. raise
  36. def new(self):
  37. return lazy(self.create_producer)
  38. def setup(self):
  39. if self.limit:
  40. for _ in range(self.limit):
  41. self._resource.put_nowait(self.new())
  42. def close_resource(self, resource):
  43. pass
  44. def prepare(self, p):
  45. if callable(p):
  46. p = p()
  47. if p._channel is None:
  48. conn = self._acquire_connection()
  49. try:
  50. p.revive(conn)
  51. except BaseException:
  52. conn.release()
  53. raise
  54. return p
  55. def release(self, resource):
  56. if resource.__connection__:
  57. resource.__connection__.release()
  58. resource.channel = None
  59. super(ProducerPool, self).release(resource)
  60. class PoolGroup(EqualityDict):
  61. """Collection of resource pools."""
  62. def __init__(self, limit=None, close_after_fork=True):
  63. self.limit = limit
  64. self.close_after_fork = close_after_fork
  65. if self.close_after_fork and register_after_fork is not None:
  66. register_after_fork(self, _after_fork_cleanup_group)
  67. def create(self, resource, limit):
  68. raise NotImplementedError('PoolGroups must define ``create``')
  69. def __missing__(self, resource):
  70. limit = self.limit
  71. if limit is use_global_limit:
  72. limit = get_limit()
  73. k = self[resource] = self.create(resource, limit)
  74. return k
  75. def register_group(group):
  76. """Register group (can be used as decorator)."""
  77. _groups.append(group)
  78. return group
  79. class Connections(PoolGroup):
  80. """Collection of connection pools."""
  81. def create(self, connection, limit):
  82. return connection.Pool(limit=limit)
  83. connections = register_group(Connections(limit=use_global_limit)) # noqa: E305
  84. class Producers(PoolGroup):
  85. """Collection of producer pools."""
  86. def create(self, connection, limit):
  87. return ProducerPool(connections[connection], limit=limit)
  88. producers = register_group(Producers(limit=use_global_limit)) # noqa: E305
  89. def _all_pools():
  90. return chain(*[(values(g) if g else iter([])) for g in _groups])
  91. def get_limit():
  92. """Get current connection pool limit."""
  93. return _limit[0]
  94. def set_limit(limit, force=False, reset_after=False, ignore_errors=False):
  95. """Set new connection pool limit."""
  96. limit = limit or 0
  97. glimit = _limit[0] or 0
  98. if limit != glimit:
  99. _limit[0] = limit
  100. for pool in _all_pools():
  101. pool.resize(limit)
  102. return limit
  103. def reset(*args, **kwargs):
  104. """Reset all pools by closing open resources."""
  105. for pool in _all_pools():
  106. try:
  107. pool.force_close_all()
  108. except Exception:
  109. pass
  110. for group in _groups:
  111. group.clear()