consumer.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854
  1. # -*- coding: utf-8 -*-
  2. """
  3. celery.worker.consumer
  4. ~~~~~~~~~~~~~~~~~~~~~~
  5. This module contains the components responsible for consuming messages
  6. from the broker, processing the messages and keeping the broker connections
  7. up and running.
  8. """
  9. from __future__ import absolute_import
  10. import errno
  11. import kombu
  12. import logging
  13. import os
  14. import socket
  15. from collections import defaultdict
  16. from functools import partial
  17. from heapq import heappush
  18. from operator import itemgetter
  19. from time import sleep
  20. from billiard.common import restart_state
  21. from billiard.exceptions import RestartFreqExceeded
  22. from kombu.async.semaphore import DummyLock
  23. from kombu.common import QoS, ignore_errors
  24. from kombu.five import buffer_t, items, values
  25. from kombu.syn import _detect_environment
  26. from kombu.utils.encoding import safe_repr, bytes_t
  27. from kombu.utils.limits import TokenBucket
  28. from celery import bootsteps
  29. from celery.app.trace import build_tracer
  30. from celery.canvas import signature
  31. from celery.exceptions import InvalidTaskError
  32. from celery.utils.functional import noop
  33. from celery.utils.log import get_logger
  34. from celery.utils.text import truncate
  35. from celery.utils.timeutils import humanize_seconds, rate
  36. from . import heartbeat, loops, pidbox
  37. from .state import task_reserved, maybe_shutdown, revoked, reserved_requests
  38. __all__ = [
  39. 'Consumer', 'Connection', 'Events', 'Heart', 'Control',
  40. 'Tasks', 'Evloop', 'Agent', 'Mingle', 'Gossip', 'dump_body',
  41. ]
  42. CLOSE = bootsteps.CLOSE
  43. logger = get_logger(__name__)
  44. debug, info, warn, error, crit = (logger.debug, logger.info, logger.warning,
  45. logger.error, logger.critical)
  46. CONNECTION_RETRY = """\
  47. consumer: Connection to broker lost. \
  48. Trying to re-establish the connection...\
  49. """
  50. CONNECTION_RETRY_STEP = """\
  51. Trying again {when}...\
  52. """
  53. CONNECTION_ERROR = """\
  54. consumer: Cannot connect to %s: %s.
  55. %s
  56. """
  57. CONNECTION_FAILOVER = """\
  58. Will retry using next failover.\
  59. """
  60. UNKNOWN_FORMAT = """\
  61. Received and deleted unknown message. Wrong destination?!?
  62. The full contents of the message body was: %s
  63. """
  64. #: Error message for when an unregistered task is received.
  65. UNKNOWN_TASK_ERROR = """\
  66. Received unregistered task of type %s.
  67. The message has been ignored and discarded.
  68. Did you remember to import the module containing this task?
  69. Or maybe you are using relative imports?
  70. Please see http://bit.ly/gLye1c for more information.
  71. The full contents of the message body was:
  72. %s
  73. """
  74. #: Error message for when an invalid task message is received.
  75. INVALID_TASK_ERROR = """\
  76. Received invalid task message: %s
  77. The message has been ignored and discarded.
  78. Please ensure your message conforms to the task
  79. message protocol as described here: http://bit.ly/hYj41y
  80. The full contents of the message body was:
  81. %s
  82. """
  83. MESSAGE_DECODE_ERROR = """\
  84. Can't decode message body: %r [type:%r encoding:%r headers:%s]
  85. body: %s
  86. """
  87. MESSAGE_REPORT = """\
  88. body: {0}
  89. {{content_type:{1} content_encoding:{2}
  90. delivery_info:{3} headers={4}}}
  91. """
  92. MINGLE_GET_FIELDS = itemgetter('clock', 'revoked')
  93. def dump_body(m, body):
  94. # v2 protocol does not deserialize body
  95. body = m.body if body is None else body
  96. if isinstance(body, buffer_t):
  97. body = bytes_t(body)
  98. return '{0} ({1}b)'.format(truncate(safe_repr(body), 1024),
  99. len(m.body))
  100. class Consumer(object):
  101. Strategies = dict
  102. #: set when consumer is shutting down.
  103. in_shutdown = False
  104. #: Optional callback called the first time the worker
  105. #: is ready to receive tasks.
  106. init_callback = None
  107. #: The current worker pool instance.
  108. pool = None
  109. #: A timer used for high-priority internal tasks, such
  110. #: as sending heartbeats.
  111. timer = None
  112. restart_count = -1 # first start is the same as a restart
  113. class Blueprint(bootsteps.Blueprint):
  114. name = 'Consumer'
  115. default_steps = [
  116. 'celery.worker.consumer:Connection',
  117. 'celery.worker.consumer:Mingle',
  118. 'celery.worker.consumer:Events',
  119. 'celery.worker.consumer:Gossip',
  120. 'celery.worker.consumer:Heart',
  121. 'celery.worker.consumer:Control',
  122. 'celery.worker.consumer:Tasks',
  123. 'celery.worker.consumer:Evloop',
  124. 'celery.worker.consumer:Agent',
  125. ]
  126. def shutdown(self, parent):
  127. self.send_all(parent, 'shutdown')
  128. def __init__(self, on_task_request,
  129. init_callback=noop, hostname=None,
  130. pool=None, app=None,
  131. timer=None, controller=None, hub=None, amqheartbeat=None,
  132. worker_options=None, disable_rate_limits=False,
  133. initial_prefetch_count=2, prefetch_multiplier=1, **kwargs):
  134. self.app = app
  135. self.controller = controller
  136. self.init_callback = init_callback
  137. self.hostname = hostname or socket.gethostname()
  138. self.pid = os.getpid()
  139. self.pool = pool
  140. self.timer = timer
  141. self.strategies = self.Strategies()
  142. conninfo = self.app.connection()
  143. self.connection_errors = conninfo.connection_errors
  144. self.channel_errors = conninfo.channel_errors
  145. self._restart_state = restart_state(maxR=5, maxT=1)
  146. self._does_info = logger.isEnabledFor(logging.INFO)
  147. self._limit_order = 0
  148. self.on_task_request = on_task_request
  149. self.on_task_message = set()
  150. self.amqheartbeat_rate = self.app.conf.BROKER_HEARTBEAT_CHECKRATE
  151. self.disable_rate_limits = disable_rate_limits
  152. self.initial_prefetch_count = initial_prefetch_count
  153. self.prefetch_multiplier = prefetch_multiplier
  154. # this contains a tokenbucket for each task type by name, used for
  155. # rate limits, or None if rate limits are disabled for that task.
  156. self.task_buckets = defaultdict(lambda: None)
  157. self.reset_rate_limits()
  158. self.hub = hub
  159. if self.hub:
  160. self.amqheartbeat = amqheartbeat
  161. if self.amqheartbeat is None:
  162. self.amqheartbeat = self.app.conf.BROKER_HEARTBEAT
  163. else:
  164. self.amqheartbeat = 0
  165. if not hasattr(self, 'loop'):
  166. self.loop = loops.asynloop if hub else loops.synloop
  167. if _detect_environment() == 'gevent':
  168. # there's a gevent bug that causes timeouts to not be reset,
  169. # so if the connection timeout is exceeded once, it can NEVER
  170. # connect again.
  171. self.app.conf.BROKER_CONNECTION_TIMEOUT = None
  172. self.steps = []
  173. self.blueprint = self.Blueprint(
  174. app=self.app, on_close=self.on_close,
  175. )
  176. self.blueprint.apply(self, **dict(worker_options or {}, **kwargs))
  177. def bucket_for_task(self, type):
  178. limit = rate(getattr(type, 'rate_limit', None))
  179. return TokenBucket(limit, capacity=1) if limit else None
  180. def reset_rate_limits(self):
  181. self.task_buckets.update(
  182. (n, self.bucket_for_task(t)) for n, t in items(self.app.tasks)
  183. )
  184. def _update_prefetch_count(self, index=0):
  185. """Update prefetch count after pool/shrink grow operations.
  186. Index must be the change in number of processes as a positive
  187. (increasing) or negative (decreasing) number.
  188. .. note::
  189. Currently pool grow operations will end up with an offset
  190. of +1 if the initial size of the pool was 0 (e.g.
  191. ``--autoscale=1,0``).
  192. """
  193. num_processes = self.pool.num_processes
  194. if not self.initial_prefetch_count or not num_processes:
  195. return # prefetch disabled
  196. self.initial_prefetch_count = (
  197. self.pool.num_processes * self.prefetch_multiplier
  198. )
  199. return self._update_qos_eventually(index)
  200. def _update_qos_eventually(self, index):
  201. return (self.qos.decrement_eventually if index < 0
  202. else self.qos.increment_eventually)(
  203. abs(index) * self.prefetch_multiplier)
  204. def _limit_move_to_pool(self, request):
  205. task_reserved(request)
  206. self.on_task_request(request)
  207. def _limit_task(self, request, bucket, tokens):
  208. if not bucket.can_consume(tokens):
  209. hold = bucket.expected_time(tokens)
  210. pri = self._limit_order = (self._limit_order + 1) % 10
  211. self.timer.call_after(
  212. hold, self._limit_move_to_pool, (request, ),
  213. priority=pri,
  214. )
  215. else:
  216. task_reserved(request)
  217. self.on_task_request(request)
  218. def start(self):
  219. blueprint = self.blueprint
  220. while blueprint.state != CLOSE:
  221. self.restart_count += 1
  222. maybe_shutdown()
  223. try:
  224. blueprint.start(self)
  225. except self.connection_errors as exc:
  226. if isinstance(exc, OSError) and exc.errno == errno.EMFILE:
  227. raise # Too many open files
  228. maybe_shutdown()
  229. try:
  230. self._restart_state.step()
  231. except RestartFreqExceeded as exc:
  232. crit('Frequent restarts detected: %r', exc, exc_info=1)
  233. sleep(1)
  234. if blueprint.state != CLOSE and self.connection:
  235. warn(CONNECTION_RETRY, exc_info=True)
  236. try:
  237. self.connection.collect()
  238. except Exception:
  239. pass
  240. self.on_close()
  241. blueprint.restart(self)
  242. def register_with_event_loop(self, hub):
  243. self.blueprint.send_all(
  244. self, 'register_with_event_loop', args=(hub, ),
  245. description='Hub.register',
  246. )
  247. def shutdown(self):
  248. self.in_shutdown = True
  249. self.blueprint.shutdown(self)
  250. def stop(self):
  251. self.blueprint.stop(self)
  252. def on_ready(self):
  253. callback, self.init_callback = self.init_callback, None
  254. if callback:
  255. callback(self)
  256. def loop_args(self):
  257. return (self, self.connection, self.task_consumer,
  258. self.blueprint, self.hub, self.qos, self.amqheartbeat,
  259. self.app.clock, self.amqheartbeat_rate)
  260. def on_decode_error(self, message, exc):
  261. """Callback called if an error occurs while decoding
  262. a message received.
  263. Simply logs the error and acknowledges the message so it
  264. doesn't enter a loop.
  265. :param message: The message with errors.
  266. :param exc: The original exception instance.
  267. """
  268. crit(MESSAGE_DECODE_ERROR,
  269. exc, message.content_type, message.content_encoding,
  270. safe_repr(message.headers), dump_body(message, message.body),
  271. exc_info=1)
  272. message.ack()
  273. def on_close(self):
  274. # Clear internal queues to get rid of old messages.
  275. # They can't be acked anyway, as a delivery tag is specific
  276. # to the current channel.
  277. if self.controller and self.controller.semaphore:
  278. self.controller.semaphore.clear()
  279. if self.timer:
  280. self.timer.clear()
  281. reserved_requests.clear()
  282. if self.pool and self.pool.flush:
  283. self.pool.flush()
  284. def connect(self):
  285. """Establish the broker connection.
  286. Will retry establishing the connection if the
  287. :setting:`BROKER_CONNECTION_RETRY` setting is enabled
  288. """
  289. conn = self.app.connection(heartbeat=self.amqheartbeat)
  290. # Callback called for each retry while the connection
  291. # can't be established.
  292. def _error_handler(exc, interval, next_step=CONNECTION_RETRY_STEP):
  293. if getattr(conn, 'alt', None) and interval == 0:
  294. next_step = CONNECTION_FAILOVER
  295. error(CONNECTION_ERROR, conn.as_uri(), exc,
  296. next_step.format(when=humanize_seconds(interval, 'in', ' ')))
  297. # remember that the connection is lazy, it won't establish
  298. # until needed.
  299. if not self.app.conf.BROKER_CONNECTION_RETRY:
  300. # retry disabled, just call connect directly.
  301. conn.connect()
  302. return conn
  303. conn = conn.ensure_connection(
  304. _error_handler, self.app.conf.BROKER_CONNECTION_MAX_RETRIES,
  305. callback=maybe_shutdown,
  306. )
  307. if self.hub:
  308. conn.transport.register_with_event_loop(conn.connection, self.hub)
  309. return conn
  310. def _flush_events(self):
  311. if self.event_dispatcher:
  312. self.event_dispatcher.flush()
  313. def on_send_event_buffered(self):
  314. if self.hub:
  315. self.hub._ready.add(self._flush_events)
  316. def add_task_queue(self, queue, exchange=None, exchange_type=None,
  317. routing_key=None, **options):
  318. cset = self.task_consumer
  319. queues = self.app.amqp.queues
  320. # Must use in' here, as __missing__ will automatically
  321. # create queues when CELERY_CREATE_MISSING_QUEUES is enabled.
  322. # (Issue #1079)
  323. if queue in queues:
  324. q = queues[queue]
  325. else:
  326. exchange = queue if exchange is None else exchange
  327. exchange_type = ('direct' if exchange_type is None
  328. else exchange_type)
  329. q = queues.select_add(queue,
  330. exchange=exchange,
  331. exchange_type=exchange_type,
  332. routing_key=routing_key, **options)
  333. if not cset.consuming_from(queue):
  334. cset.add_queue(q)
  335. cset.consume()
  336. info('Started consuming from %s', queue)
  337. def cancel_task_queue(self, queue):
  338. info('Cancelling queue %s', queue)
  339. self.app.amqp.queues.deselect(queue)
  340. self.task_consumer.cancel_by_queue(queue)
  341. def apply_eta_task(self, task):
  342. """Method called by the timer to apply a task with an
  343. ETA/countdown."""
  344. task_reserved(task)
  345. self.on_task_request(task)
  346. self.qos.decrement_eventually()
  347. def _message_report(self, body, message):
  348. return MESSAGE_REPORT.format(dump_body(message, body),
  349. safe_repr(message.content_type),
  350. safe_repr(message.content_encoding),
  351. safe_repr(message.delivery_info),
  352. safe_repr(message.headers))
  353. def on_unknown_message(self, body, message):
  354. warn(UNKNOWN_FORMAT, self._message_report(body, message))
  355. message.reject_log_error(logger, self.connection_errors)
  356. def on_unknown_task(self, body, message, exc):
  357. error(UNKNOWN_TASK_ERROR, exc, dump_body(message, body), exc_info=True)
  358. message.reject_log_error(logger, self.connection_errors)
  359. def on_invalid_task(self, body, message, exc):
  360. error(INVALID_TASK_ERROR, exc, dump_body(message, body), exc_info=True)
  361. message.reject_log_error(logger, self.connection_errors)
  362. def update_strategies(self):
  363. loader = self.app.loader
  364. for name, task in items(self.app.tasks):
  365. self.strategies[name] = task.start_strategy(self.app, self)
  366. task.__trace__ = build_tracer(name, task, loader, self.hostname,
  367. app=self.app)
  368. def create_task_handler(self):
  369. strategies = self.strategies
  370. on_unknown_message = self.on_unknown_message
  371. on_unknown_task = self.on_unknown_task
  372. on_invalid_task = self.on_invalid_task
  373. callbacks = self.on_task_message
  374. def on_task_received(message):
  375. # payload will only be set for v1 protocol, since v2
  376. # will defer deserializing the message body to the pool.
  377. payload = None
  378. try:
  379. type_ = message.headers['task'] # protocol v2
  380. except TypeError:
  381. return on_unknown_message(None, message)
  382. except KeyError:
  383. payload = message.payload
  384. try:
  385. type_, payload = payload['task'], payload # protocol v1
  386. except (TypeError, KeyError):
  387. return on_unknown_message(payload, message)
  388. try:
  389. strategy = strategies[type_]
  390. except KeyError as exc:
  391. return on_unknown_task(payload, message, exc)
  392. else:
  393. try:
  394. strategy(
  395. message, payload, message.ack_log_error,
  396. message.reject_log_error, callbacks,
  397. )
  398. except InvalidTaskError as exc:
  399. return on_invalid_task(payload, message, exc)
  400. except MemoryError:
  401. raise
  402. except Exception as exc:
  403. # XXX handle as internal error?
  404. return on_invalid_task(payload, message, exc)
  405. return on_task_received
  406. def __repr__(self):
  407. return '<Consumer: {self.hostname} ({state})>'.format(
  408. self=self, state=self.blueprint.human_state(),
  409. )
  410. class Connection(bootsteps.StartStopStep):
  411. def __init__(self, c, **kwargs):
  412. c.connection = None
  413. def start(self, c):
  414. c.connection = c.connect()
  415. info('Connected to %s', c.connection.as_uri())
  416. def shutdown(self, c):
  417. # We must set self.connection to None here, so
  418. # that the green pidbox thread exits.
  419. connection, c.connection = c.connection, None
  420. if connection:
  421. ignore_errors(connection, connection.close)
  422. def info(self, c, params='N/A'):
  423. if c.connection:
  424. params = c.connection.info()
  425. params.pop('password', None) # don't send password.
  426. return {'broker': params}
  427. class Events(bootsteps.StartStopStep):
  428. requires = (Connection, )
  429. def __init__(self, c, send_events=None, **kwargs):
  430. self.send_events = True
  431. self.groups = None if send_events else ['worker']
  432. c.event_dispatcher = None
  433. def start(self, c):
  434. # flush events sent while connection was down.
  435. prev = self._close(c)
  436. dis = c.event_dispatcher = c.app.events.Dispatcher(
  437. c.connect(), hostname=c.hostname,
  438. enabled=self.send_events, groups=self.groups,
  439. buffer_group=['task'] if c.hub else None,
  440. on_send_buffered=c.on_send_event_buffered if c.hub else None,
  441. )
  442. if prev:
  443. dis.extend_buffer(prev)
  444. dis.flush()
  445. def stop(self, c):
  446. pass
  447. def _close(self, c):
  448. if c.event_dispatcher:
  449. dispatcher = c.event_dispatcher
  450. # remember changes from remote control commands:
  451. self.groups = dispatcher.groups
  452. # close custom connection
  453. if dispatcher.connection:
  454. ignore_errors(c, dispatcher.connection.close)
  455. ignore_errors(c, dispatcher.close)
  456. c.event_dispatcher = None
  457. return dispatcher
  458. def shutdown(self, c):
  459. self._close(c)
  460. class Heart(bootsteps.StartStopStep):
  461. requires = (Events, )
  462. def __init__(self, c, without_heartbeat=False, heartbeat_interval=None,
  463. **kwargs):
  464. self.enabled = not without_heartbeat
  465. self.heartbeat_interval = heartbeat_interval
  466. c.heart = None
  467. def start(self, c):
  468. c.heart = heartbeat.Heart(
  469. c.timer, c.event_dispatcher, self.heartbeat_interval,
  470. )
  471. c.heart.start()
  472. def stop(self, c):
  473. c.heart = c.heart and c.heart.stop()
  474. shutdown = stop
  475. class Mingle(bootsteps.StartStopStep):
  476. label = 'Mingle'
  477. requires = (Events, )
  478. compatible_transports = {'amqp', 'redis'}
  479. def __init__(self, c, without_mingle=False, **kwargs):
  480. self.enabled = not without_mingle and self.compatible_transport(c.app)
  481. def compatible_transport(self, app):
  482. with app.connection() as conn:
  483. return conn.transport.driver_type in self.compatible_transports
  484. def start(self, c):
  485. info('mingle: searching for neighbors')
  486. I = c.app.control.inspect(timeout=1.0, connection=c.connection)
  487. replies = I.hello(c.hostname, revoked._data) or {}
  488. replies.pop(c.hostname, None)
  489. if replies:
  490. info('mingle: sync with %s nodes',
  491. len([reply for reply, value in items(replies) if value]))
  492. for reply in values(replies):
  493. if reply:
  494. try:
  495. other_clock, other_revoked = MINGLE_GET_FIELDS(reply)
  496. except KeyError: # reply from pre-3.1 worker
  497. pass
  498. else:
  499. c.app.clock.adjust(other_clock)
  500. revoked.update(other_revoked)
  501. info('mingle: sync complete')
  502. else:
  503. info('mingle: all alone')
  504. class Tasks(bootsteps.StartStopStep):
  505. requires = (Mingle, )
  506. def __init__(self, c, **kwargs):
  507. c.task_consumer = c.qos = None
  508. def start(self, c):
  509. c.update_strategies()
  510. # - RabbitMQ 3.3 completely redefines how basic_qos works..
  511. # This will detect if the new qos smenatics is in effect,
  512. # and if so make sure the 'apply_global' flag is set on qos updates.
  513. qos_global = not c.connection.qos_semantics_matches_spec
  514. # set initial prefetch count
  515. c.connection.default_channel.basic_qos(
  516. 0, c.initial_prefetch_count, qos_global,
  517. )
  518. c.task_consumer = c.app.amqp.TaskConsumer(
  519. c.connection, on_decode_error=c.on_decode_error,
  520. )
  521. def set_prefetch_count(prefetch_count):
  522. return c.task_consumer.qos(
  523. prefetch_count=prefetch_count,
  524. apply_global=qos_global,
  525. )
  526. c.qos = QoS(set_prefetch_count, c.initial_prefetch_count)
  527. def stop(self, c):
  528. if c.task_consumer:
  529. debug('Cancelling task consumer...')
  530. ignore_errors(c, c.task_consumer.cancel)
  531. def shutdown(self, c):
  532. if c.task_consumer:
  533. self.stop(c)
  534. debug('Closing consumer channel...')
  535. ignore_errors(c, c.task_consumer.close)
  536. c.task_consumer = None
  537. def info(self, c):
  538. return {'prefetch_count': c.qos.value if c.qos else 'N/A'}
  539. class Agent(bootsteps.StartStopStep):
  540. conditional = True
  541. requires = (Connection, )
  542. def __init__(self, c, **kwargs):
  543. self.agent_cls = self.enabled = c.app.conf.CELERYD_AGENT
  544. def create(self, c):
  545. agent = c.agent = self.instantiate(self.agent_cls, c.connection)
  546. return agent
  547. class Control(bootsteps.StartStopStep):
  548. requires = (Tasks, )
  549. def __init__(self, c, **kwargs):
  550. self.is_green = c.pool is not None and c.pool.is_green
  551. self.box = (pidbox.gPidbox if self.is_green else pidbox.Pidbox)(c)
  552. self.start = self.box.start
  553. self.stop = self.box.stop
  554. self.shutdown = self.box.shutdown
  555. def include_if(self, c):
  556. return c.app.conf.CELERY_ENABLE_REMOTE_CONTROL
  557. class Gossip(bootsteps.ConsumerStep):
  558. label = 'Gossip'
  559. requires = (Mingle, )
  560. _cons_stamp_fields = itemgetter(
  561. 'id', 'clock', 'hostname', 'pid', 'topic', 'action', 'cver',
  562. )
  563. compatible_transports = {'amqp', 'redis'}
  564. def __init__(self, c, without_gossip=False,
  565. interval=5.0, heartbeat_interval=2.0, **kwargs):
  566. self.enabled = not without_gossip and self.compatible_transport(c.app)
  567. self.app = c.app
  568. c.gossip = self
  569. self.Receiver = c.app.events.Receiver
  570. self.hostname = c.hostname
  571. self.full_hostname = '.'.join([self.hostname, str(c.pid)])
  572. self.timer = c.timer
  573. if self.enabled:
  574. self.state = c.app.events.State(
  575. on_node_join=self.on_node_join,
  576. on_node_leave=self.on_node_leave,
  577. max_tasks_in_memory=1,
  578. )
  579. if c.hub:
  580. c._mutex = DummyLock()
  581. self.update_state = self.state.event
  582. self.interval = interval
  583. self.heartbeat_interval = heartbeat_interval
  584. self._tref = None
  585. self.consensus_requests = defaultdict(list)
  586. self.consensus_replies = {}
  587. self.event_handlers = {
  588. 'worker.elect': self.on_elect,
  589. 'worker.elect.ack': self.on_elect_ack,
  590. }
  591. self.clock = c.app.clock
  592. self.election_handlers = {
  593. 'task': self.call_task
  594. }
  595. def compatible_transport(self, app):
  596. with app.connection() as conn:
  597. return conn.transport.driver_type in self.compatible_transports
  598. def election(self, id, topic, action=None):
  599. self.consensus_replies[id] = []
  600. self.dispatcher.send(
  601. 'worker-elect',
  602. id=id, topic=topic, action=action, cver=1,
  603. )
  604. def call_task(self, task):
  605. try:
  606. signature(task, app=self.app).apply_async()
  607. except Exception as exc:
  608. error('Could not call task: %r', exc, exc_info=1)
  609. def on_elect(self, event):
  610. try:
  611. (id_, clock, hostname, pid,
  612. topic, action, _) = self._cons_stamp_fields(event)
  613. except KeyError as exc:
  614. return error('election request missing field %s', exc, exc_info=1)
  615. heappush(
  616. self.consensus_requests[id_],
  617. (clock, '%s.%s' % (hostname, pid), topic, action),
  618. )
  619. self.dispatcher.send('worker-elect-ack', id=id_)
  620. def start(self, c):
  621. super(Gossip, self).start(c)
  622. self.dispatcher = c.event_dispatcher
  623. def on_elect_ack(self, event):
  624. id = event['id']
  625. try:
  626. replies = self.consensus_replies[id]
  627. except KeyError:
  628. return # not for us
  629. alive_workers = self.state.alive_workers()
  630. replies.append(event['hostname'])
  631. if len(replies) >= len(alive_workers):
  632. _, leader, topic, action = self.clock.sort_heap(
  633. self.consensus_requests[id],
  634. )
  635. if leader == self.full_hostname:
  636. info('I won the election %r', id)
  637. try:
  638. handler = self.election_handlers[topic]
  639. except KeyError:
  640. error('Unknown election topic %r', topic, exc_info=1)
  641. else:
  642. handler(action)
  643. else:
  644. info('node %s elected for %r', leader, id)
  645. self.consensus_requests.pop(id, None)
  646. self.consensus_replies.pop(id, None)
  647. def on_node_join(self, worker):
  648. debug('%s joined the party', worker.hostname)
  649. def on_node_leave(self, worker):
  650. debug('%s left', worker.hostname)
  651. def on_node_lost(self, worker):
  652. info('missed heartbeat from %s', worker.hostname)
  653. def register_timer(self):
  654. if self._tref is not None:
  655. self._tref.cancel()
  656. self._tref = self.timer.call_repeatedly(self.interval, self.periodic)
  657. def periodic(self):
  658. workers = self.state.workers
  659. dirty = set()
  660. for worker in values(workers):
  661. if not worker.alive:
  662. dirty.add(worker)
  663. self.on_node_lost(worker)
  664. for worker in dirty:
  665. workers.pop(worker.hostname, None)
  666. def get_consumers(self, channel):
  667. self.register_timer()
  668. ev = self.Receiver(channel, routing_key='worker.#',
  669. queue_ttl=self.heartbeat_interval)
  670. return [kombu.Consumer(
  671. channel,
  672. queues=[ev.queue],
  673. on_message=partial(self.on_message, ev.event_from_message),
  674. no_ack=True
  675. )]
  676. def on_message(self, prepare, message):
  677. _type = message.delivery_info['routing_key']
  678. # For redis when `fanout_patterns=False` (See Issue #1882)
  679. if _type.split('.', 1)[0] == 'task':
  680. return
  681. try:
  682. handler = self.event_handlers[_type]
  683. except KeyError:
  684. pass
  685. else:
  686. return handler(message.payload)
  687. hostname = (message.headers.get('hostname') or
  688. message.payload['hostname'])
  689. if hostname != self.hostname:
  690. type, event = prepare(message.payload)
  691. self.update_state(event)
  692. else:
  693. self.clock.forward()
  694. class Evloop(bootsteps.StartStopStep):
  695. label = 'event loop'
  696. last = True
  697. def start(self, c):
  698. self.patch_all(c)
  699. c.loop(*c.loop_args())
  700. def patch_all(self, c):
  701. c.qos._mutex = DummyLock()