simple.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. """Simple messaging interface."""
  2. from __future__ import absolute_import, unicode_literals
  3. import socket
  4. from collections import deque
  5. from . import entity
  6. from . import messaging
  7. from .connection import maybe_channel
  8. from .five import Empty, monotonic
  9. __all__ = ('SimpleQueue', 'SimpleBuffer')
  10. class SimpleBase(object):
  11. Empty = Empty
  12. _consuming = False
  13. def __enter__(self):
  14. return self
  15. def __exit__(self, *exc_info):
  16. self.close()
  17. def __init__(self, channel, producer, consumer, no_ack=False):
  18. self.channel = maybe_channel(channel)
  19. self.producer = producer
  20. self.consumer = consumer
  21. self.no_ack = no_ack
  22. self.queue = self.consumer.queues[0]
  23. self.buffer = deque()
  24. self.consumer.register_callback(self._receive)
  25. def get(self, block=True, timeout=None):
  26. if not block:
  27. return self.get_nowait()
  28. self._consume()
  29. time_start = monotonic()
  30. remaining = timeout
  31. while True:
  32. if self.buffer:
  33. return self.buffer.popleft()
  34. if remaining is not None and remaining <= 0.0:
  35. raise self.Empty()
  36. try:
  37. # The `drain_events` method will
  38. # block on the socket connection to rabbitmq. if any
  39. # application-level messages are received, it will put them
  40. # into `self.buffer`.
  41. # * The method will block for UP TO `timeout` milliseconds.
  42. # * The method may raise a socket.timeout exception; or...
  43. # * The method may return without having put anything on
  44. # `self.buffer`. This is because internal heartbeat
  45. # messages are sent over the same socket; also POSIX makes
  46. # no guarantees against socket calls returning early.
  47. self.channel.connection.client.drain_events(timeout=remaining)
  48. except socket.timeout:
  49. raise self.Empty()
  50. if remaining is not None:
  51. elapsed = monotonic() - time_start
  52. remaining = timeout - elapsed
  53. def get_nowait(self):
  54. m = self.queue.get(no_ack=self.no_ack)
  55. if not m:
  56. raise self.Empty()
  57. return m
  58. def put(self, message, serializer=None, headers=None, compression=None,
  59. routing_key=None, **kwargs):
  60. self.producer.publish(message,
  61. serializer=serializer,
  62. routing_key=routing_key,
  63. headers=headers,
  64. compression=compression,
  65. **kwargs)
  66. def clear(self):
  67. return self.consumer.purge()
  68. def qsize(self):
  69. _, size, _ = self.queue.queue_declare(passive=True)
  70. return size
  71. def close(self):
  72. self.consumer.cancel()
  73. def _receive(self, message_data, message):
  74. self.buffer.append(message)
  75. def _consume(self):
  76. if not self._consuming:
  77. self.consumer.consume(no_ack=self.no_ack)
  78. self._consuming = True
  79. def __len__(self):
  80. """`len(self) -> self.qsize()`."""
  81. return self.qsize()
  82. def __bool__(self):
  83. return True
  84. __nonzero__ = __bool__
  85. class SimpleQueue(SimpleBase):
  86. """Simple API for persistent queues."""
  87. no_ack = False
  88. queue_opts = {}
  89. queue_args = {}
  90. exchange_opts = {'type': 'direct'}
  91. def __init__(self, channel, name, no_ack=None, queue_opts=None,
  92. queue_args=None, exchange_opts=None, serializer=None,
  93. compression=None, **kwargs):
  94. queue = name
  95. queue_opts = dict(self.queue_opts, **queue_opts or {})
  96. queue_args = dict(self.queue_args, **queue_args or {})
  97. exchange_opts = dict(self.exchange_opts, **exchange_opts or {})
  98. if no_ack is None:
  99. no_ack = self.no_ack
  100. if not isinstance(queue, entity.Queue):
  101. exchange = entity.Exchange(name, **exchange_opts)
  102. queue = entity.Queue(name, exchange, name,
  103. queue_arguments=queue_args,
  104. **queue_opts)
  105. routing_key = name
  106. else:
  107. exchange = queue.exchange
  108. routing_key = queue.routing_key
  109. consumer = messaging.Consumer(channel, queue)
  110. producer = messaging.Producer(channel, exchange,
  111. serializer=serializer,
  112. routing_key=routing_key,
  113. compression=compression)
  114. super(SimpleQueue, self).__init__(channel, producer,
  115. consumer, no_ack, **kwargs)
  116. class SimpleBuffer(SimpleQueue):
  117. """Simple API for ephemeral queues."""
  118. no_ack = True
  119. queue_opts = {'durable': False,
  120. 'auto_delete': True}
  121. exchange_opts = {'durable': False,
  122. 'delivery_mode': 'transient',
  123. 'auto_delete': True}