def test__release_resource_task_count_two(self):
"""
Test _release_resource() with a resource that has a task count of two. This should simply
decrement the task_count for the resource, but should not remove it from the database.
"""
# Set up two workers
now = datetime.utcnow()
worker_1 = Worker(WORKER_1, now)
worker_1.save()
worker_2 = Worker(WORKER_2, now)
worker_2.save()
# Set up two resource reservations, using our workers from above
reserved_resource_1 = ReservedResource('resource_1', worker_1.name, 7)
reserved_resource_1.save()
reserved_resource_2 = ReservedResource('resource_2', worker_2.name, 2)
reserved_resource_2.save()
# This should reduce the reserved_resource_2 num_reservations to 1.
tasks._release_resource('resource_2')
# Make sure the ReservedResources are also correct
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 2)
rr_1 = rrc.find_one({'_id': reserved_resource_1.name})
self.assertEqual(rr_1['assigned_queue'],
reserved_resource_1.assigned_queue)
self.assertEqual(rr_1['num_reservations'], 7)
rr_2 = rrc.find_one({'_id': reserved_resource_2.name})
self.assertEqual(rr_2['assigned_queue'],
reserved_resource_2.assigned_queue)
self.assertEqual(rr_2['num_reservations'], 1)
def test_get(self):
"""
Test for the case when the requested resource does exist.
"""
# Let's add two ReservedResources just to make sure that it doesn't return the wrong
# resource.
rr_1 = ReservedResource('resource_1')
rr_1.save()
rr_2 = ReservedResource('resource_2', 'some_queue', 7)
rr_2.save()
rr_2 = resources.get_or_create_reserved_resource('resource_2')
# Assert that the returned instance is correct
self.assertEqual(type(rr_2), ReservedResource)
self.assertEqual(rr_2.name, 'resource_2')
self.assertEqual(rr_2.assigned_queue, 'some_queue')
# The resource should have 7 reservations
self.assertEqual(rr_2.num_reservations, 7)
# Now we need to assert that the DB is still correct
rrc = rr_2.get_collection()
self.assertEqual(
rrc.find_one({'_id': 'resource_2'})['num_reservations'], 7)
self.assertEqual(
rrc.find_one({'_id': 'resource_2'})['assigned_queue'],
'some_queue')
def test__release_resource_task_count_one(self):
"""
Test _release_resource() with a resource that has a task count of one. This should remove
the resource from the database.
"""
# Set up two workers
now = datetime.utcnow()
worker_1 = Worker(WORKER_1, now)
worker_1.save()
worker_2 = Worker(WORKER_2, now)
worker_2.save()
# Set up two reserved resources
reserved_resource_1 = ReservedResource('resource_1', worker_1.name, 7)
reserved_resource_1.save()
reserved_resource_2 = ReservedResource('resource_2', worker_2.name, 1)
reserved_resource_2.save()
# This should remove resource_2 from the _resource_map.
tasks._release_resource('resource_2')
# resource_2 should have been removed from the database
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 1)
rr_1 = rrc.find_one({'_id': reserved_resource_1.name})
self.assertEqual(rr_1['assigned_queue'],
reserved_resource_1.assigned_queue)
self.assertEqual(rr_1['num_reservations'], 7)
def test_resource_in_resource_map(self):
"""
Test _release_resource() with a valid resource. This should remove the resource from the
database.
"""
# Set up two workers
now = datetime.utcnow()
worker_1 = Worker(WORKER_1, now)
worker_1.save()
worker_2 = Worker(WORKER_2, now)
worker_2.save()
# Set up two reserved resources
reserved_resource_1 = ReservedResource(uuid.uuid4(), worker_1.name,
'resource_1')
reserved_resource_1.save()
reserved_resource_2 = ReservedResource(uuid.uuid4(), worker_2.name,
'resource_2')
reserved_resource_2.save()
# This should remove resource_2 from the _resource_map.
tasks._release_resource(reserved_resource_2.task_id)
# resource_2 should have been removed from the database
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 1)
rr_1 = rrc.find_one({'_id': reserved_resource_1.task_id})
self.assertEqual(rr_1['worker_name'], reserved_resource_1.worker_name)
self.assertEqual(rr_1['resource_id'], 'resource_1')
def test__reserve_resource_with_existing_reservation(self):
"""
Test _reserve_resource() with a resource that has an existing reservation in the database.
It should return the queue listed in the database, and increment the reservation counter.
"""
# Set up a worker with a reservation count of 1
now = datetime.utcnow()
worker_1 = Worker(WORKER_1, now)
worker_1.save()
# Set up a resource reservation, using our worker from above
reserved_resource_1 = ReservedResource('resource_1',
worker_1.queue_name, 1)
reserved_resource_1.save()
# This should increase the reserved_resource_1 num_reservations to 2. worker_1's name should
# be returned
queue = tasks._reserve_resource('resource_1')
self.assertEqual(queue, WORKER_1_QUEUE)
# Make sure the ReservedResource is correct
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 1)
rr_1 = rrc.find_one({'_id': reserved_resource_1.name})
self.assertEqual(rr_1['assigned_queue'], WORKER_1_QUEUE)
self.assertEqual(rr_1['num_reservations'], 2)
def _queue_reserved_task(name, task_id, resource_id, inner_args, inner_kwargs):
"""
A task that encapsulates another task to be dispatched later. This task being encapsulated is
called the "inner" task, and a task name, UUID, and accepts a list of positional args
and keyword args for the inner task. These arguments are named inner_args and inner_kwargs.
inner_args is a list, and inner_kwargs is a dictionary passed to the inner task as positional
and keyword arguments using the * and ** operators.
The inner task is dispatched into a dedicated queue for a worker that is decided at dispatch
time. The logic deciding which queue receives a task is controlled through the
find_worker function.
:param name: The name of the task to be called
:type name: basestring
:param inner_task_id: The UUID to be set on the task being called. By providing
the UUID, the caller can have an asynchronous reference to the inner task
that will be dispatched.
:type inner_task_id: basestring
:param resource_id: The name of the resource you wish to reserve for your task. The system
will ensure that no other tasks that want that same reservation will run
concurrently with yours.
:type resource_id: basestring
:return: None
"""
while True:
try:
worker = resources.get_worker_for_reservation(resource_id)
except NoWorkers:
pass
else:
break
try:
worker = resources.get_unreserved_worker()
except NoWorkers:
pass
else:
break
# No worker is ready for this work, so we need to wait
time.sleep(0.25)
ReservedResource(task_id, worker['name'], resource_id).save()
inner_kwargs['routing_key'] = worker.name
inner_kwargs['exchange'] = DEDICATED_QUEUE_EXCHANGE
inner_kwargs['task_id'] = task_id
try:
celery.tasks[name].apply_async(*inner_args, **inner_kwargs)
finally:
_release_resource.apply_async((task_id, ),
routing_key=worker.name,
exchange=DEDICATED_QUEUE_EXCHANGE)
def test__release_resource_not_in__resource_map(self):
"""
Test _release_resource() with a resource that is not in the database. This should be
gracefully handled, and result in no changes to the database.
"""
# Set up two workers
worker_1 = Worker(WORKER_1, datetime.utcnow())
worker_1.save()
worker_2 = Worker(WORKER_2, datetime.utcnow())
worker_2.save()
# Set up two resource reservations, using our workers from above
reserved_resource_1 = ReservedResource('resource_1', worker_1.name, 7)
reserved_resource_1.save()
reserved_resource_2 = ReservedResource('resource_2', worker_2.name, 3)
reserved_resource_2.save()
# This should not raise any Exception, but should also not alter either the Worker
# collection or the ReservedResource collection
tasks._release_resource('made_up_resource_id')
# Make sure that the workers collection has not been altered
worker_collection = Worker.get_collection()
self.assertEqual(worker_collection.count(), 2)
worker_1 = worker_collection.find_one({'_id': worker_1.name})
self.assertTrue(worker_1)
worker_2 = worker_collection.find_one({'_id': worker_2.name})
self.assertTrue(worker_2)
# Make sure that the reserved resources collection has not been altered
rrc = ReservedResource.get_collection()
self.assertEqual(rrc.count(), 2)
rr_1 = rrc.find_one({'_id': reserved_resource_1.name})
self.assertEqual(rr_1['assigned_queue'],
reserved_resource_1.assigned_queue)
self.assertEqual(rr_1['num_reservations'], 7)
rr_2 = rrc.find_one({'_id': reserved_resource_2.name})
self.assertEqual(rr_2['assigned_queue'],
reserved_resource_2.assigned_queue)
self.assertEqual(rr_2['num_reservations'], 3)
def test_ignores_queues_that_arent_workers(self):
"""
It is possible for the assigned_queue in a ReservedResource to reference a queue that is not
in the workers collection. This test ensures that this queue is properly ignored, even if it
is the most "enticing" choice.
"""
# Set up three Workers, with the least busy one in the middle so that we can
# demonstrate that it did pick the least busy and not the last or first.
now = datetime.utcnow()
worker_1 = Worker('busy_worker', now)
worker_2 = Worker('less_busy_worker', now)
worker_3 = Worker('most_busy_worker', now)
for worker in (worker_1, worker_2, worker_3):
worker.save()
# Now we need to make some reservations against these Workers' queues. We'll give worker_1
# 8 reservations, putting it in the middle of busyness.
rr_1 = ReservedResource(name='resource_1',
assigned_queue=worker_1.queue_name,
num_reservations=8)
# These next two will give worker_2 a total of 7 reservations, so it should get picked.
rr_2 = ReservedResource(name='resource_2',
assigned_queue=worker_2.queue_name,
num_reservations=3)
rr_3 = ReservedResource(name='resource_3',
assigned_queue=worker_2.queue_name,
num_reservations=4)
# These next three will give worker_3 a total of 9 reservations, so it should be the most
# busy.
rr_4 = ReservedResource(name='resource_4',
assigned_queue=worker_3.queue_name,
num_reservations=2)
rr_5 = ReservedResource(name='resource_5',
assigned_queue=worker_3.queue_name,
num_reservations=3)
rr_6 = ReservedResource(name='resource_6',
assigned_queue=worker_3.queue_name,
num_reservations=4)
# Now we will make a ReservedResource that references a queue that does not correspond to a
# Worker and has the lowest reservation count. This RR should be ignored.
rr_7 = ReservedResource(name='resource_7',
assigned_queue='doesnt_exist',
num_reservations=1)
for rr in (rr_1, rr_2, rr_3, rr_4, rr_5, rr_6, rr_7):
rr.save()
worker = resources.get_least_busy_worker()
self.assertEqual(type(worker), Worker)
self.assertEqual(worker.name, 'less_busy_worker')
def _release_resource(resource_id):
"""
Do not queue this task yourself, but always use the _queue_release_resource() task instead.
Please see the docblock on that function for an explanation.
When a resource-reserving task is complete, this method must be called with the
resource_id so that the we know when it is safe to unmap a resource_id from
its given queue name.
:param resource_id: The resource that is no longer in use
:type resource_id: basestring
"""
try:
reserved_resource = ReservedResource(resource_id)
reserved_resource.decrement_num_reservations()
except DoesNotExist:
# If we are trying to decrement the count on one of these objects, and they don't exist,
# that's OK
pass
def test_create(self):
"""
Test for the case when the requested resource does not exist.
"""
# Let's add an ReservedResource just to make sure that it doesn't return any existing
# resource.
rr_1 = ReservedResource('resource_1')
rr_1.save()
rr_2 = resources.get_or_create_reserved_resource('resource_2')
# Assert that the returned instance is correct
self.assertEqual(type(rr_2), ReservedResource)
self.assertEqual(rr_2.name, 'resource_2')
# By default, the assigned_queue should be set to None
self.assertEqual(rr_2.assigned_queue, None)
# A new resource should default to 1 reservations
self.assertEqual(rr_2.num_reservations, 1)
# Now we need to assert that it made it to the database as well
rrc = rr_2.get_collection()
self.assertEqual(
rrc.find_one({'_id': 'resource_2'})['num_reservations'], 1)
self.assertEqual(
rrc.find_one({'_id': 'resource_2'})['assigned_queue'], None)
def test_picks_least_busy_worker(self):
"""
Test that the function picks the least busy worker.
"""
# Set up three Workers, with the least busy one in the middle so that we can
# demonstrate that it did pick the least busy and not the last or first.
now = datetime.utcnow()
worker_1 = Worker('busy_worker', now)
worker_2 = Worker('less_busy_worker', now)
worker_3 = Worker('most_busy_worker', now)
for worker in (worker_1, worker_2, worker_3):
worker.save()
# Now we need to make some reservations against these Workers' queues. We'll give worker_1
# 8 reservations, putting it in the middle of busyness.
rr_1 = ReservedResource(name='resource_1',
assigned_queue=worker_1.queue_name,
num_reservations=8)
# These next two will give worker_2 a total of 7 reservations, so it should get picked.
rr_2 = ReservedResource(name='resource_2',
assigned_queue=worker_2.queue_name,
num_reservations=3)
rr_3 = ReservedResource(name='resource_3',
assigned_queue=worker_2.queue_name,
num_reservations=4)
# These next three will give worker_3 a total of 9 reservations, so it should be the most
# busy.
rr_4 = ReservedResource(name='resource_4',
assigned_queue=worker_3.queue_name,
num_reservations=2)
rr_5 = ReservedResource(name='resource_5',
assigned_queue=worker_3.queue_name,
num_reservations=3)
rr_6 = ReservedResource(name='resource_6',
assigned_queue=worker_3.queue_name,
num_reservations=4)
for rr in (rr_1, rr_2, rr_3, rr_4, rr_5, rr_6):
rr.save()
worker = resources.get_least_busy_worker()
self.assertEqual(type(worker), Worker)
self.assertEqual(worker.name, 'less_busy_worker')
