common.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414
  1. """Common Utilities."""
  2. from __future__ import absolute_import, unicode_literals
  3. import os
  4. import socket
  5. import threading
  6. from collections import deque
  7. from contextlib import contextmanager
  8. from functools import partial
  9. from itertools import count
  10. from uuid import uuid5, uuid4, uuid3, NAMESPACE_OID
  11. from amqp import RecoverableConnectionError
  12. from .entity import Exchange, Queue
  13. from .five import bytes_if_py2, range
  14. from .log import get_logger
  15. from .serialization import registry as serializers
  16. from .utils.uuid import uuid
  17. try:
  18. from _thread import get_ident
  19. except ImportError: # pragma: no cover
  20. try: # noqa
  21. from thread import get_ident # noqa
  22. except ImportError: # pragma: no cover
  23. from dummy_thread import get_ident # noqa
  24. __all__ = ('Broadcast', 'maybe_declare', 'uuid',
  25. 'itermessages', 'send_reply',
  26. 'collect_replies', 'insured', 'drain_consumer',
  27. 'eventloop')
  28. #: Prefetch count can't exceed short.
  29. PREFETCH_COUNT_MAX = 0xFFFF
  30. logger = get_logger(__name__)
  31. _node_id = None
  32. def get_node_id():
  33. global _node_id
  34. if _node_id is None:
  35. _node_id = uuid4().int
  36. return _node_id
  37. def generate_oid(node_id, process_id, thread_id, instance):
  38. ent = bytes_if_py2('%x-%x-%x-%x' % (
  39. node_id, process_id, thread_id, id(instance)))
  40. try:
  41. ret = str(uuid3(NAMESPACE_OID, ent))
  42. except ValueError:
  43. ret = str(uuid5(NAMESPACE_OID, ent))
  44. return ret
  45. def oid_from(instance, threads=True):
  46. return generate_oid(
  47. get_node_id(),
  48. os.getpid(),
  49. get_ident() if threads else 0,
  50. instance,
  51. )
  52. class Broadcast(Queue):
  53. """Broadcast queue.
  54. Convenience class used to define broadcast queues.
  55. Every queue instance will have a unique name,
  56. and both the queue and exchange is configured with auto deletion.
  57. Arguments:
  58. name (str): This is used as the name of the exchange.
  59. queue (str): By default a unique id is used for the queue
  60. name for every consumer. You can specify a custom
  61. queue name here.
  62. **kwargs (Any): See :class:`~kombu.Queue` for a list
  63. of additional keyword arguments supported.
  64. """
  65. attrs = Queue.attrs + (('queue', None),)
  66. def __init__(self, name=None, queue=None, auto_delete=True,
  67. exchange=None, alias=None, **kwargs):
  68. queue = '{0}.{1}'.format(queue or 'bcast', uuid())
  69. return super(Broadcast, self).__init__(
  70. alias=alias or name,
  71. queue=queue,
  72. name=queue,
  73. auto_delete=auto_delete,
  74. exchange=(exchange if exchange is not None
  75. else Exchange(name, type='fanout')),
  76. **kwargs
  77. )
  78. def declaration_cached(entity, channel):
  79. return entity in channel.connection.client.declared_entities
  80. def maybe_declare(entity, channel=None, retry=False, **retry_policy):
  81. """Declare entity (cached)."""
  82. if retry:
  83. return _imaybe_declare(entity, channel, **retry_policy)
  84. return _maybe_declare(entity, channel)
  85. def _maybe_declare(entity, channel):
  86. is_bound = entity.is_bound
  87. orig = entity
  88. if not is_bound:
  89. assert channel
  90. entity = entity.bind(channel)
  91. if channel is None:
  92. assert is_bound
  93. channel = entity.channel
  94. declared = ident = None
  95. if channel.connection and entity.can_cache_declaration:
  96. declared = channel.connection.client.declared_entities
  97. ident = hash(entity)
  98. if ident in declared:
  99. return False
  100. if not channel.connection:
  101. raise RecoverableConnectionError('channel disconnected')
  102. entity.declare(channel=channel)
  103. if declared is not None and ident:
  104. declared.add(ident)
  105. if orig is not None:
  106. orig.name = entity.name
  107. return True
  108. def _imaybe_declare(entity, channel, **retry_policy):
  109. return entity.channel.connection.client.ensure(
  110. entity, _maybe_declare, **retry_policy)(entity, channel)
  111. def drain_consumer(consumer, limit=1, timeout=None, callbacks=None):
  112. """Drain messages from consumer instance."""
  113. acc = deque()
  114. def on_message(body, message):
  115. acc.append((body, message))
  116. consumer.callbacks = [on_message] + (callbacks or [])
  117. with consumer:
  118. for _ in eventloop(consumer.channel.connection.client,
  119. limit=limit, timeout=timeout, ignore_timeouts=True):
  120. try:
  121. yield acc.popleft()
  122. except IndexError:
  123. pass
  124. def itermessages(conn, channel, queue, limit=1, timeout=None,
  125. callbacks=None, **kwargs):
  126. """Iterator over messages."""
  127. return drain_consumer(
  128. conn.Consumer(queues=[queue], channel=channel, **kwargs),
  129. limit=limit, timeout=timeout, callbacks=callbacks,
  130. )
  131. def eventloop(conn, limit=None, timeout=None, ignore_timeouts=False):
  132. """Best practice generator wrapper around ``Connection.drain_events``.
  133. Able to drain events forever, with a limit, and optionally ignoring
  134. timeout errors (a timeout of 1 is often used in environments where
  135. the socket can get "stuck", and is a best practice for Kombu consumers).
  136. ``eventloop`` is a generator.
  137. Examples:
  138. >>> from kombu.common import eventloop
  139. >>> def run(conn):
  140. ... it = eventloop(conn, timeout=1, ignore_timeouts=True)
  141. ... next(it) # one event consumed, or timed out.
  142. ...
  143. ... for _ in eventloop(conn, timeout=1, ignore_timeouts=True):
  144. ... pass # loop forever.
  145. It also takes an optional limit parameter, and timeout errors
  146. are propagated by default::
  147. for _ in eventloop(connection, limit=1, timeout=1):
  148. pass
  149. See Also:
  150. :func:`itermessages`, which is an event loop bound to one or more
  151. consumers, that yields any messages received.
  152. """
  153. for i in limit and range(limit) or count():
  154. try:
  155. yield conn.drain_events(timeout=timeout)
  156. except socket.timeout:
  157. if timeout and not ignore_timeouts: # pragma: no cover
  158. raise
  159. def send_reply(exchange, req, msg,
  160. producer=None, retry=False, retry_policy=None, **props):
  161. """Send reply for request.
  162. Arguments:
  163. exchange (kombu.Exchange, str): Reply exchange
  164. req (~kombu.Message): Original request, a message with
  165. a ``reply_to`` property.
  166. producer (kombu.Producer): Producer instance
  167. retry (bool): If true must retry according to
  168. the ``reply_policy`` argument.
  169. retry_policy (Dict): Retry settings.
  170. **props (Any): Extra properties.
  171. """
  172. return producer.publish(
  173. msg, exchange=exchange,
  174. retry=retry, retry_policy=retry_policy,
  175. **dict({'routing_key': req.properties['reply_to'],
  176. 'correlation_id': req.properties.get('correlation_id'),
  177. 'serializer': serializers.type_to_name[req.content_type],
  178. 'content_encoding': req.content_encoding}, **props)
  179. )
  180. def collect_replies(conn, channel, queue, *args, **kwargs):
  181. """Generator collecting replies from ``queue``."""
  182. no_ack = kwargs.setdefault('no_ack', True)
  183. received = False
  184. try:
  185. for body, message in itermessages(conn, channel, queue,
  186. *args, **kwargs):
  187. if not no_ack:
  188. message.ack()
  189. received = True
  190. yield body
  191. finally:
  192. if received:
  193. channel.after_reply_message_received(queue.name)
  194. def _ensure_errback(exc, interval):
  195. logger.error(
  196. 'Connection error: %r. Retry in %ss\n', exc, interval,
  197. exc_info=True,
  198. )
  199. @contextmanager
  200. def _ignore_errors(conn):
  201. try:
  202. yield
  203. except conn.connection_errors + conn.channel_errors:
  204. pass
  205. def ignore_errors(conn, fun=None, *args, **kwargs):
  206. """Ignore connection and channel errors.
  207. The first argument must be a connection object, or any other object
  208. with ``connection_error`` and ``channel_error`` attributes.
  209. Can be used as a function:
  210. .. code-block:: python
  211. def example(connection):
  212. ignore_errors(connection, consumer.channel.close)
  213. or as a context manager:
  214. .. code-block:: python
  215. def example(connection):
  216. with ignore_errors(connection):
  217. consumer.channel.close()
  218. Note:
  219. Connection and channel errors should be properly handled,
  220. and not ignored. Using this function is only acceptable in a cleanup
  221. phase, like when a connection is lost or at shutdown.
  222. """
  223. if fun:
  224. with _ignore_errors(conn):
  225. return fun(*args, **kwargs)
  226. return _ignore_errors(conn)
  227. def revive_connection(connection, channel, on_revive=None):
  228. if on_revive:
  229. on_revive(channel)
  230. def insured(pool, fun, args, kwargs, errback=None, on_revive=None, **opts):
  231. """Function wrapper to handle connection errors.
  232. Ensures function performing broker commands completes
  233. despite intermittent connection failures.
  234. """
  235. errback = errback or _ensure_errback
  236. with pool.acquire(block=True) as conn:
  237. conn.ensure_connection(errback=errback)
  238. # we cache the channel for subsequent calls, this has to be
  239. # reset on revival.
  240. channel = conn.default_channel
  241. revive = partial(revive_connection, conn, on_revive=on_revive)
  242. insured = conn.autoretry(fun, channel, errback=errback,
  243. on_revive=revive, **opts)
  244. retval, _ = insured(*args, **dict(kwargs, connection=conn))
  245. return retval
  246. class QoS(object):
  247. """Thread safe increment/decrement of a channels prefetch_count.
  248. Arguments:
  249. callback (Callable): Function used to set new prefetch count,
  250. e.g. ``consumer.qos`` or ``channel.basic_qos``. Will be called
  251. with a single ``prefetch_count`` keyword argument.
  252. initial_value (int): Initial prefetch count value..
  253. Example:
  254. >>> from kombu import Consumer, Connection
  255. >>> connection = Connection('amqp://')
  256. >>> consumer = Consumer(connection)
  257. >>> qos = QoS(consumer.qos, initial_prefetch_count=2)
  258. >>> qos.update() # set initial
  259. >>> qos.value
  260. 2
  261. >>> def in_some_thread():
  262. ... qos.increment_eventually()
  263. >>> def in_some_other_thread():
  264. ... qos.decrement_eventually()
  265. >>> while 1:
  266. ... if qos.prev != qos.value:
  267. ... qos.update() # prefetch changed so update.
  268. It can be used with any function supporting a ``prefetch_count`` keyword
  269. argument::
  270. >>> channel = connection.channel()
  271. >>> QoS(channel.basic_qos, 10)
  272. >>> def set_qos(prefetch_count):
  273. ... print('prefetch count now: %r' % (prefetch_count,))
  274. >>> QoS(set_qos, 10)
  275. """
  276. prev = None
  277. def __init__(self, callback, initial_value):
  278. self.callback = callback
  279. self._mutex = threading.RLock()
  280. self.value = initial_value or 0
  281. def increment_eventually(self, n=1):
  282. """Increment the value, but do not update the channels QoS.
  283. Note:
  284. The MainThread will be responsible for calling :meth:`update`
  285. when necessary.
  286. """
  287. with self._mutex:
  288. if self.value:
  289. self.value = self.value + max(n, 0)
  290. return self.value
  291. def decrement_eventually(self, n=1):
  292. """Decrement the value, but do not update the channels QoS.
  293. Note:
  294. The MainThread will be responsible for calling :meth:`update`
  295. when necessary.
  296. """
  297. with self._mutex:
  298. if self.value:
  299. self.value -= n
  300. if self.value < 1:
  301. self.value = 1
  302. return self.value
  303. def set(self, pcount):
  304. """Set channel prefetch_count setting."""
  305. if pcount != self.prev:
  306. new_value = pcount
  307. if pcount > PREFETCH_COUNT_MAX:
  308. logger.warning('QoS: Disabled: prefetch_count exceeds %r',
  309. PREFETCH_COUNT_MAX)
  310. new_value = 0
  311. logger.debug('basic.qos: prefetch_count->%s', new_value)
  312. self.callback(prefetch_count=new_value)
  313. self.prev = pcount
  314. return pcount
  315. def update(self):
  316. """Update prefetch count with current value."""
  317. with self._mutex:
  318. return self.set(self.value)