def handle_worker_heartbeat(event):
"""
Celery event handler for 'worker-heartbeat' events.
The event is first parsed and logged. Then the existing Worker objects are
searched for one to update. If an existing one is found, it is updated.
Otherwise a new Worker entry is created. Logging at the info and debug
level is also done.
:param event: A celery event to handle.
:type event: dict
"""
event_info = _parse_and_log_event(event)
find_worker_criteria = Criteria(filters={'_id': event_info['worker_name']},
fields=('_id', 'last_heartbeat'))
find_worker_list = list(resources.filter_workers(find_worker_criteria))
if find_worker_list:
Worker.get_collection().find_and_modify(
query={'_id': event_info['worker_name']},
update={'$set': {'last_heartbeat': event_info['timestamp']}}
)
else:
new_worker = Worker(event_info['worker_name'], event_info['timestamp'])
msg = _("New worker '%(worker_name)s' discovered") % event_info
_logger.info(msg)
new_worker.save()
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 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__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__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_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 test_post(self, mock_publish, mock_get_worker_for_reservation):
"""
Test that publish repo group creates a task for a worker.
"""
# Setup
group_id = 'group-1'
distributor_id = 'dist-1'
mock_publish.apply_async_with_reservation.return_value = AsyncResult(
'1234')
self.manager.create_repo_group(group_id)
self.distributor_manager.add_distributor(group_id,
'dummy-group-distributor', {},
distributor_id=distributor_id)
mock_get_worker_for_reservation.return_value = Worker(
'some_queue', datetime.datetime.now())
# Test
data = {'id': distributor_id}
status, body = self.post(
'/v2/repo_groups/%s/actions/publish/' % group_id, data)
# Verify
self.assertEqual(202, status)
self.assertEqual(mock_publish.apply_async_with_reservation.call_count,
1)
def test_get_unreserved_worker_breaks_out_of_loop(self):
self.mock_get_worker_for_reservation.side_effect = NoWorkers()
self.mock_get_unreserved_worker.return_value = Worker(
'worker1', datetime.utcnow())
tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id',
[1, 2], {'a': 2})
self.assertTrue(not self.mock_time.sleep.called)
def test_update_repo_and_plugins(self, distributor_update,
mock_get_worker_for_reservation):
"""
Tests the aggregate call to update a repo and its plugins.
"""
mock_get_worker_for_reservation.return_value = Worker(
'some_queue', datetime.datetime.now())
self.manager.create_repo('repo-1', 'Original', 'Original Description')
importer_manager = manager_factory.repo_importer_manager()
distributor_manager = manager_factory.repo_distributor_manager()
importer_manager.set_importer('repo-1', 'mock-importer',
{'key-i1': 'orig-1'})
distributor_manager.add_distributor('repo-1',
'mock-distributor',
{'key-d1': 'orig-1'},
True,
distributor_id='dist-1')
distributor_manager.add_distributor('repo-1',
'mock-distributor',
{'key-d2': 'orig-2'},
True,
distributor_id='dist-2')
# Test
repo_delta = {'display_name': 'Updated'}
new_importer_config = {'key-i1': 'updated-1', 'key-i2': 'new-1'}
new_distributor_configs = {
'dist-1': {
'key-d1': 'updated-1'
},
} # only update one of the two distributors
result = self.manager.update_repo_and_plugins('repo-1', repo_delta,
new_importer_config,
new_distributor_configs)
self.assertTrue(isinstance(result, TaskResult))
self.assertEquals(None, result.error)
repo = result.return_value
# Verify
self.assertEqual(repo['id'], 'repo-1')
self.assertEqual(repo['display_name'], 'Updated')
self.assertEqual(repo['description'], 'Original Description')
importer = importer_manager.get_importer('repo-1')
self.assertEqual(importer['config'], new_importer_config)
dist_1 = distributor_manager.get_distributor('repo-1', 'dist-1')
self.assertEqual(dist_1['config'], new_distributor_configs['dist-1'])
dist_2 = distributor_manager.get_distributor('repo-1', 'dist-2')
self.assertEqual(dist_2['config'], {'key-d2': 'orig-2'})
# There should have been a spawned task for the new distributor config
expected_task_id = dispatch.TaskStatus.get_collection().find_one(
{'tags': 'pulp:repository_distributor:dist-1'})['task_id']
self.assertEqual(result.spawned_tasks, [{'task_id': expected_task_id}])
def test_post_with_override_config(self, mock_get_worker_for_reservation, mock_uuid,
mock_apply_async):
# Setup
uuid_list = [uuid.uuid4() for i in range(10)]
mock_uuid.uuid4.side_effect = copy.deepcopy(uuid_list)
expected_async_result = AsyncResult(str(uuid_list[0]))
mock_get_worker_for_reservation.return_value = Worker('some_queue', datetime.datetime.now())
upload_id = self.upload_manager.initialize_upload()
self.upload_manager.save_data(upload_id, 0, 'string data')
repo_manager = manager_factory.repo_manager()
repo_manager.create_repo('repo-upload')
importer_manager = manager_factory.repo_importer_manager()
importer_manager.set_importer('repo-upload', 'dummy-importer', {})
# Test
test_override_config = {'key1': 'value1', 'key2': 'value2'}
body = {
'upload_id' : upload_id,
'unit_type_id' : 'dummy-type',
'unit_key' : {'name' : 'foo'},
'unit_metadata' : {'stuff' : 'bar'},
'override_config': test_override_config,
}
status, body = self.post('/v2/repositories/repo-upload/actions/import_upload/', body)
# Verify
self.assertEqual(202, status)
assert_body_matches_async_task(body, expected_async_result)
exepcted_call_args = ['repo-upload', 'dummy-type',
{'name': 'foo'}, {'stuff': 'bar'},
upload_id, test_override_config]
self.assertEqual(exepcted_call_args, mock_apply_async.call_args[0][0])
def test_dispatches__release_resource(self):
self.mock_get_worker_for_reservation.return_value = Worker(
'worker1', datetime.utcnow())
tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id',
[1, 2], {'a': 2})
self.mock__release_resource.apply_async.assert_called_once_with(
('my_task_id', ), routing_key='worker1', exchange='C.dq')
def test_creates_and_saves_reserved_resource(self):
self.mock_get_worker_for_reservation.return_value = Worker(
'worker1', datetime.utcnow())
tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id',
[1, 2], {'a': 2})
self.mock_reserved_resource.assert_called_once_with(
'my_task_id', 'worker1', 'my_resource_id')
self.mock_reserved_resource.return_value.save.assert_called_once_with()
def test__delete_worker(self, logger, cancel, mock_add_consumer):
"""
Assert that the correct Tasks get canceled when their Worker is deleted, and that the Worker
is removed from the database.
"""
# cause two workers to be added to the database as having workers
worker_watcher.handle_worker_heartbeat({
'timestamp': time.time(),
'type': 'worker-heartbeat',
'hostname': WORKER_1,
})
worker_watcher.handle_worker_heartbeat({
'timestamp': time.time(),
'type': 'worker-heartbeat',
'hostname': WORKER_2,
})
# Let's simulate three tasks being assigned to WORKER_2, with two of them being
# in an incomplete state and one in a complete state. We will delete WORKER_2,
# which should cause the two to get canceled. Let's put task_1 in progress
TaskStatusManager.create_task_status('task_1',
WORKER_2_QUEUE,
state=CALL_RUNNING_STATE)
TaskStatusManager.create_task_status('task_2',
WORKER_2_QUEUE,
state=CALL_WAITING_STATE)
# This task shouldn't get canceled because it isn't in an incomplete state
TaskStatusManager.create_task_status('task_3',
WORKER_2_QUEUE,
state=CALL_FINISHED_STATE)
# Let's make a task in a worker that is still present just to make sure it isn't touched.
TaskStatusManager.create_task_status('task_4',
WORKER_1_QUEUE,
state=CALL_RUNNING_STATE)
# Let's just make sure the setup worked and that we have a Worker with RR2
worker_collection = Worker.get_collection()
self.assertEqual(worker_collection.find({'_id': WORKER_2}).count(), 1)
# Now let's delete the Worker named WORKER_2
tasks._delete_worker.apply_async(args=(WORKER_2, ),
queue=tasks.RESOURCE_MANAGER_QUEUE)
# cancel() should have been called twice with task_1 and task_2 as parameters
self.assertEqual(cancel.call_count, 2)
# Let's build a set out of the two times that cancel was called. We can't know for sure
# which order the Tasks got canceled in, but we can assert that the correct two tasks were
# canceled (task_3 should not appear in this set).
cancel_param_set = set([c[1] for c in cancel.mock_calls])
self.assertEqual(cancel_param_set, set([('task_1', ), ('task_2', )]))
# We should have logged that we are canceling the tasks
self.assertEqual(logger.call_count, 0)
self.assertTrue(WORKER_2 in logger.mock_calls[0][1][0])
self.assertTrue('Canceling the tasks' in logger.mock_calls[0][1][0])
# The Worker should have been deleted
self.assertEqual(worker_collection.find({'_id': WORKER_2}).count(), 0)
# the Worker for RW1 should remain
self.assertEqual(worker_collection.find({'_id': WORKER_1}).count(), 1)
def test__reserve_resource_without_existing_reservation(self):
"""
Test _reserve_resource() with a resource that does not have an existing reservation in the
database. It should find the least busy worker, add a reservation to the database with that
worker's queue, and then return the queue name.
"""
# Set up a worker
worker_1 = Worker(WORKER_1, datetime.utcnow())
worker_1.save()
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': 'resource_1'})
self.assertEqual(rr_1['assigned_queue'], WORKER_1_QUEUE)
self.assertEqual(rr_1['num_reservations'], 1)
def test__reserve_resource_without_existing_reservation(self):
"""
Test _reserve_resource() with a resource that does not have an existing reservation in the
database. It should find the least busy worker, add a reservation to the database with that
worker's queue, and then return the queue name.
"""
# Set up a worker
worker_1 = Worker(WORKER_1, datetime.utcnow())
worker_1.save()
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': 'resource_1'})
self.assertEqual(rr_1['assigned_queue'], WORKER_1_QUEUE)
self.assertEqual(rr_1['num_reservations'], 1)
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')
def test_post(self, mock_get_worker_for_reservation):
# Setup
group_id = 'group-1'
distributor_id = 'dist-1'
self.manager.create_repo_group(group_id)
self.distributor_manager.add_distributor(group_id, 'dummy-group-distributor', {}, distributor_id=distributor_id)
mock_get_worker_for_reservation.return_value = Worker('some_queue', datetime.datetime.now())
# Test
data = {'id' : distributor_id}
status, body = self.post('/v2/repo_groups/%s/actions/publish/' % group_id, data)
def GET(self, task_id):
try:
task = TaskStatus.objects.get(task_id=task_id)
except DoesNotExist:
raise MissingResource(task_id)
task_dict = task_serializer(task)
if 'worker_name' in task_dict:
queue_name = Worker(task_dict['worker_name'], datetime.now()).queue_name
task_dict.update({'queue': queue_name})
return self.ok(task_dict)
def test__delete_worker(self, logger, cancel, mock_add_consumer):
"""
Assert that the correct Tasks get canceled when their Worker is deleted, and that the Worker
is removed from the database.
"""
# cause two workers to be added to the database as having workers
worker_watcher.handle_worker_heartbeat({
'timestamp': time.time(),
'type': 'worker-heartbeat',
'hostname': WORKER_1,
})
worker_watcher.handle_worker_heartbeat({
'timestamp': time.time(),
'type': 'worker-heartbeat',
'hostname': WORKER_2,
})
# Let's simulate three tasks being assigned to WORKER_2, with two of them being
# in an incomplete state and one in a complete state. We will delete WORKER_2,
# which should cause the two to get canceled. Let's put task_1 in progress
TaskStatusManager.create_task_status('task_1', WORKER_2_QUEUE,
state=CALL_RUNNING_STATE)
TaskStatusManager.create_task_status('task_2', WORKER_2_QUEUE,
state=CALL_WAITING_STATE)
# This task shouldn't get canceled because it isn't in an incomplete state
TaskStatusManager.create_task_status('task_3', WORKER_2_QUEUE,
state=CALL_FINISHED_STATE)
# Let's make a task in a worker that is still present just to make sure it isn't touched.
TaskStatusManager.create_task_status('task_4', WORKER_1_QUEUE,
state=CALL_RUNNING_STATE)
# Let's just make sure the setup worked and that we have a Worker with RR2
worker_collection = Worker.get_collection()
self.assertEqual(worker_collection.find({'_id': WORKER_2}).count(), 1)
# Now let's delete the Worker named WORKER_2
tasks._delete_worker.apply_async(args=(WORKER_2,),
queue=tasks.RESOURCE_MANAGER_QUEUE)
# cancel() should have been called twice with task_1 and task_2 as parameters
self.assertEqual(cancel.call_count, 2)
# Let's build a set out of the two times that cancel was called. We can't know for sure
# which order the Tasks got canceled in, but we can assert that the correct two tasks were
# canceled (task_3 should not appear in this set).
cancel_param_set = set([c[1] for c in cancel.mock_calls])
self.assertEqual(cancel_param_set, set([('task_1',), ('task_2',)]))
# We should have logged that we are canceling the tasks
self.assertEqual(logger.call_count, 0)
self.assertTrue(WORKER_2 in logger.mock_calls[0][1][0])
self.assertTrue('Canceling the tasks' in logger.mock_calls[0][1][0])
# The Worker should have been deleted
self.assertEqual(worker_collection.find({'_id': WORKER_2}).count(), 0)
# the Worker for RW1 should remain
self.assertEqual(worker_collection.find({'_id': WORKER_1}).count(), 1)
def test_DELETE_completed_celery_task(self):
"""
Test the DELETE() method raises a TaskComplete exception if the task is already complete.
"""
task_id = '1234abcd'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(
task_id, test_worker.name, state=constants.CALL_FINISHED_STATE)
self.assertRaises(PulpCodedException, self.task_resource.DELETE,
task_id)
def test_cancel_after_task_finished(self, logger, revoke):
task_id = '1234abcd'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(task_id,
test_worker.name,
state=CALL_FINISHED_STATE)
self.assertRaises(PulpCodedException, tasks.cancel, task_id)
task_status = TaskStatusManager.find_by_task_id(task_id)
self.assertEqual(task_status['state'], CALL_FINISHED_STATE)
def GET(self, task_id):
task = TaskStatusManager.find_by_task_id(task_id)
if task is None:
raise MissingResource(task_id)
else:
link = serialization.link.link_obj('/pulp/api/v2/tasks/%s/' % task_id)
task.update(link)
task.update(serialization.dispatch.spawned_tasks(task))
if 'worker_name' in task:
queue_name = Worker(task['worker_name'], datetime.now()).queue_name
task.update({'queue': queue_name})
return self.ok(task)
def test_dispatches_inner_task(self):
self.mock_get_worker_for_reservation.return_value = Worker(
'worker1', datetime.utcnow())
tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id',
[1, 2], {'a': 2})
apply_async = self.mock_celery.tasks['task_name'].apply_async
apply_async.assert_called_once_with(1,
2,
a=2,
routing_key='worker1',
task_id='my_task_id',
exchange='C.dq')
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__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 handle_worker_heartbeat(event):
"""
Celery event handler for 'worker-heartbeat' events.
The event is first parsed and logged. If this event is from the resource manager, there is
no further processing to be done. Then the existing Worker objects are searched
for one to update. If an existing one is found, it is updated. Otherwise a new
Worker entry is created. Logging at the info and debug level is also done.
:param event: A celery event to handle.
:type event: dict
"""
event_info = _parse_and_log_event(event)
# if this is the resource_manager do nothing
if _is_resource_manager(event):
return
find_worker_criteria = Criteria(filters={'_id': event_info['worker_name']},
fields=('_id', 'last_heartbeat',
'num_reservations'))
find_worker_list = list(resources.filter_workers(find_worker_criteria))
if find_worker_list:
Worker.get_collection().find_and_modify(
query={'_id': event_info['worker_name']},
update={'$set': {
'last_heartbeat': event_info['timestamp']
}})
else:
new_worker = Worker(event_info['worker_name'], event_info['timestamp'])
msg = _("New worker '%(worker_name)s' discovered") % event_info
_logger.info(msg)
new_worker.save()
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_DELETE_celery_task(self, revoke):
"""
Test the DELETE() method with a UUID that does not correspond to a UUID that the
coordinator is aware of. This should cause a revoke call to Celery's Controller.
"""
task_id = '1234abcd'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(task_id, test_worker.name)
self.task_resource.DELETE(task_id)
revoke.assert_called_once_with(task_id, terminate=True)
def test_cancel_successful(self, logger, revoke):
task_id = '1234abcd'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(task_id, test_worker.name)
tasks.cancel(task_id)
revoke.assert_called_once_with(task_id, terminate=True)
self.assertEqual(logger.info.call_count, 1)
log_msg = logger.info.mock_calls[0][1][0]
self.assertTrue(task_id in log_msg)
self.assertTrue('Task canceled' in log_msg)
task_status = TaskStatusManager.find_by_task_id(task_id)
self.assertEqual(task_status['state'], CALL_CANCELED_STATE)
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 _delete_worker(name, normal_shutdown=False):
"""
Delete the Worker with _id name from the database, cancel any associated tasks and reservations
If the worker shutdown normally, no message is logged, otherwise an error level message is
logged. Default is to assume the worker did not shut down normally.
Any resource reservations associated with this worker are cleaned up by this function.
Any tasks associated with this worker are explicitly canceled.
:param name: The name of the worker you wish to delete. In the database, the _id
field is the name.
:type name: basestring
:param normal_shutdown: True if the worker shutdown normally, False otherwise. Defaults to
False.
:type normal_shutdown: bool
"""
if normal_shutdown is False:
msg = _(
'The worker named %(name)s is missing. Canceling the tasks in its queue.'
)
msg = msg % {'name': name}
logger.error(msg)
# Delete the worker document
worker_list = list(
resources.filter_workers(Criteria(filters={'_id': name})))
if len(worker_list) > 0:
worker_document = worker_list[0]
worker_document.delete()
# Delete all reserved_resource documents for the worker
ReservedResource.get_collection().remove({'worker_name': name})
# Cancel all of the tasks that were assigned to this worker's queue
worker = Worker.from_bson({'_id': name})
for task in TaskStatusManager.find_by_criteria(
Criteria(
filters={
'worker_name': worker.name,
'state': {
'$in': constants.CALL_INCOMPLETE_STATES
}
})):
cancel(task['task_id'])
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_filter(self):
"""
Test a filter operation to make sure the results appear to be correct.
"""
# Make three workers. We'll filter for two of them.
now = datetime.utcnow()
kw_1 = Worker('worker_1', now)
kw_1.save()
kw_2 = Worker('worker_2', now)
kw_2.save()
kw_3 = Worker('worker_3', now)
kw_3.save()
criteria = Criteria(filters={'_id': {'$gt': 'worker_1'}}, sort=[('_id', pymongo.ASCENDING)])
workers = resources.filter_workers(criteria)
# Let's assert that workers is a generator, and then let's cast it to a list so it's easier
# to test that we got the correct instances back.
self.assertEqual(type(workers), types.GeneratorType)
workers = list(workers)
self.assertEqual(all([isinstance(w, Worker) for w in workers]), True)
self.assertEqual(workers[0].name, 'worker_2')
self.assertEqual(workers[1].name, 'worker_3')
def test_apply_async_task_canceled(self, apply_async):
"""
Assert that apply_async() honors 'canceled' task status.
"""
args = [1, 'b', 'iii']
kwargs = {'1': 'for the money', 'tags': ['test_tags']}
task_id = 'test_task_id'
now = datetime.utcnow()
TaskStatusManager.create_task_status(task_id,
Worker('test-worker', now),
state=CALL_CANCELED_STATE)
apply_async.return_value = celery.result.AsyncResult(task_id)
task = tasks.Task()
task.apply_async(*args, **kwargs)
task_status = TaskStatusManager.find_by_task_id(task_id)
self.assertEqual(task_status['state'], 'canceled')
self.assertEqual(task_status['start_time'], None)
def _delete_worker(name, normal_shutdown=False):
"""
Delete the Worker with _id name from the database, cancel any associated tasks and reservations
If the worker shutdown normally, no message is logged, otherwise an error level message is
logged. Default is to assume the worker did not shut down normally.
Any resource reservations associated with this worker are cleaned up by this function.
Any tasks associated with this worker are explicitly canceled.
:param name: The name of the worker you wish to delete. In the database, the _id
field is the name.
:type name: basestring
:param normal_shutdown: True if the worker shutdown normally, False otherwise. Defaults to
False.
:type normal_shutdown: bool
"""
if normal_shutdown is False:
msg = _('The worker named %(name)s is missing. Canceling the tasks in its queue.')
msg = msg % {'name': name}
_logger.error(msg)
# Delete the worker document
worker_list = list(resources.filter_workers(Criteria(filters={'_id': name})))
if len(worker_list) > 0:
worker_document = worker_list[0]
worker_document.delete()
# Delete all reserved_resource documents for the worker
ReservedResource.get_collection().remove({'worker_name': name})
# Cancel all of the tasks that were assigned to this worker's queue
worker = Worker.from_bson({'_id': name})
for task_status in TaskStatus.objects(worker_name=worker.name,
state__in=constants.CALL_INCOMPLETE_STATES):
cancel(task_status['task_id'])
def test_DELETE_doesnt_cancel_spawned_celery_task(self, revoke):
"""
Test the DELETE() which should cause a revoke call to Celery's Controller.
This also tests that the spawned tasks are canceled as well.
"""
task_id = '1234abcd'
spawned_task_id = 'spawned_task'
spawned_by_spawned_task_id = 'spawned_by_spawned_task'
now = datetime.utcnow()
test_worker = Worker('test_worker', now)
TaskStatusManager.create_task_status(task_id, test_worker.queue_name)
TaskStatusManager.create_task_status(spawned_task_id,
test_worker.queue_name)
TaskStatusManager.create_task_status(spawned_by_spawned_task_id,
test_worker.queue_name)
TaskStatusManager.update_task_status(
task_id, delta={'spawned_tasks': [spawned_task_id]})
TaskStatusManager.update_task_status(
spawned_task_id,
delta={'spawned_tasks': [spawned_by_spawned_task_id]})
self.task_resource.DELETE(task_id)
self.assertEqual(revoke.call_count, 1)
revoke.assert_called_once_with(task_id, terminate=True)
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(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(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 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.task_id})
self.assertEqual(rr_1['worker_name'], reserved_resource_1.worker_name)
self.assertEqual(rr_1['resource_id'], 'resource_1')
rr_2 = rrc.find_one({'_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 tearDown(self):
Worker.get_collection().remove()
ReservedResource.get_collection().remove()
TaskStatus.objects().delete()
def tearDown(self):
Worker.get_collection().remove()
ReservedResource.get_collection().remove()
TaskStatus.get_collection().remove()
def test_filter(self):
"""
Test a filter operation to make sure the results appear to be correct.
"""
# Make three workers. We'll filter for two of them.
now = datetime.utcnow()
kw_1 = Worker('worker_1', now)
kw_1.save()
kw_2 = Worker('worker_2', now)
kw_2.save()
kw_3 = Worker('worker_3', now)
kw_3.save()
criteria = Criteria(filters={'_id': {
'$gt': 'worker_1'
}},
sort=[('_id', pymongo.ASCENDING)])
workers = resources.filter_workers(criteria)
# Let's assert that workers is a generator, and then let's cast it to a list so it's easier
# to test that we got the correct instances back.
self.assertEqual(type(workers), types.GeneratorType)
workers = list(workers)
self.assertEqual(all([isinstance(w, Worker) for w in workers]), True)
self.assertEqual(workers[0].name, 'worker_2')
self.assertEqual(workers[1].name, 'worker_3')
