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(name=WORKER_1, last_heartbeat=now)
worker_1.save()
worker_2 = Worker(name=WORKER_2, last_heartbeat=now)
worker_2.save()
# Set up two reserved resources
reserved_resource_1 = ReservedResource(task_id=str(uuid.uuid4()),
worker_name=worker_1.name, resource_id='resource_1')
reserved_resource_1.save()
reserved_resource_2 = ReservedResource(task_id=str(uuid.uuid4()),
worker_name=worker_2.name, resource_id='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
self.assertEqual(ReservedResource.objects.count(), 1)
rr_1 = ReservedResource.objects.get(task_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 _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 = get_worker_for_reservation(resource_id)
except NoWorkers:
pass
else:
break
try:
worker = _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=task_id,
worker_name=worker['name'],
resource_id=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_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(name=WORKER_1, last_heartbeat=datetime.utcnow())
worker_1.save()
worker_2 = Worker(name=WORKER_2, last_heartbeat=datetime.utcnow())
worker_2.save()
# Set up two resource reservations, using our workers from above
reserved_resource_1 = ReservedResource(task_id=str(uuid.uuid4()),
worker_name=worker_1.name,
resource_id='resource_1')
reserved_resource_1.save()
reserved_resource_2 = ReservedResource(task_id=str(uuid.uuid4()),
worker_name=worker_2.name,
resource_id='resource_2')
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
self.assertEqual(Worker.objects().count(), 2)
worker_1 = Worker.objects().get(name=worker_1.name)
self.assertTrue(worker_1)
worker_2 = Worker.objects().get(name=worker_2.name)
self.assertTrue(worker_2)
# Make sure that the reserved resources collection has not been altered
self.assertEqual(ReservedResource.objects.count(), 2)
rr_1 = ReservedResource.objects.get(
task_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')
rr_2 = ReservedResource.objects.get(
task_id=reserved_resource_2.task_id)
self.assertEqual(rr_2['worker_name'], reserved_resource_2.worker_name)
self.assertEqual(rr_2['resource_id'], 'resource_2')
def _queue_reserved_task_list(name, task_id, resource_id_list, inner_args,
inner_kwargs):
"""
A task that allows multiple resources to be reserved before dispatching a second, "inner", task.
See _queue_reserved_task for details on the inner workings.
: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_list: A list of names of the resources you wish to reserve for your task.
The system will ensure that no other tasks that want any of the same
reservations will run concurrently with yours.
:type resource_id_list: list
:return: None
"""
_logger.debug('_queue_reserved_task_list for task %s and ids [%s]' %
(task_id, resource_id_list))
# Find a/the available Worker for processing our list of resources
worker = get_worker_for_reservation_list(resource_id_list)
# Reserve each resource, associating them with that Worker
for rid in resource_id_list:
_logger.debug('...saving RR for RID %s' % rid)
ReservedResource(task_id=task_id,
worker_name=worker['name'],
resource_id=rid).save()
# Dispatch the Worker
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:
# Arrange to release all held reserved-resources
_release_resource.apply_async((task_id, ),
routing_key=worker.name,
exchange=DEDICATED_QUEUE_EXCHANGE)
