diff --git a/awx/main/tasks/jobs.py b/awx/main/tasks/jobs.py index 08ed0a45f5..234e0e55f0 100644 --- a/awx/main/tasks/jobs.py +++ b/awx/main/tasks/jobs.py @@ -649,6 +649,27 @@ class BaseTask(object): def should_use_fact_cache(self): return False + def _finalize_and_release(self, pk, status, receptor_job, owner_task_id): + """Write the job's final status, then release its work unit. + + Return False when another task adopted the job while it ran; its final status + and work unit are then that task's, so neither is touched here. + """ + release_work = True + try: + self.instance = _finalize_job_run(self.model, pk, self.runner_callback, status, owner_task_id=owner_task_id) + except NotOwner as exc: + logger.warning(f'{self.instance.log_format} was taken over by task {exc.current_task_id}; not writing status {status} or releasing its work unit') + release_work = False + finally: + # Guarantee work unit release even if finalization throws + if release_work and receptor_job and getattr(receptor_job, 'work', None): + try: + receptor_job._receptor_release_work(receptor_job.work, status) + except Exception: + logger.exception(f'Failed to release work unit {getattr(receptor_job, "unit_id", "unknown")}') + return release_work + @with_path_cleanup @with_signal_handling def run(self, pk, **kwargs): @@ -859,21 +880,8 @@ class BaseTask(object): logger.exception('{} Post run hook errored.'.format(self.instance.log_format)) self.private_data_dir = private_data_dir - release_work = True - try: - self.instance = _finalize_job_run(self.model, pk, self.runner_callback, status, owner_task_id=owner_task_id) - except NotOwner as exc: - # Another task adopted this job while it ran; its final status and work unit are theirs. - logger.warning(f'{self.instance.log_format} was taken over by task {exc.current_task_id}; not writing status {status} or releasing its work unit') - release_work = False + if not self._finalize_and_release(pk, status, receptor_job, owner_task_id): return - finally: - # Guarantee work unit release even if finalization throws - if release_work and receptor_job and getattr(receptor_job, 'work', None): - try: - receptor_job._receptor_release_work(receptor_job.work, status) - except Exception: - logger.exception(f'Failed to release work unit {getattr(receptor_job, "unit_id", "unknown")}') if not self.instance: logger.error(f'Unified job pk={pk} appears to be deleted while running') diff --git a/awx/main/tests/functional/tasks/test_adoption_ownership.py b/awx/main/tests/functional/tasks/test_adoption_ownership.py index 9b8f2ac27e..8171998b72 100644 --- a/awx/main/tests/functional/tasks/test_adoption_ownership.py +++ b/awx/main/tests/functional/tasks/test_adoption_ownership.py @@ -12,7 +12,7 @@ import pytest from awx.main.models import Job from awx.main.tasks.callback import RunnerCallback from awx.main.tasks.jobs import _finalize_job_run -from awx.main.tasks.system import _process_running_jobs, adopt_job, claim_for_adoption +from awx.main.tasks.system import _dispatch_adoption, _process_running_jobs, adopt_job, claim_for_adoption from awx.main.tests.fake_receptor import FakeReceptorWork, FakeWorkUnit, RecordingDispatcher, make_events from awx.main.utils.update_model import NotOwner, update_model @@ -82,11 +82,23 @@ def test_write_without_an_owner_is_not_fenced(running_job): def test_claim_records_a_new_task_id(running_job): task_id = claim_for_adoption(running_job) - assert task_id and task_id != ORIGINAL + assert task_id is not None + assert task_id != ORIGINAL running_job.refresh_from_db() assert running_job.celery_task_id == task_id +@pytest.mark.django_db +def test_dispatch_after_a_lost_claim_sends_nothing(running_job): + stale_read = Job.objects.get(pk=running_job.pk) + Job.objects.filter(pk=running_job.pk).update(celery_task_id='another-heartbeat') + + with mock.patch('awx.main.tasks.system.adopt_job_async.apply_async') as apply_async: + _dispatch_adoption(stale_read) + + apply_async.assert_not_called() + + @pytest.mark.django_db def test_only_one_of_two_claims_from_the_same_read_wins(running_job): # Two heartbeats read the job before either claims it. diff --git a/awx/main/tests/functional/utils/test_update_model.py b/awx/main/tests/functional/utils/test_update_model.py index 94304cca30..e4e4787f31 100644 --- a/awx/main/tests/functional/utils/test_update_model.py +++ b/awx/main/tests/functional/utils/test_update_model.py @@ -45,3 +45,22 @@ def test_deleted_job(normal_job): job_pk = normal_job.pk normal_job.delete() assert update_model(Job, job_pk) is None + + +@pytest.mark.django_db +def test_abort_signal_stops_the_retry(normal_job, mocker): + mocker.patch.object(Job.objects, 'get', side_effect=DatabaseError) + mocker.patch('awx.main.utils.update_model.time.sleep') + mocker.patch('awx.main.utils.update_model.signal_callback', return_value=True) + with pytest.raises(RuntimeError, match='abort signal'): + update_model(Job, normal_job.pk) + + +@pytest.mark.django_db +def test_retry_keeps_the_owner_fence(normal_job, mocker): + Job.objects.filter(pk=normal_job.pk).update(celery_task_id='adoption-task') + write = mocker.patch('awx.main.utils.update_model._write_fields', side_effect=[DatabaseError, normal_job]) + mocker.patch('awx.main.utils.update_model.time.sleep') + + assert update_model(Job, normal_job.pk, owner_task_id='adoption-task', status='failed') is normal_job + assert write.call_args_list[1].args == (Job, normal_job.pk, True, 'adoption-task', {'status': 'failed'}) diff --git a/awx/main/tests/unit/tasks/test_jobs.py b/awx/main/tests/unit/tasks/test_jobs.py index 0519c98fd8..e7185c0cc2 100644 --- a/awx/main/tests/unit/tasks/test_jobs.py +++ b/awx/main/tests/unit/tasks/test_jobs.py @@ -24,6 +24,7 @@ from awx.main.models import ( AdHocCommand, ) from awx.main.tasks import jobs +from awx.main.utils.update_model import NotOwner from ansible_base.lib.workload_identity.controller import AutomationControllerJobScope @@ -723,6 +724,44 @@ def test_work_unit_not_released_if_no_receptor_work(): assert receptor_job._receptor_release_work.call_count == 0 +def _finalizing_task(): + task = jobs.BaseTask.__new__(jobs.BaseTask) + task.model = Job + task.instance = mock.MagicMock(log_format='job 1 (running)') + task.runner_callback = mock.MagicMock() + return task + + +@mock.patch('awx.main.tasks.jobs._finalize_job_run') +def test_finalize_and_release_releases_the_work_unit(mock_finalize): + task = _finalizing_task() + receptor_job = mock.MagicMock() + + assert task._finalize_and_release(1, 'successful', receptor_job, 'task-1') is True + mock_finalize.assert_called_once_with(Job, 1, task.runner_callback, 'successful', owner_task_id='task-1') + receptor_job._receptor_release_work.assert_called_once_with(receptor_job.work, 'successful') + + +@mock.patch('awx.main.tasks.jobs._finalize_job_run', side_effect=NotOwner(1, 'task-1', 'adoption-task')) +def test_finalize_and_release_leaves_an_adopted_job_alone(mock_finalize): + task = _finalizing_task() + receptor_job = mock.MagicMock() + + assert task._finalize_and_release(1, 'successful', receptor_job, 'task-1') is False + receptor_job._receptor_release_work.assert_not_called() + + +@mock.patch('awx.main.tasks.jobs._finalize_job_run', side_effect=RuntimeError('db down')) +def test_finalize_and_release_still_releases_when_finalize_raises(mock_finalize): + task = _finalizing_task() + receptor_job = mock.MagicMock() + receptor_job._receptor_release_work.side_effect = Exception('receptor gone') + + with pytest.raises(RuntimeError): + task._finalize_and_release(1, 'failed', receptor_job, None) + receptor_job._receptor_release_work.assert_called_once_with(receptor_job.work, 'failed') + + @mock.patch('awx.main.tasks.jobs.ScheduleWorkflowManager') @mock.patch('awx.main.tasks.jobs.ScheduleTaskManager') @mock.patch('awx.main.tasks.jobs.events_processed_hook') diff --git a/awx/main/utils/update_model.py b/awx/main/utils/update_model.py index 5db4d38761..9f5c2eb498 100644 --- a/awx/main/utils/update_model.py +++ b/awx/main/utils/update_model.py @@ -19,6 +19,39 @@ class NotOwner(Exception): super().__init__(f'pk={pk} is owned by task {current_task_id!r}, not {owner_task_id!r}') +def _write_fields(model, pk, select_for_update, owner_task_id, updates): + """Load the row, check ownership, and save the given fields in one transaction.""" + with transaction.atomic(): + # Retrieve the model instance. + if select_for_update: + instance = model.objects.select_for_update().get(pk=pk) + else: + instance = model.objects.get(pk=pk) + + if owner_task_id and instance.celery_task_id != owner_task_id: + raise NotOwner(pk, owner_task_id, instance.celery_task_id) + + # Update the appropriate fields and save the model + # instance, then return the new instance. + if updates: + update_fields = ['modified'] + for field, value in updates.items(): + setattr(instance, field, value) + update_fields.append(field) + if field == 'status': + update_fields.append('failed') + instance.save(update_fields=update_fields) + return instance + + +def _wait_before_retry(pk): + """Sleep five seconds before a retry, giving up early on an abort signal.""" + for i in range(5): + time.sleep(1) + if signal_callback(): + raise RuntimeError(f'Could not fetch {pk} because of receiving abort signal') + + def update_model(model, pk, _attempt=0, _max_attempts=5, select_for_update=False, owner_task_id=None, **updates): """Reload the model instance from the database and update the given fields. @@ -29,27 +62,7 @@ def update_model(model, pk, _attempt=0, _max_attempts=5, select_for_update=False if owner_task_id: select_for_update = True try: - with transaction.atomic(): - # Retrieve the model instance. - if select_for_update: - instance = model.objects.select_for_update().get(pk=pk) - else: - instance = model.objects.get(pk=pk) - - if owner_task_id and instance.celery_task_id != owner_task_id: - raise NotOwner(pk, owner_task_id, instance.celery_task_id) - - # Update the appropriate fields and save the model - # instance, then return the new instance. - if updates: - update_fields = ['modified'] - for field, value in updates.items(): - setattr(instance, field, value) - update_fields.append(field) - if field == 'status': - update_fields.append('failed') - instance.save(update_fields=update_fields) - return instance + return _write_fields(model, pk, select_for_update, owner_task_id, updates) except ObjectDoesNotExist: return None except (DatabaseError, InterfaceError) as e: @@ -58,14 +71,10 @@ def update_model(model, pk, _attempt=0, _max_attempts=5, select_for_update=False # Attempt to retry the update, assuming we haven't already # tried too many times. - if _attempt < _max_attempts: - for i in range(5): - time.sleep(1) - if signal_callback(): - raise RuntimeError(f'Could not fetch {pk} because of receiving abort signal') - return update_model( - model, pk, _attempt=_attempt + 1, _max_attempts=_max_attempts, select_for_update=select_for_update, owner_task_id=owner_task_id, **updates - ) - else: + if _attempt >= _max_attempts: logger.warning(f'Failed to update {model._meta.object_name} pk={pk} after {_attempt} retries.') raise + _wait_before_retry(pk) + return update_model( + model, pk, _attempt=_attempt + 1, _max_attempts=_max_attempts, select_for_update=select_for_update, owner_task_id=owner_task_id, **updates + )