consumer.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685
  1. # -*- coding: utf-8 -*-
  2. """
  3. celery.worker.consumer
  4. ~~~~~~~~~~~~~~~~~~~~~~
  5. This module contains the component responsible for consuming messages
  6. from the broker, processing the messages and keeping the broker connections
  7. up and running.
  8. :copyright: (c) 2009 - 2012 by Ask Solem.
  9. :license: BSD, see LICENSE for more details.
  10. * :meth:`~Consumer.start` is an infinite loop, which only iterates
  11. again if the connection is lost. For each iteration (at start, or if the
  12. connection is lost) it calls :meth:`~Consumer.reset_connection`,
  13. and starts the consumer by calling :meth:`~Consumer.consume_messages`.
  14. * :meth:`~Consumer.reset_connection`, clears the internal queues,
  15. establishes a new connection to the broker, sets up the task
  16. consumer (+ QoS), and the broadcast remote control command consumer.
  17. Also if events are enabled it configures the event dispatcher and starts
  18. up the heartbeat thread.
  19. * Finally it can consume messages. :meth:`~Consumer.consume_messages`
  20. is simply an infinite loop waiting for events on the AMQP channels.
  21. Both the task consumer and the broadcast consumer uses the same
  22. callback: :meth:`~Consumer.receive_message`.
  23. * So for each message received the :meth:`~Consumer.receive_message`
  24. method is called, this checks the payload of the message for either
  25. a `task` key or a `control` key.
  26. If the message is a task, it verifies the validity of the message
  27. converts it to a :class:`celery.worker.job.Request`, and sends
  28. it to :meth:`~Consumer.on_task`.
  29. If the message is a control command the message is passed to
  30. :meth:`~Consumer.on_control`, which in turn dispatches
  31. the control command using the control dispatcher.
  32. It also tries to handle malformed or invalid messages properly,
  33. so the worker doesn't choke on them and die. Any invalid messages
  34. are acknowledged immediately and logged, so the message is not resent
  35. again, and again.
  36. * If the task has an ETA/countdown, the task is moved to the `eta_schedule`
  37. so the :class:`timer2.Timer` can schedule it at its
  38. deadline. Tasks without an eta are moved immediately to the `ready_queue`,
  39. so they can be picked up by the :class:`~celery.worker.mediator.Mediator`
  40. to be sent to the pool.
  41. * When a task with an ETA is received the QoS prefetch count is also
  42. incremented, so another message can be reserved. When the ETA is met
  43. the prefetch count is decremented again, though this cannot happen
  44. immediately because amqplib doesn't support doing broker requests
  45. across threads. Instead the current prefetch count is kept as a
  46. shared counter, so as soon as :meth:`~Consumer.consume_messages`
  47. detects that the value has changed it will send out the actual
  48. QoS event to the broker.
  49. * Notice that when the connection is lost all internal queues are cleared
  50. because we can no longer ack the messages reserved in memory.
  51. However, this is not dangerous as the broker will resend them
  52. to another worker when the channel is closed.
  53. * **WARNING**: :meth:`~Consumer.stop` does not close the connection!
  54. This is because some pre-acked messages may be in processing,
  55. and they need to be finished before the channel is closed.
  56. For celeryd this means the pool must finish the tasks it has acked
  57. early, *then* close the connection.
  58. """
  59. from __future__ import absolute_import
  60. from __future__ import with_statement
  61. import logging
  62. import socket
  63. import threading
  64. from kombu.utils.encoding import safe_repr
  65. from celery.app import app_or_default
  66. from celery.datastructures import AttributeDict
  67. from celery.exceptions import InvalidTaskError
  68. from celery.utils import timer2
  69. from celery.utils.functional import noop
  70. from celery.utils.log import get_logger
  71. from . import state
  72. from .abstract import StartStopComponent
  73. from .control import Panel
  74. from .heartbeat import Heart
  75. RUN = 0x1
  76. CLOSE = 0x2
  77. #: Prefetch count can't exceed short.
  78. PREFETCH_COUNT_MAX = 0xFFFF
  79. UNKNOWN_FORMAT = """\
  80. Received and deleted unknown message. Wrong destination?!?
  81. The full contents of the message body was: %s
  82. """
  83. #: Error message for when an unregistered task is received.
  84. UNKNOWN_TASK_ERROR = """\
  85. Received unregistered task of type %s.
  86. The message has been ignored and discarded.
  87. Did you remember to import the module containing this task?
  88. Or maybe you are using relative imports?
  89. Please see http://bit.ly/gLye1c for more information.
  90. The full contents of the message body was:
  91. %s
  92. """
  93. #: Error message for when an invalid task message is received.
  94. INVALID_TASK_ERROR = """\
  95. Received invalid task message: %s
  96. The message has been ignored and discarded.
  97. Please ensure your message conforms to the task
  98. message protocol as described here: http://bit.ly/hYj41y
  99. The full contents of the message body was:
  100. %s
  101. """
  102. MESSAGE_REPORT_FMT = """\
  103. body: %s {content_type:%s content_encoding:%s delivery_info:%s}\
  104. """
  105. RETRY_CONNECTION = """\
  106. Consumer: Connection to broker lost. \
  107. Trying to re-establish the connection...\
  108. """
  109. logger = get_logger(__name__)
  110. info, warn, error, crit = (logger.info, logger.warn,
  111. logger.error, logger.critical)
  112. def debug(msg, *args, **kwargs):
  113. logger.debug("Consumer: %s" % (msg, ), *args, **kwargs)
  114. class Component(StartStopComponent):
  115. name = "worker.consumer"
  116. last = True
  117. def create(self, w):
  118. prefetch_count = w.concurrency * w.prefetch_multiplier
  119. c = w.consumer = self.instantiate(
  120. w.consumer_cls, w.ready_queue, w.scheduler,
  121. hostname=w.hostname,
  122. send_events=w.send_events,
  123. init_callback=w.ready_callback,
  124. initial_prefetch_count=prefetch_count,
  125. pool=w.pool,
  126. priority_timer=w.priority_timer,
  127. app=w.app,
  128. controller=w)
  129. return c
  130. class QoS(object):
  131. """Quality of Service for Channel.
  132. For thread-safe increment/decrement of a channels prefetch count value.
  133. :param consumer: A :class:`kombu.messaging.Consumer` instance.
  134. :param initial_value: Initial prefetch count value.
  135. """
  136. prev = None
  137. def __init__(self, consumer, initial_value):
  138. self.consumer = consumer
  139. self._mutex = threading.RLock()
  140. self.value = initial_value
  141. def increment(self, n=1):
  142. """Increment the current prefetch count value by n."""
  143. with self._mutex:
  144. if self.value:
  145. new_value = self.value + max(n, 0)
  146. self.value = self.set(new_value)
  147. return self.value
  148. def _sub(self, n=1):
  149. assert self.value - n > 1
  150. self.value -= n
  151. def decrement(self, n=1):
  152. """Decrement the current prefetch count value by n."""
  153. with self._mutex:
  154. if self.value:
  155. self._sub(n)
  156. self.set(self.value)
  157. return self.value
  158. def decrement_eventually(self, n=1):
  159. """Decrement the value, but do not update the qos.
  160. The MainThread will be responsible for calling :meth:`update`
  161. when necessary.
  162. """
  163. with self._mutex:
  164. if self.value:
  165. self._sub(n)
  166. def set(self, pcount):
  167. """Set channel prefetch_count setting."""
  168. if pcount != self.prev:
  169. new_value = pcount
  170. if pcount > PREFETCH_COUNT_MAX:
  171. warn("QoS: Disabled: prefetch_count exceeds %r",
  172. PREFETCH_COUNT_MAX)
  173. new_value = 0
  174. debug("basic.qos: prefetch_count->%s", new_value)
  175. self.consumer.qos(prefetch_count=new_value)
  176. self.prev = pcount
  177. return pcount
  178. def update(self):
  179. """Update prefetch count with current value."""
  180. with self._mutex:
  181. return self.set(self.value)
  182. class Consumer(object):
  183. """Listen for messages received from the broker and
  184. move them to the ready queue for task processing.
  185. :param ready_queue: See :attr:`ready_queue`.
  186. :param eta_schedule: See :attr:`eta_schedule`.
  187. """
  188. #: The queue that holds tasks ready for immediate processing.
  189. ready_queue = None
  190. #: Timer for tasks with an ETA/countdown.
  191. eta_schedule = None
  192. #: Enable/disable events.
  193. send_events = False
  194. #: Optional callback to be called when the connection is established.
  195. #: Will only be called once, even if the connection is lost and
  196. #: re-established.
  197. init_callback = None
  198. #: The current hostname. Defaults to the system hostname.
  199. hostname = None
  200. #: Initial QoS prefetch count for the task channel.
  201. initial_prefetch_count = 0
  202. #: A :class:`celery.events.EventDispatcher` for sending events.
  203. event_dispatcher = None
  204. #: The thread that sends event heartbeats at regular intervals.
  205. #: The heartbeats are used by monitors to detect that a worker
  206. #: went offline/disappeared.
  207. heart = None
  208. #: The broker connection.
  209. connection = None
  210. #: The consumer used to consume task messages.
  211. task_consumer = None
  212. #: The consumer used to consume broadcast commands.
  213. broadcast_consumer = None
  214. #: The process mailbox (kombu pidbox node).
  215. pidbox_node = None
  216. _pidbox_node_shutdown = None # used for greenlets
  217. _pidbox_node_stopped = None # used for greenlets
  218. #: The current worker pool instance.
  219. pool = None
  220. #: A timer used for high-priority internal tasks, such
  221. #: as sending heartbeats.
  222. priority_timer = None
  223. # Consumer state, can be RUN or CLOSE.
  224. _state = None
  225. def __init__(self, ready_queue, eta_schedule,
  226. init_callback=noop, send_events=False, hostname=None,
  227. initial_prefetch_count=2, pool=None, app=None,
  228. priority_timer=None, controller=None, **kwargs):
  229. self.app = app_or_default(app)
  230. self.connection = None
  231. self.task_consumer = None
  232. self.controller = controller
  233. self.broadcast_consumer = None
  234. self.ready_queue = ready_queue
  235. self.eta_schedule = eta_schedule
  236. self.send_events = send_events
  237. self.init_callback = init_callback
  238. self.hostname = hostname or socket.gethostname()
  239. self.initial_prefetch_count = initial_prefetch_count
  240. self.event_dispatcher = None
  241. self.heart = None
  242. self.pool = pool
  243. self.priority_timer = priority_timer or timer2.default_timer
  244. pidbox_state = AttributeDict(app=self.app,
  245. hostname=self.hostname,
  246. listener=self, # pre 2.2
  247. consumer=self)
  248. self.pidbox_node = self.app.control.mailbox.Node(self.hostname,
  249. state=pidbox_state,
  250. handlers=Panel.data)
  251. conninfo = self.app.broker_connection()
  252. self.connection_errors = conninfo.connection_errors
  253. self.channel_errors = conninfo.channel_errors
  254. self._does_info = logger.isEnabledFor(logging.INFO)
  255. self.strategies = {}
  256. def update_strategies(self):
  257. S = self.strategies
  258. for task in self.app.tasks.itervalues():
  259. S[task.name] = task.start_strategy(self.app, self)
  260. def start(self):
  261. """Start the consumer.
  262. Automatically survives intermittent connection failure,
  263. and will retry establishing the connection and restart
  264. consuming messages.
  265. """
  266. self.init_callback(self)
  267. while self._state != CLOSE:
  268. try:
  269. self.reset_connection()
  270. self.consume_messages()
  271. except self.connection_errors + self.channel_errors:
  272. error(RETRY_CONNECTION, exc_info=True)
  273. def consume_messages(self):
  274. """Consume messages forever (or until an exception is raised)."""
  275. debug("Starting message consumer...")
  276. self.task_consumer.consume()
  277. debug("Ready to accept tasks!")
  278. while self._state != CLOSE and self.connection:
  279. if self.qos.prev != self.qos.value: # pragma: no cover
  280. self.qos.update()
  281. try:
  282. self.connection.drain_events(timeout=1)
  283. except socket.timeout:
  284. pass
  285. except socket.error:
  286. if self._state != CLOSE: # pragma: no cover
  287. raise
  288. def on_task(self, task):
  289. """Handle received task.
  290. If the task has an `eta` we enter it into the ETA schedule,
  291. otherwise we move it the ready queue for immediate processing.
  292. """
  293. if task.revoked():
  294. return
  295. if self._does_info:
  296. info("Got task from broker: %s", task.shortinfo())
  297. if self.event_dispatcher.enabled:
  298. self.event_dispatcher.send("task-received", uuid=task.id,
  299. name=task.name, args=safe_repr(task.args),
  300. kwargs=safe_repr(task.kwargs),
  301. retries=task.request_dict.get("retries", 0),
  302. eta=task.eta and task.eta.isoformat(),
  303. expires=task.expires and task.expires.isoformat())
  304. if task.eta:
  305. try:
  306. eta = timer2.to_timestamp(task.eta)
  307. except OverflowError, exc:
  308. error("Couldn't convert eta %s to timestamp: %r. Task: %r",
  309. task.eta, exc, task.info(safe=True), exc_info=True)
  310. task.acknowledge()
  311. else:
  312. self.qos.increment()
  313. self.eta_schedule.apply_at(eta,
  314. self.apply_eta_task, (task, ))
  315. else:
  316. state.task_reserved(task)
  317. self.ready_queue.put(task)
  318. def on_control(self, body, message):
  319. """Process remote control command message."""
  320. try:
  321. self.pidbox_node.handle_message(body, message)
  322. except KeyError, exc:
  323. error("No such control command: %s", exc)
  324. except Exception, exc:
  325. error("Control command error: %r", exc, exc_info=True)
  326. self.reset_pidbox_node()
  327. def apply_eta_task(self, task):
  328. """Method called by the timer to apply a task with an
  329. ETA/countdown."""
  330. state.task_reserved(task)
  331. self.ready_queue.put(task)
  332. self.qos.decrement_eventually()
  333. def _message_report(self, body, message):
  334. return MESSAGE_REPORT_FMT % (safe_repr(body),
  335. safe_repr(message.content_type),
  336. safe_repr(message.content_encoding),
  337. safe_repr(message.delivery_info))
  338. def receive_message(self, body, message):
  339. """Handles incoming messages.
  340. :param body: The message body.
  341. :param message: The kombu message object.
  342. """
  343. try:
  344. name = body["task"]
  345. except (KeyError, TypeError):
  346. warn(UNKNOWN_FORMAT, self._message_report(body, message))
  347. message.reject_log_error(logger, self.connection_errors)
  348. return
  349. try:
  350. self.strategies[name](message, body, message.ack_log_error)
  351. except KeyError, exc:
  352. error(UNKNOWN_TASK_ERROR, exc, safe_repr(body), exc_info=True)
  353. message.reject_log_error(logger, self.connection_errors)
  354. except InvalidTaskError, exc:
  355. error(INVALID_TASK_ERROR, str(exc), safe_repr(body), exc_info=True)
  356. message.reject_log_error(logger, self.connection_errors)
  357. def maybe_conn_error(self, fun):
  358. """Applies function but ignores any connection or channel
  359. errors raised."""
  360. try:
  361. fun()
  362. except (AttributeError, ) + \
  363. self.connection_errors + \
  364. self.channel_errors:
  365. pass
  366. def close_connection(self):
  367. """Closes the current broker connection and all open channels."""
  368. # We must set self.connection to None here, so
  369. # that the green pidbox thread exits.
  370. connection, self.connection = self.connection, None
  371. if self.task_consumer:
  372. debug("Closing consumer channel...")
  373. self.task_consumer = \
  374. self.maybe_conn_error(self.task_consumer.close)
  375. self.stop_pidbox_node()
  376. if connection:
  377. debug("Closing broker connection...")
  378. self.maybe_conn_error(connection.close)
  379. def stop_consumers(self, close_connection=True):
  380. """Stop consuming tasks and broadcast commands, also stops
  381. the heartbeat thread and event dispatcher.
  382. :keyword close_connection: Set to False to skip closing the broker
  383. connection.
  384. """
  385. if not self._state == RUN:
  386. return
  387. if self.heart:
  388. # Stop the heartbeat thread if it's running.
  389. debug("Heart: Going into cardiac arrest...")
  390. self.heart = self.heart.stop()
  391. debug("Cancelling task consumer...")
  392. if self.task_consumer:
  393. self.maybe_conn_error(self.task_consumer.cancel)
  394. if self.event_dispatcher:
  395. debug("Shutting down event dispatcher...")
  396. self.event_dispatcher = \
  397. self.maybe_conn_error(self.event_dispatcher.close)
  398. debug("Cancelling broadcast consumer...")
  399. if self.broadcast_consumer:
  400. self.maybe_conn_error(self.broadcast_consumer.cancel)
  401. if close_connection:
  402. self.close_connection()
  403. def on_decode_error(self, message, exc):
  404. """Callback called if an error occurs while decoding
  405. a message received.
  406. Simply logs the error and acknowledges the message so it
  407. doesn't enter a loop.
  408. :param message: The message with errors.
  409. :param exc: The original exception instance.
  410. """
  411. crit("Can't decode message body: %r (type:%r encoding:%r raw:%r')",
  412. exc, message.content_type, message.content_encoding,
  413. safe_repr(message.body))
  414. message.ack()
  415. def reset_pidbox_node(self):
  416. """Sets up the process mailbox."""
  417. self.stop_pidbox_node()
  418. # close previously opened channel if any.
  419. if self.pidbox_node.channel:
  420. try:
  421. self.pidbox_node.channel.close()
  422. except self.connection_errors + self.channel_errors:
  423. pass
  424. if self.pool is not None and self.pool.is_green:
  425. return self.pool.spawn_n(self._green_pidbox_node)
  426. self.pidbox_node.channel = self.connection.channel()
  427. self.broadcast_consumer = self.pidbox_node.listen(
  428. callback=self.on_control)
  429. self.broadcast_consumer.consume()
  430. def stop_pidbox_node(self):
  431. if self._pidbox_node_stopped:
  432. self._pidbox_node_shutdown.set()
  433. debug("Waiting for broadcast thread to shutdown...")
  434. self._pidbox_node_stopped.wait()
  435. self._pidbox_node_stopped = self._pidbox_node_shutdown = None
  436. elif self.broadcast_consumer:
  437. debug("Closing broadcast channel...")
  438. self.broadcast_consumer = \
  439. self.maybe_conn_error(self.broadcast_consumer.channel.close)
  440. def _green_pidbox_node(self):
  441. """Sets up the process mailbox when running in a greenlet
  442. environment."""
  443. # THIS CODE IS TERRIBLE
  444. # Luckily work has already started rewriting the Consumer for 3.0.
  445. self._pidbox_node_shutdown = threading.Event()
  446. self._pidbox_node_stopped = threading.Event()
  447. try:
  448. with self._open_connection() as conn:
  449. self.pidbox_node.channel = conn.default_channel
  450. self.broadcast_consumer = self.pidbox_node.listen(
  451. callback=self.on_control)
  452. with self.broadcast_consumer:
  453. while not self._pidbox_node_shutdown.isSet():
  454. try:
  455. conn.drain_events(timeout=1.0)
  456. except socket.timeout:
  457. pass
  458. finally:
  459. self._pidbox_node_stopped.set()
  460. def reset_connection(self):
  461. """Re-establish the broker connection and set up consumers,
  462. heartbeat and the event dispatcher."""
  463. debug("Re-establishing connection to the broker...")
  464. self.stop_consumers()
  465. # Clear internal queues to get rid of old messages.
  466. # They can't be acked anyway, as a delivery tag is specific
  467. # to the current channel.
  468. self.ready_queue.clear()
  469. self.eta_schedule.clear()
  470. # Re-establish the broker connection and setup the task consumer.
  471. self.connection = self._open_connection()
  472. debug("Connection established.")
  473. self.task_consumer = self.app.amqp.get_task_consumer(self.connection,
  474. on_decode_error=self.on_decode_error)
  475. # QoS: Reset prefetch window.
  476. self.qos = QoS(self.task_consumer, self.initial_prefetch_count)
  477. self.qos.update()
  478. # receive_message handles incoming messages.
  479. self.task_consumer.register_callback(self.receive_message)
  480. # Setup the process mailbox.
  481. self.reset_pidbox_node()
  482. # Flush events sent while connection was down.
  483. prev_event_dispatcher = self.event_dispatcher
  484. self.event_dispatcher = self.app.events.Dispatcher(self.connection,
  485. hostname=self.hostname,
  486. enabled=self.send_events)
  487. if prev_event_dispatcher:
  488. self.event_dispatcher.copy_buffer(prev_event_dispatcher)
  489. self.event_dispatcher.flush()
  490. # Restart heartbeat thread.
  491. self.restart_heartbeat()
  492. # reload all task's execution strategies.
  493. self.update_strategies()
  494. # We're back!
  495. self._state = RUN
  496. def restart_heartbeat(self):
  497. """Restart the heartbeat thread.
  498. This thread sends heartbeat events at intervals so monitors
  499. can tell if the worker is off-line/missing.
  500. """
  501. self.heart = Heart(self.priority_timer, self.event_dispatcher)
  502. self.heart.start()
  503. def _open_connection(self):
  504. """Establish the broker connection.
  505. Will retry establishing the connection if the
  506. :setting:`BROKER_CONNECTION_RETRY` setting is enabled
  507. """
  508. # Callback called for each retry while the connection
  509. # can't be established.
  510. def _error_handler(exc, interval):
  511. error("Consumer: Connection Error: %s. "
  512. "Trying again in %d seconds...", exc, interval)
  513. # remember that the connection is lazy, it won't establish
  514. # until it's needed.
  515. conn = self.app.broker_connection()
  516. if not self.app.conf.BROKER_CONNECTION_RETRY:
  517. # retry disabled, just call connect directly.
  518. conn.connect()
  519. return conn
  520. return conn.ensure_connection(_error_handler,
  521. self.app.conf.BROKER_CONNECTION_MAX_RETRIES)
  522. def stop(self):
  523. """Stop consuming.
  524. Does not close the broker connection, so be sure to call
  525. :meth:`close_connection` when you are finished with it.
  526. """
  527. # Notifies other threads that this instance can't be used
  528. # anymore.
  529. self._state = CLOSE
  530. debug("Stopping consumers...")
  531. self.stop_consumers(close_connection=False)
  532. @property
  533. def info(self):
  534. """Returns information about this consumer instance
  535. as a dict.
  536. This is also the consumer related info returned by
  537. ``celeryctl stats``.
  538. """
  539. conninfo = {}
  540. if self.connection:
  541. conninfo = self.connection.info()
  542. conninfo.pop("password", None) # don't send password.
  543. return {"broker": conninfo, "prefetch_count": self.qos.value}