Development of an internal social media platform with personalised dashboards for students
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

pidbox.py 12KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364
  1. """
  2. kombu.pidbox
  3. ===============
  4. Generic process mailbox.
  5. """
  6. from __future__ import absolute_import
  7. import socket
  8. import warnings
  9. from collections import defaultdict, deque
  10. from copy import copy
  11. from itertools import count
  12. from threading import local
  13. from time import time
  14. from . import Exchange, Queue, Consumer, Producer
  15. from .clocks import LamportClock
  16. from .common import maybe_declare, oid_from
  17. from .exceptions import InconsistencyError
  18. from .five import range
  19. from .log import get_logger
  20. from .utils import cached_property, kwdict, uuid, reprcall
  21. REPLY_QUEUE_EXPIRES = 10
  22. W_PIDBOX_IN_USE = """\
  23. A node named {node.hostname} is already using this process mailbox!
  24. Maybe you forgot to shutdown the other node or did not do so properly?
  25. Or if you meant to start multiple nodes on the same host please make sure
  26. you give each node a unique node name!
  27. """
  28. __all__ = ['Node', 'Mailbox']
  29. logger = get_logger(__name__)
  30. debug, error = logger.debug, logger.error
  31. class Node(object):
  32. #: hostname of the node.
  33. hostname = None
  34. #: the :class:`Mailbox` this is a node for.
  35. mailbox = None
  36. #: map of method name/handlers.
  37. handlers = None
  38. #: current context (passed on to handlers)
  39. state = None
  40. #: current channel.
  41. channel = None
  42. def __init__(self, hostname, state=None, channel=None,
  43. handlers=None, mailbox=None):
  44. self.channel = channel
  45. self.mailbox = mailbox
  46. self.hostname = hostname
  47. self.state = state
  48. self.adjust_clock = self.mailbox.clock.adjust
  49. if handlers is None:
  50. handlers = {}
  51. self.handlers = handlers
  52. def Consumer(self, channel=None, no_ack=True, accept=None, **options):
  53. queue = self.mailbox.get_queue(self.hostname)
  54. def verify_exclusive(name, messages, consumers):
  55. if consumers:
  56. warnings.warn(W_PIDBOX_IN_USE.format(node=self))
  57. queue.on_declared = verify_exclusive
  58. return Consumer(
  59. channel or self.channel, [queue], no_ack=no_ack,
  60. accept=self.mailbox.accept if accept is None else accept,
  61. **options
  62. )
  63. def handler(self, fun):
  64. self.handlers[fun.__name__] = fun
  65. return fun
  66. def on_decode_error(self, message, exc):
  67. error('Cannot decode message: %r', exc, exc_info=1)
  68. def listen(self, channel=None, callback=None):
  69. consumer = self.Consumer(channel=channel,
  70. callbacks=[callback or self.handle_message],
  71. on_decode_error=self.on_decode_error)
  72. consumer.consume()
  73. return consumer
  74. def dispatch(self, method, arguments=None,
  75. reply_to=None, ticket=None, **kwargs):
  76. arguments = arguments or {}
  77. debug('pidbox received method %s [reply_to:%s ticket:%s]',
  78. reprcall(method, (), kwargs=arguments), reply_to, ticket)
  79. handle = reply_to and self.handle_call or self.handle_cast
  80. try:
  81. reply = handle(method, kwdict(arguments))
  82. except SystemExit:
  83. raise
  84. except Exception as exc:
  85. error('pidbox command error: %r', exc, exc_info=1)
  86. reply = {'error': repr(exc)}
  87. if reply_to:
  88. self.reply({self.hostname: reply},
  89. exchange=reply_to['exchange'],
  90. routing_key=reply_to['routing_key'],
  91. ticket=ticket)
  92. return reply
  93. def handle(self, method, arguments={}):
  94. return self.handlers[method](self.state, **arguments)
  95. def handle_call(self, method, arguments):
  96. return self.handle(method, arguments)
  97. def handle_cast(self, method, arguments):
  98. return self.handle(method, arguments)
  99. def handle_message(self, body, message=None):
  100. destination = body.get('destination')
  101. if message:
  102. self.adjust_clock(message.headers.get('clock') or 0)
  103. if not destination or self.hostname in destination:
  104. return self.dispatch(**kwdict(body))
  105. dispatch_from_message = handle_message
  106. def reply(self, data, exchange, routing_key, ticket, **kwargs):
  107. self.mailbox._publish_reply(data, exchange, routing_key, ticket,
  108. channel=self.channel,
  109. serializer=self.mailbox.serializer)
  110. class Mailbox(object):
  111. node_cls = Node
  112. exchange_fmt = '%s.pidbox'
  113. reply_exchange_fmt = 'reply.%s.pidbox'
  114. #: Name of application.
  115. namespace = None
  116. #: Connection (if bound).
  117. connection = None
  118. #: Exchange type (usually direct, or fanout for broadcast).
  119. type = 'direct'
  120. #: mailbox exchange (init by constructor).
  121. exchange = None
  122. #: exchange to send replies to.
  123. reply_exchange = None
  124. #: Only accepts json messages by default.
  125. accept = ['json']
  126. #: Message serializer
  127. serializer = None
  128. def __init__(self, namespace,
  129. type='direct', connection=None, clock=None,
  130. accept=None, serializer=None):
  131. self.namespace = namespace
  132. self.connection = connection
  133. self.type = type
  134. self.clock = LamportClock() if clock is None else clock
  135. self.exchange = self._get_exchange(self.namespace, self.type)
  136. self.reply_exchange = self._get_reply_exchange(self.namespace)
  137. self._tls = local()
  138. self.unclaimed = defaultdict(deque)
  139. self.accept = self.accept if accept is None else accept
  140. self.serializer = self.serializer if serializer is None else serializer
  141. def __call__(self, connection):
  142. bound = copy(self)
  143. bound.connection = connection
  144. return bound
  145. def Node(self, hostname=None, state=None, channel=None, handlers=None):
  146. hostname = hostname or socket.gethostname()
  147. return self.node_cls(hostname, state, channel, handlers, mailbox=self)
  148. def call(self, destination, command, kwargs={},
  149. timeout=None, callback=None, channel=None):
  150. return self._broadcast(command, kwargs, destination,
  151. reply=True, timeout=timeout,
  152. callback=callback,
  153. channel=channel)
  154. def cast(self, destination, command, kwargs={}):
  155. return self._broadcast(command, kwargs, destination, reply=False)
  156. def abcast(self, command, kwargs={}):
  157. return self._broadcast(command, kwargs, reply=False)
  158. def multi_call(self, command, kwargs={}, timeout=1,
  159. limit=None, callback=None, channel=None):
  160. return self._broadcast(command, kwargs, reply=True,
  161. timeout=timeout, limit=limit,
  162. callback=callback,
  163. channel=channel)
  164. def get_reply_queue(self):
  165. oid = self.oid
  166. return Queue(
  167. '%s.%s' % (oid, self.reply_exchange.name),
  168. exchange=self.reply_exchange,
  169. routing_key=oid,
  170. durable=False,
  171. auto_delete=True,
  172. queue_arguments={'x-expires': int(REPLY_QUEUE_EXPIRES * 1000)},
  173. )
  174. @cached_property
  175. def reply_queue(self):
  176. return self.get_reply_queue()
  177. def get_queue(self, hostname):
  178. return Queue('%s.%s.pidbox' % (hostname, self.namespace),
  179. exchange=self.exchange,
  180. durable=False,
  181. auto_delete=True)
  182. def _publish_reply(self, reply, exchange, routing_key, ticket,
  183. channel=None, **opts):
  184. chan = channel or self.connection.default_channel
  185. exchange = Exchange(exchange, exchange_type='direct',
  186. delivery_mode='transient',
  187. durable=False)
  188. producer = Producer(chan, auto_declare=False)
  189. try:
  190. producer.publish(
  191. reply, exchange=exchange, routing_key=routing_key,
  192. declare=[exchange], headers={
  193. 'ticket': ticket, 'clock': self.clock.forward(),
  194. },
  195. **opts
  196. )
  197. except InconsistencyError:
  198. pass # queue probably deleted and no one is expecting a reply.
  199. def _publish(self, type, arguments, destination=None,
  200. reply_ticket=None, channel=None, timeout=None,
  201. serializer=None):
  202. message = {'method': type,
  203. 'arguments': arguments,
  204. 'destination': destination}
  205. chan = channel or self.connection.default_channel
  206. exchange = self.exchange
  207. if reply_ticket:
  208. maybe_declare(self.reply_queue(channel))
  209. message.update(ticket=reply_ticket,
  210. reply_to={'exchange': self.reply_exchange.name,
  211. 'routing_key': self.oid})
  212. serializer = serializer or self.serializer
  213. producer = Producer(chan, auto_declare=False)
  214. producer.publish(
  215. message, exchange=exchange.name, declare=[exchange],
  216. headers={'clock': self.clock.forward(),
  217. 'expires': time() + timeout if timeout else 0},
  218. serializer=serializer,
  219. )
  220. def _broadcast(self, command, arguments=None, destination=None,
  221. reply=False, timeout=1, limit=None,
  222. callback=None, channel=None, serializer=None):
  223. if destination is not None and \
  224. not isinstance(destination, (list, tuple)):
  225. raise ValueError(
  226. 'destination must be a list/tuple not {0}'.format(
  227. type(destination)))
  228. arguments = arguments or {}
  229. reply_ticket = reply and uuid() or None
  230. chan = channel or self.connection.default_channel
  231. # Set reply limit to number of destinations (if specified)
  232. if limit is None and destination:
  233. limit = destination and len(destination) or None
  234. serializer = serializer or self.serializer
  235. self._publish(command, arguments, destination=destination,
  236. reply_ticket=reply_ticket,
  237. channel=chan,
  238. timeout=timeout,
  239. serializer=serializer)
  240. if reply_ticket:
  241. return self._collect(reply_ticket, limit=limit,
  242. timeout=timeout,
  243. callback=callback,
  244. channel=chan)
  245. def _collect(self, ticket,
  246. limit=None, timeout=1, callback=None,
  247. channel=None, accept=None):
  248. if accept is None:
  249. accept = self.accept
  250. chan = channel or self.connection.default_channel
  251. queue = self.reply_queue
  252. consumer = Consumer(channel, [queue], accept=accept, no_ack=True)
  253. responses = []
  254. unclaimed = self.unclaimed
  255. adjust_clock = self.clock.adjust
  256. try:
  257. return unclaimed.pop(ticket)
  258. except KeyError:
  259. pass
  260. def on_message(body, message):
  261. # ticket header added in kombu 2.5
  262. header = message.headers.get
  263. adjust_clock(header('clock') or 0)
  264. expires = header('expires')
  265. if expires and time() > expires:
  266. return
  267. this_id = header('ticket', ticket)
  268. if this_id == ticket:
  269. if callback:
  270. callback(body)
  271. responses.append(body)
  272. else:
  273. unclaimed[this_id].append(body)
  274. consumer.register_callback(on_message)
  275. try:
  276. with consumer:
  277. for i in limit and range(limit) or count():
  278. try:
  279. self.connection.drain_events(timeout=timeout)
  280. except socket.timeout:
  281. break
  282. return responses
  283. finally:
  284. chan.after_reply_message_received(queue.name)
  285. def _get_exchange(self, namespace, type):
  286. return Exchange(self.exchange_fmt % namespace,
  287. type=type,
  288. durable=False,
  289. delivery_mode='transient')
  290. def _get_reply_exchange(self, namespace):
  291. return Exchange(self.reply_exchange_fmt % namespace,
  292. type='direct',
  293. durable=False,
  294. delivery_mode='transient')
  295. @cached_property
  296. def oid(self):
  297. try:
  298. return self._tls.OID
  299. except AttributeError:
  300. oid = self._tls.OID = oid_from(self)
  301. return oid