def start(self, c):
c.update_strategies()
c.task_consumer = c.app.amqp.TaskConsumer(
c.connection,
on_decode_error=c.on_decode_error,
)
c.qos = QoS(c.task_consumer.qos, self.initial_prefetch_count)
c.qos.update() # set initial prefetch count
def test_qos_exceeds_16bit(self):
with patch('kombu.common.logger') as logger:
callback = Mock()
qos = QoS(callback, 10)
qos.prev = 100
qos.set(2**32)
self.assertTrue(logger.warn.called)
callback.assert_called_with(prefetch_count=0)
def test_consumer_decrement_eventually(self):
mconsumer = Mock()
qos = QoS(mconsumer.qos, 10)
qos.decrement_eventually()
assert qos.value == 9
qos.value = 0
qos.decrement_eventually()
assert qos.value == 0
def test_consumer_decrement_eventually(self):
mconsumer = Mock()
qos = QoS(mconsumer.qos, 10)
qos.decrement_eventually()
self.assertEqual(qos.value, 9)
qos.value = 0
qos.decrement_eventually()
self.assertEqual(qos.value, 0)
def test_apply_eta_task(self):
c = self.NoopConsumer()
c.qos = QoS(None, 10)
task = Mock(name='task', id='1234213')
qos = c.qos.value
c.apply_eta_task(task)
self.assertIn(task, state.reserved_requests)
self.assertEqual(c.qos.value, qos - 1)
self.assertIs(self.buffer.get_nowait(), task)
def test_apply_eta_task(self):
c = self.NoopConsumer()
c.qos = QoS(None, 10)
task = Mock(name='task', id='1234213')
qos = c.qos.value
c.apply_eta_task(task)
assert task in state.reserved_requests
assert c.qos.value == qos - 1
assert self.buffer.get_nowait() is task
def test_info(self):
l = MyKombuConsumer(self.ready_queue, timer=self.timer)
l.qos = QoS(l.task_consumer, 10)
info = l.info
self.assertEqual(info['prefetch_count'], 10)
self.assertFalse(info['broker'])
l.connection = current_app.connection()
info = l.info
self.assertTrue(info['broker'])
def test_qos_exceeds_16bit(self):
with patch('kombu.common.logger') as logger:
callback = Mock()
qos = QoS(callback, 10)
qos.prev = 100
# cannot use 2 ** 32 because of a bug on macOS Py2.5:
# https://jira.mongodb.org/browse/PYTHON-389
qos.set(4294967296)
logger.warning.assert_called()
callback.assert_called_with(prefetch_count=0)
def test_with_file_descriptor_safety(self):
# Given: a test celery worker instance
worker = self.create_worker(
autoscale=[10, 5],
use_eventloop=True,
timer_cls='celery.utils.timer2.Timer',
threads=False,
)
# Given: This test requires a QoS defined on the worker consumer
worker.consumer.qos = qos = QoS(lambda prefetch_count: prefetch_count,
2)
qos.update()
# Given: We have started the worker pool
worker.pool.start()
# Given: Utilize kombu to get the global hub state
hub = get_event_loop()
# Given: Initial call the Async Pool to register events works fine
worker.pool.register_with_event_loop(hub)
# Given: Mock the Hub to return errors for add and remove
def throw_file_not_found_error(*args, **kwargs):
raise OSError()
hub.add = throw_file_not_found_error
hub.add_reader = throw_file_not_found_error
hub.remove = throw_file_not_found_error
# When: Calling again to register with event loop ...
worker.pool.register_with_event_loop(hub)
worker.pool._pool.register_with_event_loop(hub)
# Then: test did not raise OSError
# Note: worker.pool is prefork.TaskPool whereas
# worker.pool._pool is the asynpool.AsynPool class.
# When: Calling the tic method on_poll_start
worker.pool._pool.on_poll_start()
# Then: test did not raise OSError
# Given: a mock object that fakes whats required to do whats next
proc = Mock(_sentinel_poll=42)
# When: Calling again to register with event loop ...
worker.pool._pool._track_child_process(proc, hub)
# Then: test did not raise OSError
# Given:
worker.pool._pool._flush_outqueue = throw_file_not_found_error
# Finally: Clean up so the threads before/after fixture passes
worker.terminate()
worker.pool.terminate()
def test_apply_eta_task(self):
from celery.worker import state
l = MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)
l.qos = QoS(None, 10)
task = object()
qos = l.qos.value
l.apply_eta_task(task)
self.assertIn(task, state.reserved_requests)
self.assertEqual(l.qos.value, qos - 1)
self.assertIs(self.buffer.get_nowait(), task)
def test_exceeds_short(self):
qos = QoS(Mock(), PREFETCH_COUNT_MAX - 1)
qos.update()
self.assertEqual(qos.value, PREFETCH_COUNT_MAX - 1)
qos.increment_eventually()
self.assertEqual(qos.value, PREFETCH_COUNT_MAX)
qos.increment_eventually()
self.assertEqual(qos.value, PREFETCH_COUNT_MAX + 1)
qos.decrement_eventually()
self.assertEqual(qos.value, PREFETCH_COUNT_MAX)
qos.decrement_eventually()
self.assertEqual(qos.value, PREFETCH_COUNT_MAX - 1)
def test_exceeds_short(self):
qos = QoS(Mock(), PREFETCH_COUNT_MAX - 1)
qos.update()
assert qos.value == PREFETCH_COUNT_MAX - 1
qos.increment_eventually()
assert qos.value == PREFETCH_COUNT_MAX
qos.increment_eventually()
assert qos.value == PREFETCH_COUNT_MAX + 1
qos.decrement_eventually()
assert qos.value == PREFETCH_COUNT_MAX
qos.decrement_eventually()
assert qos.value == PREFETCH_COUNT_MAX - 1
def test_info(self):
l = MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)
l.task_consumer = Mock()
l.qos = QoS(l.task_consumer.qos, 10)
l.connection = Mock()
l.connection.info.return_value = {'foo': 'bar'}
l.controller = l.app.WorkController()
l.controller.pool = Mock()
l.controller.pool.info.return_value = [Mock(), Mock()]
l.controller.consumer = l
info = l.controller.stats()
self.assertEqual(info['prefetch_count'], 10)
self.assertTrue(info['broker'])
def test_loop_ignores_socket_timeout(self):
class Connection(self.app.connection_for_read().__class__):
obj = None
def drain_events(self, **kwargs):
self.obj.connection = None
raise socket.timeout(10)
c = self.NoopConsumer()
c.connection = Connection()
c.connection.obj = c
c.qos = QoS(c.task_consumer.qos, 10)
c.loop(*c.loop_args())
def test_loop_ignores_socket_timeout(self):
class Connection(self.app.connection().__class__):
obj = None
def drain_events(self, **kwargs):
self.obj.connection = None
raise socket.timeout(10)
l = MyKombuConsumer(self.buffer.put, timer=self.timer, app=self.app)
l.connection = Connection()
l.task_consumer = Mock()
l.connection.obj = l
l.qos = QoS(l.task_consumer.qos, 10)
l.loop(*l.loop_args())
def test_loop_ignores_socket_timeout(self):
class Connection(current_app.connection().__class__):
obj = None
def drain_events(self, **kwargs):
self.obj.connection = None
raise socket.timeout(10)
l = MyKombuConsumer(self.ready_queue, timer=self.timer)
l.connection = Connection()
l.task_consumer = Mock()
l.connection.obj = l
l.qos = QoS(l.task_consumer, 10)
l.loop(*l.loop_args())
def test_with_autoscaler_file_descriptor_safety(self):
# Given: a test celery worker instance with auto scaling
worker = self.create_worker(
autoscale=[10, 5],
use_eventloop=True,
timer_cls='celery.utils.timer2.Timer',
threads=False,
)
# Given: This test requires a QoS defined on the worker consumer
worker.consumer.qos = qos = QoS(lambda prefetch_count: prefetch_count,
2)
qos.update()
# Given: We have started the worker pool
worker.pool.start()
# Then: the worker pool is the same as the autoscaler pool
auto_scaler = worker.autoscaler
assert worker.pool == auto_scaler.pool
# Given: Utilize kombu to get the global hub state
hub = get_event_loop()
# Given: Initial call the Async Pool to register events works fine
worker.pool.register_with_event_loop(hub)
# Create some mock queue message and read from them
_keep = [Mock(name=f'req{i}') for i in range(20)]
[state.task_reserved(m) for m in _keep]
auto_scaler.body()
# Simulate a file descriptor from the list is closed by the OS
# auto_scaler.force_scale_down(5)
# This actually works -- it releases the semaphore properly
# Same with calling .terminate() on the process directly
for fd, proc in worker.pool._pool._fileno_to_outq.items():
# however opening this fd as a file and closing it will do it
queue_worker_socket = open(str(fd), "w")
queue_worker_socket.close()
break # Only need to do this once
# When: Calling again to register with event loop ...
worker.pool.register_with_event_loop(hub)
# Then: test did not raise "OSError: [Errno 9] Bad file descriptor!"
# Finally: Clean up so the threads before/after fixture passes
worker.terminate()
worker.pool.terminate()
def test_loop_when_socket_error(self):
class Connection(self.app.connection_for_read().__class__):
obj = None
def drain_events(self, **kwargs):
self.obj.connection = None
raise socket.error('foo')
c = self.LoopConsumer()
c.blueprint.state = RUN
conn = c.connection = Connection()
c.connection.obj = c
c.qos = QoS(c.task_consumer.qos, 10)
with pytest.raises(socket.error):
c.loop(*c.loop_args())
c.blueprint.state = CLOSE
c.connection = conn
c.loop(*c.loop_args())
def test_loop_when_socket_error(self):
class Connection(current_app.connection().__class__):
obj = None
def drain_events(self, **kwargs):
self.obj.connection = None
raise socket.error('foo')
l = Consumer(self.ready_queue, timer=self.timer)
l.namespace.state = RUN
c = l.connection = Connection()
l.connection.obj = l
l.task_consumer = Mock()
l.qos = QoS(l.task_consumer.qos, 10)
with self.assertRaises(socket.error):
l.loop(*l.loop_args())
l.namespace.state = CLOSE
l.connection = c
l.loop(*l.loop_args())
def test_loop_when_socket_error(self):
class Connection(self.app.connection().__class__):
obj = None
def drain_events(self, **kwargs):
self.obj.connection = None
raise socket.error('foo')
l = Consumer(self.buffer.put, timer=self.timer, app=self.app)
l.blueprint.state = RUN
c = l.connection = Connection()
l.connection.obj = l
l.task_consumer = Mock()
l.qos = QoS(l.task_consumer.qos, 10)
with self.assertRaises(socket.error):
l.loop(*l.loop_args())
l.blueprint.state = CLOSE
l.connection = c
l.loop(*l.loop_args())
def reset_connection(self):
"""Re-establish the broker connection and set up consumers,
heartbeat and the event dispatcher."""
debug('Re-establishing connection to the broker...')
self.stop_consumers(join=False)
# Clear internal queues to get rid of old messages.
# They can't be acked anyway, as a delivery tag is specific
# to the current channel.
self.ready_queue.clear()
self.timer.clear()
# Re-establish the broker connection and setup the task consumer.
self.connection = self._open_connection()
info('consumer: Connected to %s.', self.connection.as_uri())
self.task_consumer = self.app.amqp.TaskConsumer(self.connection,
on_decode_error=self.on_decode_error)
# QoS: Reset prefetch window.
self.qos = QoS(self.task_consumer, self.initial_prefetch_count)
self.qos.update()
# Setup the process mailbox.
self.reset_pidbox_node()
# Flush events sent while connection was down.
prev_event_dispatcher = self.event_dispatcher
self.event_dispatcher = self.app.events.Dispatcher(self.connection,
hostname=self.hostname,
enabled=self.send_events)
if prev_event_dispatcher:
self.event_dispatcher.copy_buffer(prev_event_dispatcher)
self.event_dispatcher.flush()
# Restart heartbeat thread.
self.restart_heartbeat()
# reload all task's execution strategies.
self.update_strategies()
# We're back!
self._state = RUN
def test_consumer_increment_decrement(self):
mconsumer = Mock()
qos = QoS(mconsumer.qos, 10)
qos.update()
self.assertEqual(qos.value, 10)
mconsumer.qos.assert_called_with(prefetch_count=10)
qos.decrement_eventually()
qos.update()
self.assertEqual(qos.value, 9)
mconsumer.qos.assert_called_with(prefetch_count=9)
qos.decrement_eventually()
self.assertEqual(qos.value, 8)
mconsumer.qos.assert_called_with(prefetch_count=9)
self.assertIn({'prefetch_count': 9}, mconsumer.qos.call_args)
# Does not decrement 0 value
qos.value = 0
qos.decrement_eventually()
self.assertEqual(qos.value, 0)
qos.increment_eventually()
self.assertEqual(qos.value, 0)
def test_consumer_increment_decrement(self):
mconsumer = Mock()
qos = QoS(mconsumer.qos, 10)
qos.update()
assert qos.value == 10
mconsumer.qos.assert_called_with(prefetch_count=10)
qos.decrement_eventually()
qos.update()
assert qos.value == 9
mconsumer.qos.assert_called_with(prefetch_count=9)
qos.decrement_eventually()
assert qos.value == 8
mconsumer.qos.assert_called_with(prefetch_count=9)
assert {'prefetch_count': 9} in mconsumer.qos.call_args
# Does not decrement 0 value
qos.value = 0
qos.decrement_eventually()
assert qos.value == 0
qos.increment_eventually()
assert qos.value == 0
def test_loop(self):
class Connection(self.app.connection_for_read().__class__):
obj = None
def drain_events(self, **kwargs):
self.obj.connection = None
c = self.LoopConsumer()
c.blueprint.state = RUN
c.connection = Connection()
c.connection.obj = c
c.qos = QoS(c.task_consumer.qos, 10)
c.loop(*c.loop_args())
c.loop(*c.loop_args())
self.assertTrue(c.task_consumer.consume.call_count)
c.task_consumer.qos.assert_called_with(prefetch_count=10)
self.assertEqual(c.qos.value, 10)
c.qos.decrement_eventually()
self.assertEqual(c.qos.value, 9)
c.qos.update()
self.assertEqual(c.qos.value, 9)
c.task_consumer.qos.assert_called_with(prefetch_count=9)
def test_loop(self):
class Connection(current_app.connection().__class__):
obj = None
def drain_events(self, **kwargs):
self.obj.connection = None
l = Consumer(self.buffer.put, timer=self.timer)
l.connection = Connection()
l.connection.obj = l
l.task_consumer = Mock()
l.qos = QoS(l.task_consumer.qos, 10)
l.loop(*l.loop_args())
l.loop(*l.loop_args())
self.assertTrue(l.task_consumer.consume.call_count)
l.task_consumer.qos.assert_called_with(prefetch_count=10)
self.assertEqual(l.qos.value, 10)
l.qos.decrement_eventually()
self.assertEqual(l.qos.value, 9)
l.qos.update()
self.assertEqual(l.qos.value, 9)
l.task_consumer.qos.assert_called_with(prefetch_count=9)
def start(self, c):
c.update_strategies()
# - RabbitMQ 3.3 completely redefines how basic_qos works..
# This will detect if the new qos smenatics is in effect,
# and if so make sure the 'apply_global' flag is set on qos updates.
qos_global = not c.connection.qos_semantics_matches_spec
# set initial prefetch count
c.connection.default_channel.basic_qos(
0, c.initial_prefetch_count, qos_global,
)
c.task_consumer = c.app.amqp.TaskConsumer(
c.connection, on_decode_error=c.on_decode_error,
)
def set_prefetch_count(prefetch_count):
return c.task_consumer.qos(
prefetch_count=prefetch_count,
apply_global=qos_global,
)
c.qos = QoS(set_prefetch_count, c.initial_prefetch_count)
def test_receieve_message_eta_isoformat(self):
l = MyKombuConsumer(self.ready_queue, timer=self.timer)
m = create_message(Mock(), task=foo_task.name,
eta=datetime.now().isoformat(),
args=[2, 4, 8], kwargs={})
l.task_consumer = Mock()
l.qos = QoS(l.task_consumer, l.initial_prefetch_count)
current_pcount = l.qos.value
l.event_dispatcher = Mock()
l.enabled = False
l.update_strategies()
l.receive_message(m.decode(), m)
l.timer.stop()
l.timer.join(1)
items = [entry[2] for entry in self.timer.queue]
found = 0
for item in items:
if item.args[0].name == foo_task.name:
found = True
self.assertTrue(found)
self.assertGreater(l.qos.value, current_pcount)
l.timer.stop()
def test_set(self):
mconsumer = Mock()
qos = QoS(mconsumer.qos, 10)
qos.set(12)
self.assertEqual(qos.prev, 12)
qos.set(qos.prev)
def test_set(self):
mconsumer = Mock()
qos = QoS(mconsumer.qos, 10)
qos.set(12)
assert qos.prev == 12
qos.set(qos.prev)
