mixins.py 9.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298
  1. # -*- coding: utf-8 -*-
  2. """Mixins."""
  3. from __future__ import absolute_import, unicode_literals
  4. import socket
  5. from contextlib import contextmanager
  6. from functools import partial
  7. from itertools import count
  8. from time import sleep
  9. from .common import ignore_errors
  10. from .five import range
  11. from .messaging import Consumer, Producer
  12. from .log import get_logger
  13. from .utils.compat import nested
  14. from .utils.encoding import safe_repr
  15. from .utils.limits import TokenBucket
  16. from .utils.objects import cached_property
  17. __all__ = ('ConsumerMixin', 'ConsumerProducerMixin')
  18. logger = get_logger(__name__)
  19. debug, info, warn, error = logger.debug, logger.info, logger.warn, logger.error
  20. W_CONN_LOST = """\
  21. Connection to broker lost, trying to re-establish connection...\
  22. """
  23. W_CONN_ERROR = """\
  24. Broker connection error, trying again in %s seconds: %r.\
  25. """
  26. class ConsumerMixin(object):
  27. """Convenience mixin for implementing consumer programs.
  28. It can be used outside of threads, with threads, or greenthreads
  29. (eventlet/gevent) too.
  30. The basic class would need a :attr:`connection` attribute
  31. which must be a :class:`~kombu.Connection` instance,
  32. and define a :meth:`get_consumers` method that returns a list
  33. of :class:`kombu.Consumer` instances to use.
  34. Supporting multiple consumers is important so that multiple
  35. channels can be used for different QoS requirements.
  36. Example:
  37. .. code-block:: python
  38. class Worker(ConsumerMixin):
  39. task_queue = Queue('tasks', Exchange('tasks'), 'tasks')
  40. def __init__(self, connection):
  41. self.connection = None
  42. def get_consumers(self, Consumer, channel):
  43. return [Consumer(queues=[self.task_queue],
  44. callbacks=[self.on_task])]
  45. def on_task(self, body, message):
  46. print('Got task: {0!r}'.format(body))
  47. message.ack()
  48. Methods:
  49. * :meth:`extra_context`
  50. Optional extra context manager that will be entered
  51. after the connection and consumers have been set up.
  52. Takes arguments ``(connection, channel)``.
  53. * :meth:`on_connection_error`
  54. Handler called if the connection is lost/ or
  55. is unavailable.
  56. Takes arguments ``(exc, interval)``, where interval
  57. is the time in seconds when the connection will be retried.
  58. The default handler will log the exception.
  59. * :meth:`on_connection_revived`
  60. Handler called as soon as the connection is re-established
  61. after connection failure.
  62. Takes no arguments.
  63. * :meth:`on_consume_ready`
  64. Handler called when the consumer is ready to accept
  65. messages.
  66. Takes arguments ``(connection, channel, consumers)``.
  67. Also keyword arguments to ``consume`` are forwarded
  68. to this handler.
  69. * :meth:`on_consume_end`
  70. Handler called after the consumers are canceled.
  71. Takes arguments ``(connection, channel)``.
  72. * :meth:`on_iteration`
  73. Handler called for every iteration while draining
  74. events.
  75. Takes no arguments.
  76. * :meth:`on_decode_error`
  77. Handler called if a consumer was unable to decode
  78. the body of a message.
  79. Takes arguments ``(message, exc)`` where message is the
  80. original message object.
  81. The default handler will log the error and
  82. acknowledge the message, so if you override make
  83. sure to call super, or perform these steps yourself.
  84. """
  85. #: maximum number of retries trying to re-establish the connection,
  86. #: if the connection is lost/unavailable.
  87. connect_max_retries = None
  88. #: When this is set to true the consumer should stop consuming
  89. #: and return, so that it can be joined if it is the implementation
  90. #: of a thread.
  91. should_stop = False
  92. def get_consumers(self, Consumer, channel):
  93. raise NotImplementedError('Subclass responsibility')
  94. def on_connection_revived(self):
  95. pass
  96. def on_consume_ready(self, connection, channel, consumers, **kwargs):
  97. pass
  98. def on_consume_end(self, connection, channel):
  99. pass
  100. def on_iteration(self):
  101. pass
  102. def on_decode_error(self, message, exc):
  103. error("Can't decode message body: %r (type:%r encoding:%r raw:%r')",
  104. exc, message.content_type, message.content_encoding,
  105. safe_repr(message.body))
  106. message.ack()
  107. def on_connection_error(self, exc, interval):
  108. warn(W_CONN_ERROR, interval, exc, exc_info=1)
  109. @contextmanager
  110. def extra_context(self, connection, channel):
  111. yield
  112. def run(self, _tokens=1, **kwargs):
  113. restart_limit = self.restart_limit
  114. errors = (self.connection.connection_errors +
  115. self.connection.channel_errors)
  116. while not self.should_stop:
  117. try:
  118. if restart_limit.can_consume(_tokens): # pragma: no cover
  119. for _ in self.consume(limit=None, **kwargs):
  120. pass
  121. else:
  122. sleep(restart_limit.expected_time(_tokens))
  123. except errors:
  124. warn(W_CONN_LOST, exc_info=1)
  125. @contextmanager
  126. def consumer_context(self, **kwargs):
  127. with self.Consumer() as (connection, channel, consumers):
  128. with self.extra_context(connection, channel):
  129. self.on_consume_ready(connection, channel, consumers, **kwargs)
  130. yield connection, channel, consumers
  131. def consume(self, limit=None, timeout=None, safety_interval=1, **kwargs):
  132. elapsed = 0
  133. with self.consumer_context(**kwargs) as (conn, channel, consumers):
  134. for i in limit and range(limit) or count():
  135. if self.should_stop:
  136. break
  137. self.on_iteration()
  138. try:
  139. conn.drain_events(timeout=safety_interval)
  140. except socket.timeout:
  141. conn.heartbeat_check()
  142. elapsed += safety_interval
  143. if timeout and elapsed >= timeout:
  144. raise
  145. except socket.error:
  146. if not self.should_stop:
  147. raise
  148. else:
  149. yield
  150. elapsed = 0
  151. debug('consume exiting')
  152. def maybe_conn_error(self, fun):
  153. """Use :func:`kombu.common.ignore_errors` instead."""
  154. return ignore_errors(self, fun)
  155. def create_connection(self):
  156. return self.connection.clone()
  157. @contextmanager
  158. def establish_connection(self):
  159. with self.create_connection() as conn:
  160. conn.ensure_connection(self.on_connection_error,
  161. self.connect_max_retries)
  162. yield conn
  163. @contextmanager
  164. def Consumer(self):
  165. with self.establish_connection() as conn:
  166. self.on_connection_revived()
  167. info('Connected to %s', conn.as_uri())
  168. channel = conn.default_channel
  169. cls = partial(Consumer, channel,
  170. on_decode_error=self.on_decode_error)
  171. with self._consume_from(*self.get_consumers(cls, channel)) as c:
  172. yield conn, channel, c
  173. debug('Consumers canceled')
  174. self.on_consume_end(conn, channel)
  175. debug('Connection closed')
  176. def _consume_from(self, *consumers):
  177. return nested(*consumers)
  178. @cached_property
  179. def restart_limit(self):
  180. return TokenBucket(1)
  181. @cached_property
  182. def connection_errors(self):
  183. return self.connection.connection_errors
  184. @cached_property
  185. def channel_errors(self):
  186. return self.connection.channel_errors
  187. class ConsumerProducerMixin(ConsumerMixin):
  188. """Consumer and Producer mixin.
  189. Version of ConsumerMixin having separate connection for also
  190. publishing messages.
  191. Example:
  192. .. code-block:: python
  193. class Worker(ConsumerProducerMixin):
  194. def __init__(self, connection):
  195. self.connection = connection
  196. def get_consumers(self, Consumer, channel):
  197. return [Consumer(queues=Queue('foo'),
  198. on_message=self.handle_message,
  199. accept='application/json',
  200. prefetch_count=10)]
  201. def handle_message(self, message):
  202. self.producer.publish(
  203. {'message': 'hello to you'},
  204. exchange='',
  205. routing_key=message.properties['reply_to'],
  206. correlation_id=message.properties['correlation_id'],
  207. retry=True,
  208. )
  209. """
  210. _producer_connection = None
  211. def on_consume_end(self, connection, channel):
  212. if self._producer_connection is not None:
  213. self._producer_connection.close()
  214. self._producer_connection = None
  215. @property
  216. def producer(self):
  217. return Producer(self.producer_connection)
  218. @property
  219. def producer_connection(self):
  220. if self._producer_connection is None:
  221. conn = self.connection.clone()
  222. conn.ensure_connection(self.on_connection_error,
  223. self.connect_max_retries)
  224. self._producer_connection = conn
  225. return self._producer_connection