From 274d5abed8d8013b190b541e88da3b8d9ce96352 Mon Sep 17 00:00:00 2001 From: Chris Meyers Date: Sat, 10 Oct 2026 21:35:03 -0400 Subject: [PATCH] Fence job status writes by the owning task Adopting an orphaned job now claims it first: the heartbeat replaces the job's celery_task_id with a new task id, only if it still holds the value the heartbeat read, and then dispatches adopt_job_async under that id. Previously the task was dispatched first and its id recorded after, and two heartbeats could both record one. Final status writes are fenced by that id. update_model takes owner_task_id and raises NotOwner, writing nothing, when the job's celery_task_id no longer matches. BaseTask.run passes the id it was dispatched under, and adoption passes its claimed id. A task that loses ownership neither writes its final status nor releases the work unit, which now belongs to the new owner. An adoption whose claim was replaced before it starts does nothing. Cancels follow celery_task_id, so they go to the current owner. This changes behavior: a controller that only looked dead can no longer overwrite an adopted job's result when it comes back. Co-Authored-By: Claude --- awx/main/tasks/jobs.py | 22 ++- awx/main/tasks/receptor.py | 19 +- awx/main/tasks/system.py | 46 +++-- .../tasks/test_adoption_ownership.py | 175 ++++++++++++++++++ .../functional/tasks/test_tasks_system.py | 4 +- awx/main/utils/update_model.py | 24 ++- 6 files changed, 265 insertions(+), 25 deletions(-) create mode 100644 awx/main/tests/functional/tasks/test_adoption_ownership.py 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