def test_receive_message_control_command(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, control={"command": "shutdown"}) l.event_dispatcher = MockEventDispatcher() l.control_dispatch = MockControlDispatch() l.receive_message(m.decode(), m) self.assertIn("shutdown", l.control_dispatch.commands)
def test_receieve_message_not_registered(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, task="x.X.31x", args=[2, 4, 8], kwargs={}) l.event_dispatcher = MockEventDispatcher() self.assertFalse(l.receive_message(m.decode(), m)) self.assertRaises(Empty, self.ready_queue.get_nowait) self.assertTrue(self.eta_schedule.empty())
def test_receive_message_unknown(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, unknown={"baz": "!!!"}) l.event_dispatcher = MockEventDispatcher() l.control_dispatch = MockControlDispatch() import warnings with warnings.catch_warnings(record=True) as log: l.receive_message(m.decode(), m) self.assertTrue(log) self.assertTrue("unknown message" in log[0].message.args[0])
def test_receive_message_InvalidTaskError(self): logger = MockLogger() l = CarrotListener(self.ready_queue, self.eta_schedule, logger, send_events=False) backend = MockBackend() m = create_message(backend, task=foo_task.name, args=(1, 2), kwargs="foobarbaz", id=1) l.event_dispatcher = MockEventDispatcher() l.control_dispatch = MockControlDispatch() l.receive_message(m.decode(), m) self.assertIn("Invalid task ignored", logger.logged[0])
def test_receieve_message_eta(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, task=foo_task.name, args=[2, 4, 8], kwargs={}, eta=(datetime.now() + timedelta(days=1)).isoformat()) l.reset_connection() p, conf.BROKER_CONNECTION_RETRY = conf.BROKER_CONNECTION_RETRY, False try: l.reset_connection() finally: conf.BROKER_CONNECTION_RETRY = p l.receive_message(m.decode(), m) in_hold = self.eta_schedule.queue[0] self.assertEqual(len(in_hold), 4) eta, priority, task, on_accept = in_hold self.assertIsInstance(task, TaskRequest) self.assertTrue(callable(on_accept)) self.assertEqual(task.task_name, foo_task.name) self.assertEqual(task.execute(), 2 * 4 * 8) self.assertRaises(Empty, self.ready_queue.get_nowait)
def test_receieve_message(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, task=foo_task.name, args=[2, 4, 8], kwargs={}) l.event_dispatcher = MockEventDispatcher() l.receive_message(m.decode(), m) in_bucket = self.ready_queue.get_nowait() self.assertIsInstance(in_bucket, TaskRequest) self.assertEqual(in_bucket.task_name, foo_task.name) self.assertEqual(in_bucket.execute(), 2 * 4 * 8) self.assertTrue(self.eta_schedule.empty())
def test_receieve_message_eta_isoformat(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, task=foo_task.name, eta=datetime.now().isoformat(), args=[2, 4, 8], kwargs={}) l.event_dispatcher = MockEventDispatcher() l.receive_message(m.decode(), m) items = [entry[2] for entry in self.eta_schedule.queue] found = 0 for item in items: if item.task_name == foo_task.name: found = True self.assertTrue(found)
def test_receive_message_unknown(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, unknown={"baz": "!!!"}) l.event_dispatcher = MockEventDispatcher() l.control_dispatch = MockControlDispatch() def with_catch_warnings(log): l.receive_message(m.decode(), m) self.assertTrue(log) self.assertIn("unknown message", log[0].message.args[0]) context = catch_warnings(record=True) execute_context(context, with_catch_warnings)
def test_receieve_message_eta_isoformat(self): class MockConsumer(object): prefetch_count_incremented = False def qos(self, **kwargs): self.prefetch_count_incremented = True l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, task=foo_task.name, eta=datetime.now().isoformat(), args=[2, 4, 8], kwargs={}) l.event_dispatcher = MockEventDispatcher() l.task_consumer = MockConsumer() l.qos = QoS(l.task_consumer, l.initial_prefetch_count, l.logger) l.receive_message(m.decode(), m) items = [entry[2] for entry in self.eta_schedule.queue] found = 0 for item in items: if item.task_name == foo_task.name: found = True self.assertTrue(found) self.assertTrue(l.task_consumer.prefetch_count_incremented)
def test_revoke(self): ready_queue = FastQueue() l = CarrotListener(ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() id = gen_unique_id() c = create_message(backend, control={ "command": "revoke", "task_id": id }) t = create_message(backend, task=foo_task.name, args=[2, 4, 8], kwargs={}, id=id) l.event_dispatcher = MockEventDispatcher() l.receive_message(c.decode(), c) from celery.worker.state import revoked self.assertIn(id, revoked) l.receive_message(t.decode(), t) self.assertTrue(ready_queue.empty())
def test_receieve_message_eta(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, task=foo_task.name, args=[2, 4, 8], kwargs={}, eta=(datetime.now() + timedelta(days=1)).isoformat()) l.reset_connection() p, conf.BROKER_CONNECTION_RETRY = conf.BROKER_CONNECTION_RETRY, False try: l.reset_connection() finally: conf.BROKER_CONNECTION_RETRY = p l.receive_message(m.decode(), m) in_hold = self.eta_schedule.queue[0] self.assertEqual(len(in_hold), 3) eta, priority, entry = in_hold task = entry.args[0] self.assertIsInstance(task, TaskRequest) self.assertEqual(task.task_name, foo_task.name) self.assertEqual(task.execute(), 2 * 4 * 8) self.assertRaises(Empty, self.ready_queue.get_nowait)
def test_on_decode_error(self): logger = MockLogger() l = CarrotListener(self.ready_queue, self.eta_schedule, logger, send_events=False) class MockMessage(object): content_type = "application/x-msgpack" content_encoding = "binary" body = "foobarbaz" acked = False def ack(self): self.acked = True message = MockMessage() l.on_decode_error(message, KeyError("foo")) self.assertTrue(message.acked) self.assertIn("Message decoding error", logger.logged[0])
def __init__(self, concurrency=None, logfile=None, loglevel=None, send_events=conf.SEND_EVENTS, hostname=None, ready_callback=noop, embed_clockservice=False, schedule_filename=conf.CELERYBEAT_SCHEDULE_FILENAME): # Options self.loglevel = loglevel or self.loglevel self.concurrency = concurrency or self.concurrency self.logfile = logfile or self.logfile self.logger = setup_logger(loglevel, logfile) self.hostname = hostname or socket.gethostname() self.embed_clockservice = embed_clockservice self.ready_callback = ready_callback self.send_events = send_events self._finalize = Finalize(self, self.stop, exitpriority=20) # Queues if conf.DISABLE_RATE_LIMITS: self.ready_queue = Queue() else: self.ready_queue = TaskBucket(task_registry=registry.tasks) self.eta_schedule = Scheduler(self.ready_queue) self.logger.debug("Instantiating thread components...") # Threads + Pool + Consumer self.pool = TaskPool(self.concurrency, logger=self.logger, initializer=process_initializer) self.mediator = Mediator(self.ready_queue, callback=self.process_task, logger=self.logger) self.scheduler = ScheduleController(self.eta_schedule, logger=self.logger) self.clockservice = None if self.embed_clockservice: self.clockservice = EmbeddedClockService(logger=self.logger, schedule_filename=schedule_filename) prefetch_count = self.concurrency * conf.CELERYD_PREFETCH_MULTIPLIER self.listener = CarrotListener(self.ready_queue, self.eta_schedule, logger=self.logger, hostname=self.hostname, send_events=self.send_events, init_callback=self.ready_callback, initial_prefetch_count=prefetch_count) # The order is important here; # the first in the list is the first to start, # and they must be stopped in reverse order. self.components = filter(None, (self.pool, self.mediator, self.scheduler, self.clockservice, self.listener))
def test_close_connection(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) l._state = RUN l.close_connection() l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) eventer = l.event_dispatcher = MockEventDispatcher() heart = l.heart = MockHeart() l._state = RUN l.stop_consumers() self.assertTrue(eventer.closed) self.assertTrue(heart.closed)
def test_connection(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) l.reset_connection() self.assertTrue(isinstance(l.connection, BrokerConnection)) l.close_connection() self.assertTrue(l.connection is None) self.assertTrue(l.task_consumer is None) l.reset_connection() self.assertTrue(isinstance(l.connection, BrokerConnection)) l.stop() self.assertTrue(l.connection is None) self.assertTrue(l.task_consumer is None)
def test_receieve_message_eta_isoformat(self): class MockConsumer(object): prefetch_count_incremented = False def qos(self, **kwargs): self.prefetch_count_incremented = True l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() m = create_message(backend, task=foo_task.name, eta=datetime.now().isoformat(), args=[2, 4, 8], kwargs={}) l.event_dispatcher = MockEventDispatcher() l.task_consumer = MockConsumer() l.qos = QoS(l.task_consumer, l.initial_prefetch_count, l.logger) l.receive_message(m.decode(), m) items = [entry[2] for entry in self.eta_schedule.queue] found = 0 for item in items: if item.args[0].task_name == foo_task.name: found = True self.assertTrue(found) self.assertTrue(l.task_consumer.prefetch_count_incremented)
def test_mainloop(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) class MockConnection(object): def drain_events(self): return "draining" l.connection = MockConnection() l.connection.connection = MockConnection() it = l._mainloop() self.assertTrue(it.next(), "draining") records = {} def create_recorder(key): def _recorder(*args, **kwargs): records[key] = True return _recorder l.task_consumer = PlaceHolder() l.task_consumer.iterconsume = create_recorder("consume_tasks") l.broadcast_consumer = PlaceHolder() l.broadcast_consumer.register_callback = create_recorder( "broadcast_callback") l.broadcast_consumer.iterconsume = create_recorder("consume_broadcast") l.task_consumer.add_consumer = create_recorder("consumer_add") records.clear() self.assertEqual(l._detect_wait_method(), l._mainloop) for record in ("broadcast_callback", "consume_broadcast", "consume_tasks"): self.assertTrue(records.get(record)) records.clear() l.connection.connection = PlaceHolder() self.assertIs(l._detect_wait_method(), l.task_consumer.iterconsume) self.assertTrue(records.get("consumer_add"))
def test_revoke(self): ready_queue = FastQueue() l = CarrotListener(ready_queue, self.eta_schedule, self.logger, send_events=False) backend = MockBackend() id = gen_unique_id() c = create_message(backend, control={"command": "revoke", "task_id": id}) t = create_message(backend, task=foo_task.name, args=[2, 4, 8], kwargs={}, id=id) l.event_dispatcher = MockEventDispatcher() l.receive_message(c.decode(), c) from celery.worker.state import revoked self.assertIn(id, revoked) l.receive_message(t.decode(), t) self.assertTrue(ready_queue.empty())
def test_mainloop(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) class MockConnection(object): def drain_events(self): return "draining" l.connection = MockConnection() l.connection.connection = MockConnection() it = l._mainloop() self.assertTrue(it.next(), "draining") records = {} def create_recorder(key): def _recorder(*args, **kwargs): records[key] = True return _recorder l.task_consumer = PlaceHolder() l.task_consumer.iterconsume = create_recorder("consume_tasks") l.broadcast_consumer = PlaceHolder() l.broadcast_consumer.register_callback = create_recorder( "broadcast_callback") l.broadcast_consumer.iterconsume = create_recorder( "consume_broadcast") l.task_consumer.add_consumer = create_recorder("consumer_add") records.clear() self.assertEqual(l._detect_wait_method(), l._mainloop) for record in ("broadcast_callback", "consume_broadcast", "consume_tasks"): self.assertTrue(records.get(record)) records.clear() l.connection.connection = PlaceHolder() self.assertIs(l._detect_wait_method(), l.task_consumer.iterconsume) self.assertTrue(records.get("consumer_add"))
def test_connection(self): l = CarrotListener(self.ready_queue, self.eta_schedule, self.logger, send_events=False) l.reset_connection() self.assertIsInstance(l.connection, BrokerConnection) l.stop_consumers() self.assertIsNone(l.connection) self.assertIsNone(l.task_consumer) l.reset_connection() self.assertIsInstance(l.connection, BrokerConnection) l.stop() l.close_connection() self.assertIsNone(l.connection) self.assertIsNone(l.task_consumer)
class WorkController(object): """Executes tasks waiting in the task queue. :param concurrency: see :attr:`concurrency`. :param logfile: see :attr:`logfile`. :param loglevel: see :attr:`loglevel`. :param embed_clockservice: see :attr:`run_clockservice`. :param send_events: see :attr:`send_events`. .. attribute:: concurrency The number of simultaneous processes doing work (default: :const:`celery.conf.CELERYD_CONCURRENCY`) .. attribute:: loglevel The loglevel used (default: :const:`logging.INFO`) .. attribute:: logfile The logfile used, if no logfile is specified it uses ``stderr`` (default: :const:`celery.conf.CELERYD_LOG_FILE`). .. attribute:: embed_clockservice If ``True``, celerybeat is embedded, running in the main worker process as a thread. .. attribute:: send_events Enable the sending of monitoring events, these events can be captured by monitors (celerymon). .. attribute:: logger The :class:`logging.Logger` instance used for logging. .. attribute:: pool The :class:`multiprocessing.Pool` instance used. .. attribute:: ready_queue The :class:`Queue.Queue` that holds tasks ready for immediate processing. .. attribute:: hold_queue The :class:`Queue.Queue` that holds paused tasks. Reasons for holding back the task include waiting for ``eta`` to pass or the task is being retried. .. attribute:: schedule_controller Instance of :class:`celery.worker.controllers.ScheduleController`. .. attribute:: mediator Instance of :class:`celery.worker.controllers.Mediator`. .. attribute:: listener Instance of :class:`CarrotListener`. """ loglevel = logging.ERROR concurrency = conf.CELERYD_CONCURRENCY logfile = conf.CELERYD_LOG_FILE _state = None def __init__(self, concurrency=None, logfile=None, loglevel=None, send_events=conf.SEND_EVENTS, hostname=None, ready_callback=noop, embed_clockservice=False, schedule_filename=conf.CELERYBEAT_SCHEDULE_FILENAME): # Options self.loglevel = loglevel or self.loglevel self.concurrency = concurrency or self.concurrency self.logfile = logfile or self.logfile self.logger = setup_logger(loglevel, logfile) self.hostname = hostname or socket.gethostname() self.embed_clockservice = embed_clockservice self.ready_callback = ready_callback self.send_events = send_events self._finalize = Finalize(self, self.stop, exitpriority=20) # Queues if conf.DISABLE_RATE_LIMITS: self.ready_queue = Queue() else: self.ready_queue = TaskBucket(task_registry=registry.tasks) self.eta_schedule = Scheduler(self.ready_queue) self.logger.debug("Instantiating thread components...") # Threads + Pool + Consumer self.pool = TaskPool(self.concurrency, logger=self.logger, initializer=process_initializer) self.mediator = Mediator(self.ready_queue, callback=self.process_task, logger=self.logger) self.scheduler = ScheduleController(self.eta_schedule, logger=self.logger) self.clockservice = None if self.embed_clockservice: self.clockservice = EmbeddedClockService(logger=self.logger, schedule_filename=schedule_filename) prefetch_count = self.concurrency * conf.CELERYD_PREFETCH_MULTIPLIER self.listener = CarrotListener(self.ready_queue, self.eta_schedule, logger=self.logger, hostname=self.hostname, send_events=self.send_events, init_callback=self.ready_callback, initial_prefetch_count=prefetch_count) # The order is important here; # the first in the list is the first to start, # and they must be stopped in reverse order. self.components = filter(None, (self.pool, self.mediator, self.scheduler, self.clockservice, self.listener)) def start(self): """Starts the workers main loop.""" self._state = "RUN" try: for component in self.components: self.logger.debug("Starting thread %s..." % \ component.__class__.__name__) component.start() finally: self.stop() def process_task(self, wrapper): """Process task by sending it to the pool of workers.""" try: try: wrapper.task.execute(wrapper, self.pool, self.loglevel, self.logfile) except Exception, exc: self.logger.critical("Internal error %s: %s\n%s" % ( exc.__class__, exc, traceback.format_exc())) except (SystemExit, KeyboardInterrupt): self.stop() def stop(self): """Gracefully shutdown the worker server.""" if self._state != "RUN": return signals.worker_shutdown.send(sender=self) for component in reversed(self.components): self.logger.debug("Stopping thread %s..." % ( component.__class__.__name__)) component.stop() self.listener.close_connection() self._state = "STOP" def terminate(self): """Not so gracefully shutdown the worker server.""" if self._state != "RUN": return signals.worker_shutdown.send(sender=self) for component in reversed(self.components): self.logger.debug("Terminating thread %s..." % ( component.__class__.__name__)) terminate = getattr(component, "terminate", component.stop) terminate() self.listener.close_connection() self._state = "STOP"