Split ownership-fenced writes into testable pieces

SonarCloud flagged update_model's cognitive complexity and uncovered
takeover paths. Move the transactional write and the retry wait out of
update_model, move BaseTask.run's finalize-and-release into its own
method, and cover the takeover, lost-claim and retry paths.

Co-Authored-By: Claude <noreply@anthropic.com>

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')

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.

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'})

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')

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
+        )