diff --git a/awx/main/tasks/jobs.py b/awx/main/tasks/jobs.py index a10367a076..08ed0a45f5 100644 --- a/awx/main/tasks/jobs.py +++ b/awx/main/tasks/jobs.py @@ -86,7 +86,7 @@ from awx.main.utils.common import ( ) from awx.conf.license import get_license from awx.main.utils.handlers import SpecialInventoryHandler -from awx.main.utils.update_model import update_model +from awx.main.utils.update_model import NotOwner, update_model # Django flags from flags.state import flag_enabled @@ -204,16 +204,19 @@ def dispatch_waiting_jobs(binder): UnifiedJob.objects.filter(pk=uj.pk, status='waiting').update(status='running', start_args='') -def _finalize_job_run(model, pk, runner_callback, status, extra_fields=None): +def _finalize_job_run(model, pk, runner_callback, status, extra_fields=None, owner_task_id=None): """Commit terminal status, trigger notifications, and emit websocket status. Shared by BaseTask.run() (normal path) and adoption (reattach_to_work_unit path). Uses duck typing on the instance to schedule dependent task/workflow managers. + + With owner_task_id, nothing is written and NotOwner is raised if another task has + taken the job over (adoption records its own task id in celery_task_id). """ all_fields = runner_callback.get_delayed_update_fields() if extra_fields: all_fields.update(extra_fields) - instance = update_model(model, pk, status=status, select_for_update=True, **all_fields) + instance = update_model(model, pk, status=status, select_for_update=True, owner_task_id=owner_task_id, **all_fields) if not instance: return None if (instance.host_status_counts is not None) or (not runner_callback.wrapup_event_dispatched): @@ -673,6 +676,9 @@ class BaseTask(object): return self.instance.websocket_emit_status("running") + # The task id this run was dispatched under. If adoption records a different one, + # this run no longer owns the job and must not write its final status. + owner_task_id = self.instance.celery_task_id or None status, rc = 'error', None self.runner_callback.event_ct = 0 @@ -853,11 +859,17 @@ 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) + 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 + return finally: # Guarantee work unit release even if finalization throws - if receptor_job and getattr(receptor_job, 'work', None): + if release_work and receptor_job and getattr(receptor_job, 'work', None): try: receptor_job._receptor_release_work(receptor_job.work, status) except Exception: diff --git a/awx/main/tasks/receptor.py b/awx/main/tasks/receptor.py index e57a595372..293c2bb98e 100644 --- a/awx/main/tasks/receptor.py +++ b/awx/main/tasks/receptor.py @@ -38,6 +38,7 @@ from awx.main.utils.common import ( from awx.main.constants import JOB_FOLDER_PREFIX, MAX_ISOLATED_PATH_COLON_DELIMITER from awx.main.tasks.signals import signal_state, signal_callback, SignalExit from awx.main.tasks.callback import RunnerCallback +from awx.main.utils.update_model import NotOwner from awx.main.models import Instance, InstanceLink, UnifiedJob, ReceptorAddress from awx.main.dispatch import get_task_queuename @@ -943,7 +944,7 @@ def _get_or_create_private_data_dir(job): ) -def _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock=None): +def _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock=None, owner_task_id=None): """Commit terminal status for an adopted job via the shared _finalize_job_run path. The shared finalization function uses duck typing to schedule task/workflow managers @@ -964,7 +965,7 @@ def _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock= if job.started: extra['elapsed'] = (finished_at - job.started).total_seconds() - _finalize_job_run(type(job), job.pk, callback, final_status, extra_fields=extra) + _finalize_job_run(type(job), job.pk, callback, final_status, extra_fields=extra, owner_task_id=owner_task_id) label = 'exit_code (process phase raised)' if process_phase_failed else 'adoption' logger.info(f'Job {job.id} finalized via {label}: {final_status}') @@ -1001,7 +1002,7 @@ def _compute_adoption_dedup(job): return safe_threshold, collision_zone -def reattach_to_work_unit(job, work: ReceptorWork, clock=None): +def reattach_to_work_unit(job, work: ReceptorWork, clock=None, owner_task_id=None): """Reconnect to a receptor work unit and stream events in real-time until it completes. Reconstructs the minimal process-phase context from the DB job record, then calls @@ -1058,10 +1059,11 @@ def reattach_to_work_unit(job, work: ReceptorWork, clock=None): # # _process_phase -> _handle_work_error may return None for 'exceeded quota' where # the job is already set to 'pending' — don't finalize, let the next dispatch handle it. + release_work = True try: if res is not None: exit_code = 0 if getattr(res, 'status', '') == 'successful' else 1 - _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock=clock) + _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock=clock, owner_task_id=owner_task_id) elif not process_phase_failed: logger.info(f'Job {job.id}: adoption deferred (handled in _handle_work_error)') else: @@ -1074,9 +1076,14 @@ def reattach_to_work_unit(job, work: ReceptorWork, clock=None): unit_status = {} state_name = '' exit_code = _get_adoption_exit_code(unit_status, state_name) - _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock=clock) + _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock=clock, owner_task_id=owner_task_id) + except NotOwner as exc: + # A later claim took the job over; its final status and work unit are that task's. + logger.warning(f'Job {job.id} was claimed by task {exc.current_task_id} during adoption; not finalizing or releasing unit {unit_id}') + release_work = False finally: - receptor_job._receptor_release_work(work, getattr(res, 'status', 'error')) + if release_work: + receptor_job._receptor_release_work(work, getattr(res, 'status', 'error')) shutil.rmtree(private_data_dir, ignore_errors=True) return True diff --git a/awx/main/tasks/system.py b/awx/main/tasks/system.py index 9554a8d05a..78d5a358ed 100644 --- a/awx/main/tasks/system.py +++ b/awx/main/tasks/system.py @@ -7,6 +7,7 @@ import logging import os import shutil import time +import uuid from collections import namedtuple from contextlib import redirect_stdout from packaging.version import Version @@ -978,8 +979,7 @@ def _process_startup_jobs(this_inst): for j in jobs: try: if startup_job_action(j) == JobAction.ADOPT: - obj, _ = adopt_job_async.apply_async(args=[j.id], queue=get_task_queuename()) - UnifiedJob.objects.filter(pk=j.id).update(celery_task_id=obj['uuid']) + _dispatch_adoption(j) else: reaped_ids.append(j.id) reaper.reap_job( @@ -1019,38 +1019,62 @@ def _process_running_jobs(this_inst, active_task_ids, ref_time): try: action = running_job_action(j, this_inst.hostname, active_task_ids, ref_time, workflow_ctype_id) if action == JobAction.ADOPT: - obj, _ = adopt_job_async.apply_async(args=[j.id], queue=get_task_queuename()) - # Record the adoption task UUID so subsequent heartbeats see this job - # as active and skip it — preventing redundant re-adoption dispatches. - UnifiedJob.objects.filter(pk=j.id).update(celery_task_id=obj['uuid']) + _dispatch_adoption(j) elif action == JobAction.REAP: reaper.reap_job(j, 'failed') except Exception: logger.exception(f'Failed processing job {j.id} in heartbeat job loop') +def claim_for_adoption(job): + """Make a new adoption task the job's owner. Return its task id, or None if the job changed. + + The claim replaces celery_task_id only if it still holds the value this heartbeat read, + so two claims can't both win. Once claimed, the previous owner's final status write is + refused (see NotOwner) and cancels go to the adoption task. + """ + task_id = str(uuid.uuid4()) + claimed = UnifiedJob.objects.filter(pk=job.id, status='running', celery_task_id=job.celery_task_id).update(celery_task_id=task_id) + return task_id if claimed else None + + +def _dispatch_adoption(job): + """Claim the job, then dispatch its adoption under the claimed task id.""" + task_id = claim_for_adoption(job) + if task_id is None: + logger.info(f'Job {job.id} changed since it was read; not adopting it on this heartbeat') + return + # The claim is recorded before dispatch, so subsequent heartbeats see this job as + # active and skip it, and the adoption task can check that it still owns the job. + adopt_job_async.apply_async(args=[job.id], kwargs={'owner_task_id': task_id}, uuid=task_id, queue=get_task_queuename()) + + @task(queue=get_task_queuename, timeout=3600 * 2, on_duplicate='discard') -def adopt_job_async(job_id): +def adopt_job_async(job_id, owner_task_id=None): """Adopt a single orphaned job in a background task, streaming events in real-time. Called from _process_startup_jobs and _process_running_jobs via apply_async so the heartbeat returns immediately. on_duplicate='discard' ensures only one adoption runs per job across heartbeat cycles. """ - adopt_job(job_id, open_work=lambda: ReceptorWork(get_receptor_ctl())) + adopt_job(job_id, open_work=lambda: ReceptorWork(get_receptor_ctl()), owner_task_id=owner_task_id) -def adopt_job(job_id, open_work, clock=None): +def adopt_job(job_id, open_work, clock=None, owner_task_id=None): """Adopt one orphaned job: fail it if it has been silent too long, else reattach to its work unit. open_work returns the ReceptorWork to stream from; it is only called when the job is - actually reattached. clock defaults to django.utils.timezone.now. + actually reattached. clock defaults to django.utils.timezone.now. owner_task_id is the + task id the job was claimed under; if a later claim replaced it, this adoption does nothing. """ clock = clock or now job = UnifiedJob.objects.filter(id=job_id, status='running').first() if not job: logger.debug(f'adopt_job_async: job {job_id} is no longer running, skipping') return + if owner_task_id and job.celery_task_id != owner_task_id: + logger.info(f'adopt_job_async: job {job_id} was claimed again by task {job.celery_task_id}, skipping') + return adoption_timeout = settings.HADR_JOB_ADOPTION_TIMEOUT last_event_time = job.get_event_queryset().aggregate(Max('created'))['created__max'] @@ -1061,7 +1085,7 @@ def adopt_job(job_id, open_work, clock=None): work = open_work() try: - reattach_to_work_unit(job, work, clock=clock) + reattach_to_work_unit(job, work, clock=clock, owner_task_id=owner_task_id) except Exception: logger.exception(f'adopt_job_async: adoption failed for job {job.id} (unit={job.work_unit_id})') finally: diff --git a/awx/main/tests/functional/tasks/test_adoption_ownership.py b/awx/main/tests/functional/tasks/test_adoption_ownership.py new file mode 100644 index 0000000000..9b8f2ac27e --- /dev/null +++ b/awx/main/tests/functional/tasks/test_adoption_ownership.py @@ -0,0 +1,175 @@ +"""Who may write a job's final status once adoption claims it. + +A claim records the adoption task's id in celery_task_id before the task is dispatched. +After that, a final status write from any other task is refused with NotOwner, so a +controller that only looked dead can't overwrite the adopted job when it comes back. +""" + +from unittest import mock + +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.tests.fake_receptor import FakeReceptorWork, FakeWorkUnit, RecordingDispatcher, make_events +from awx.main.utils.update_model import NotOwner, update_model + +UNIT = 'unit-1' +ORIGINAL = 'original-task' + + +@pytest.fixture(autouse=True) +def adoption_env(settings, tmp_path): + settings.AWX_ISOLATION_BASE_PATH = str(tmp_path) + RecordingDispatcher.reset() + with ( + mock.patch('awx.main.tasks.receptor.connections'), + mock.patch('awx.main.tasks.callback.CallbackQueueDispatcher', RecordingDispatcher), + mock.patch('awx.main.tasks.jobs.ScheduleTaskManager'), + mock.patch('awx.main.tasks.jobs.ScheduleWorkflowManager'), + mock.patch('awx.main.models.unified_jobs.UnifiedJob.websocket_emit_status'), + ): + yield + + +@pytest.fixture +def running_job(me_inst): + return Job.objects.create(controller_node=me_inst.hostname, status='running', work_unit_id=UNIT, celery_task_id=ORIGINAL) + + +def finalize_as(job, owner_task_id, status): + callback = RunnerCallback(model=Job) + callback.instance = job + return _finalize_job_run(Job, job.pk, callback, status, owner_task_id=owner_task_id) + + +# ── update_model ────────────────────────────────────────────────────────────── + + +@pytest.mark.django_db +def test_owner_may_write(running_job): + update_model(Job, running_job.pk, owner_task_id=ORIGINAL, job_explanation='mine') + running_job.refresh_from_db() + assert running_job.job_explanation == 'mine' + + +@pytest.mark.django_db +def test_former_owner_is_refused_and_nothing_is_written(running_job): + Job.objects.filter(pk=running_job.pk).update(celery_task_id='adoption-task') + + with pytest.raises(NotOwner) as exc: + update_model(Job, running_job.pk, owner_task_id=ORIGINAL, status='error') + + assert exc.value.current_task_id == 'adoption-task' + running_job.refresh_from_db() + assert running_job.status == 'running' + + +@pytest.mark.django_db +def test_write_without_an_owner_is_not_fenced(running_job): + Job.objects.filter(pk=running_job.pk).update(celery_task_id='adoption-task') + update_model(Job, running_job.pk, job_explanation='unfenced') + running_job.refresh_from_db() + assert running_job.job_explanation == 'unfenced' + + +# ── The claim ───────────────────────────────────────────────────────────────── + + +@pytest.mark.django_db +def test_claim_records_a_new_task_id(running_job): + task_id = claim_for_adoption(running_job) + + assert task_id and task_id != ORIGINAL + running_job.refresh_from_db() + assert running_job.celery_task_id == task_id + + +@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. + first_read = Job.objects.get(pk=running_job.pk) + second_read = Job.objects.get(pk=running_job.pk) + + assert claim_for_adoption(first_read) is not None + assert claim_for_adoption(second_read) is None + + +@pytest.mark.django_db +def test_heartbeat_claims_before_dispatching_under_the_claimed_id(running_job, me_inst): + with mock.patch('awx.main.tasks.system.adopt_job_async') as task: + _process_running_jobs(me_inst, ['some-other-task'], None) + + running_job.refresh_from_db() + kwargs = task.apply_async.call_args.kwargs + assert kwargs['uuid'] == running_job.celery_task_id + assert kwargs['kwargs'] == {'owner_task_id': running_job.celery_task_id} + + +# ── Stale writers after a claim ─────────────────────────────────────────────── + + +@pytest.mark.django_db +def test_returning_controller_cannot_overwrite_the_adopted_result(running_job): + # The original controller looked dead, so this one claimed and adopted the job. + task_id = claim_for_adoption(running_job) + work = FakeReceptorWork({UNIT: FakeWorkUnit(make_events(10))}) + adopt_job(running_job.id, open_work=lambda: work, owner_task_id=task_id) + running_job.refresh_from_db() + assert running_job.status == 'successful' + + # The original controller comes back and tries to write its own final status. + with pytest.raises(NotOwner): + finalize_as(running_job, ORIGINAL, 'error') + + running_job.refresh_from_db() + assert running_job.status == 'successful' + + +@pytest.mark.django_db +def test_original_owner_finishing_first_is_refused_while_adoption_proceeds(running_job): + # The claim lands while the original task is still streaming. When it finishes, + # its write is refused; the adoption then streams and finalizes. + task_id = claim_for_adoption(running_job) + + with pytest.raises(NotOwner): + finalize_as(running_job, ORIGINAL, 'successful') + running_job.refresh_from_db() + assert running_job.status == 'running' + + work = FakeReceptorWork({UNIT: FakeWorkUnit(make_events(10))}) + adopt_job(running_job.id, open_work=lambda: work, owner_task_id=task_id) + running_job.refresh_from_db() + assert running_job.status == 'successful' + + +@pytest.mark.django_db +def test_adoption_claimed_again_before_it_starts_does_nothing(running_job): + stale = claim_for_adoption(running_job) + claim_for_adoption(Job.objects.get(pk=running_job.pk)) + opened = mock.Mock() + + adopt_job(running_job.id, open_work=opened, owner_task_id=stale) + + opened.assert_not_called() + running_job.refresh_from_db() + assert running_job.status == 'running' + + +@pytest.mark.django_db +def test_adoption_that_loses_ownership_mid_stream_leaves_the_unit(running_job): + stale = claim_for_adoption(running_job) + work = FakeReceptorWork({UNIT: FakeWorkUnit(make_events(10))}) + + def reclaim_then_open(): + # Another claim lands after this adoption started but before it finalizes. + claim_for_adoption(Job.objects.get(pk=running_job.pk)) + return work + + adopt_job(running_job.id, open_work=reclaim_then_open, owner_task_id=stale) + + running_job.refresh_from_db() + assert running_job.status == 'running', 'the new owner finalizes, not the stale adoption' + assert UNIT in work.units, 'the stale adoption does not release the new owner\'s unit' diff --git a/awx/main/tests/functional/tasks/test_tasks_system.py b/awx/main/tests/functional/tasks/test_tasks_system.py index add20f4fb2..bb28ce3009 100644 --- a/awx/main/tests/functional/tasks/test_tasks_system.py +++ b/awx/main/tests/functional/tasks/test_tasks_system.py @@ -745,8 +745,10 @@ def test_adoption_skips_still_running_work_unit(me_inst): mock_adopt.apply_async = MagicMock() _process_startup_jobs(me_inst) - mock_adopt.apply_async.assert_called_once_with(args=[job.id], queue=mock_adopt.apply_async.call_args[1]['queue']) job.refresh_from_db() + mock_adopt.apply_async.assert_called_once_with( + args=[job.id], kwargs={'owner_task_id': job.celery_task_id}, uuid=job.celery_task_id, queue=mock_adopt.apply_async.call_args[1]['queue'] + ) assert job.status == 'running', 'job must not be reaped — adoption deferred to background task' diff --git a/awx/main/utils/update_model.py b/awx/main/utils/update_model.py index 37f35f7091..5db4d38761 100644 --- a/awx/main/utils/update_model.py +++ b/awx/main/utils/update_model.py @@ -9,10 +9,25 @@ from awx.main.tasks.signals import signal_callback logger = logging.getLogger('awx.main.tasks.utils') -def update_model(model, pk, _attempt=0, _max_attempts=5, select_for_update=False, **updates): +class NotOwner(Exception): + """The job is now owned by another task, so this task may not write it.""" + + def __init__(self, pk, owner_task_id, current_task_id): + self.pk = pk + self.owner_task_id = owner_task_id + self.current_task_id = current_task_id + super().__init__(f'pk={pk} is owned by task {current_task_id!r}, not {owner_task_id!r}') + + +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. + + With owner_task_id, the row is locked and updated only while its celery_task_id + still equals owner_task_id; otherwise NotOwner is raised and nothing is written. """ + if owner_task_id: + select_for_update = True try: with transaction.atomic(): # Retrieve the model instance. @@ -21,6 +36,9 @@ def update_model(model, pk, _attempt=0, _max_attempts=5, select_for_update=False 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: @@ -45,7 +63,9 @@ def update_model(model, pk, _attempt=0, _max_attempts=5, select_for_update=False 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, **updates) + 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: logger.warning(f'Failed to update {model._meta.object_name} pk={pk} after {_attempt} retries.') raise