Let the heartbeat and adoption take a clock

The heartbeat's judgment of lost peers and orphaned jobs, the adoption
deadline, and an adopted job's finish time now read the time from a
clock callable that defaults to django.utils.timezone.now.
cluster_node_heartbeat delegates to run_heartbeat(binder, clock), and
adopt_job_async to adopt_job(job_id, open_work, clock). Behavior is
unchanged.

awx.main.tests.fake_clock provides FakeClock. New functional tests check
the lost-peer grace period to the second, clock skew in both
directions (a fast judge finds a healthy peer lost; a dead peer with a
fast clock is found late), and the adoption deadline end to end against
the fake receptor.

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

awx/main/tasks/receptor.py

@@ -212,6 +212,9 @@ class ReceptorWork:
     def release(self, unit_id):
         return self.receptor_ctl.simple_command(f'work release {unit_id}')
 
+    def close(self):
+        self.receptor_ctl.close()
+
 
 def adopt_remote_work(receptor_ctl, node, unit_id, config_data=None):
     """Adopt a running work unit from a remote node, passing TLS/signwork when configured.
@@ -940,7 +943,7 @@ def _get_or_create_private_data_dir(job):
     )
 
 
-def _finalize_adopted_job(job, callback, exit_code, process_phase_failed):
+def _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock=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
@@ -956,7 +959,7 @@ def _finalize_adopted_job(job, callback, exit_code, process_phase_failed):
         return
 
     final_status = 'successful' if exit_code == 0 else 'failed'
-    finished_at = now()
+    finished_at = (clock or now)()
     extra = {'finished': finished_at}
     if job.started:
         extra['elapsed'] = (finished_at - job.started).total_seconds()
@@ -998,7 +1001,7 @@ def _compute_adoption_dedup(job):
     return safe_threshold, collision_zone
 
 
-def reattach_to_work_unit(job, work: ReceptorWork):
+def reattach_to_work_unit(job, work: ReceptorWork, clock=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,7 +1061,7 @@ def reattach_to_work_unit(job, work: ReceptorWork):
     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)
+            _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock=clock)
         elif not process_phase_failed:
             logger.info(f'Job {job.id}: adoption deferred (handled in _handle_work_error)')
         else:
@@ -1071,7 +1074,7 @@ def reattach_to_work_unit(job, work: ReceptorWork):
                 unit_status = {}
                 state_name = ''
             exit_code = _get_adoption_exit_code(unit_status, state_name)
-            _finalize_adopted_job(job, callback, exit_code, process_phase_failed)
+            _finalize_adopted_job(job, callback, exit_code, process_phase_failed, clock=clock)
     finally:
         receptor_job._receptor_release_work(work, getattr(res, 'status', 'error'))
         shutil.rmtree(private_data_dir, ignore_errors=True)

awx/main/tasks/system.py

@@ -679,10 +679,21 @@ def cluster_node_heartbeat(binder):
     Dispatcherd implementation.
     Uses Control API to get running tasks.
     """
+    run_heartbeat(binder)
+
+
+def run_heartbeat(binder, clock=None):
+    """One heartbeat of this instance.
+
+    clock returns the current time and defaults to django.utils.timezone.now. Tests pass
+    their own to judge peers or orphaned jobs as of any moment, or to give two simulated
+    controllers different clocks.
+    """
+    clock = clock or now
 
     # Run common instance management logic — ctl is the same receptor connection used for
     # mesh status; we reuse it for the job processing loop to avoid a second socket open.
-    this_inst, instance_list, lost_instances, _ctl = _heartbeat_instance_management()
+    this_inst, instance_list, lost_instances, _ctl = _heartbeat_instance_management(clock)
     if this_inst is None:
         return  # Early return case from instance management
 
@@ -706,7 +717,7 @@ def cluster_node_heartbeat(binder):
         return  # Failed to get task IDs, don't attempt reaping
 
     # One loop over all orphaned running jobs — adopt dispatched, reap undispatched.
-    ref_time = now()
+    ref_time = clock()
     logger.debug(f"Running job processing loop with {len(active_task_ids)} excluded UUIDs")
     _process_running_jobs(this_inst, active_task_ids, ref_time)
 
@@ -777,10 +788,10 @@ def _mesh_all_ready_nodes_visible(mesh_status):
     return True
 
 
-def _heartbeat_instance_management():
+def _heartbeat_instance_management(clock=None):
     """Common logic for heartbeat instance management."""
     logger.debug("Cluster node heartbeat task.")
-    nowtime = now()
+    nowtime = (clock or now)()
     instance_list = list(Instance.objects.filter(node_state__in=(Instance.States.READY, Instance.States.UNAVAILABLE, Instance.States.INSTALLED)))
     this_inst = None
 
@@ -1012,6 +1023,16 @@ def adopt_job_async(job_id):
     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()))
+
+
+def adopt_job(job_id, open_work, clock=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.
+    """
+    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')
@@ -1019,19 +1040,19 @@ def adopt_job_async(job_id):
 
     adoption_timeout = settings.HADR_JOB_ADOPTION_TIMEOUT
     last_event_time = job.get_event_queryset().aggregate(Max('created'))['created__max']
-    if adoption_deadline_passed(last_event_time, job.started, now(), adoption_timeout):
+    if adoption_deadline_passed(last_event_time, job.started, clock(), adoption_timeout):
         logger.error(f'Job {job.id} (unit={job.work_unit_id}) orphaned for >{adoption_timeout}s, failing')
         reaper.reap_job(job, 'failed', job_explanation='Job exceeded HADR_JOB_ADOPTION_TIMEOUT during controller restart')
         return
 
-    receptor_ctl = get_receptor_ctl()
+    work = open_work()
     try:
-        reattach_to_work_unit(job, ReceptorWork(receptor_ctl))
+        reattach_to_work_unit(job, work, clock=clock)
     except Exception:
         logger.exception(f'adopt_job_async: adoption failed for job {job.id} (unit={job.work_unit_id})')
     finally:
         try:
-            receptor_ctl.close()
+            work.close()
         except Exception:
             pass
 

awx/main/tests/fake_clock.py

@@ -0,0 +1,15 @@
+"""A clock tests can set and advance, for code that takes a `clock` callable."""
+
+from django.utils.timezone import now, timedelta
+
+
+class FakeClock:
+    def __init__(self, start=None):
+        self.current = start or now()
+
+    def __call__(self):
+        return self.current
+
+    def advance(self, seconds):
+        self.current += timedelta(seconds=seconds)
+        return self.current

awx/main/tests/fake_receptor.py

@@ -117,6 +117,9 @@ class FakeReceptorWork(ReceptorWork):
         self._unit(unit_id)
         del self.units[unit_id]
 
+    def close(self):
+        self.calls.append(('close',))
+
 
 class RecordingDispatcher:
     """Collects callback-receiver messages instead of sending them to Redis."""

awx/main/tests/functional/tasks/test_adoption_clock.py

@@ -0,0 +1,146 @@
+"""Heartbeat and adoption decisions as of a chosen time, using an injected clock.
+
+Each test sets the time instead of sleeping, so thresholds and clock skew between
+controllers are checked exactly and in milliseconds.
+"""
+
+from unittest import mock
+
+import pytest
+from django.utils.timezone import now, timedelta
+
+from awx.main.models import Instance, Job, JobEvent
+from awx.main.tasks.system import _heartbeat_instance_management, adopt_job
+from awx.main.tests.fake_clock import FakeClock
+from awx.main.tests.fake_receptor import FakeReceptorWork, FakeWorkUnit, RecordingDispatcher, make_events
+
+GRACE = 120  # CLUSTER_NODE_HEARTBEAT_PERIOD * CLUSTER_NODE_MISSED_HEARTBEAT_TOLERANCE
+TIMEOUT = 300  # HADR_JOB_ADOPTION_TIMEOUT
+UNIT = 'unit-1'
+
+
+@pytest.fixture
+def heartbeat_settings(settings):
+    settings.CLUSTER_HOST_ID = 'ctrl-0'
+    settings.AWX_AUTO_DEPROVISION_INSTANCES = False
+    settings.CLUSTER_NODE_HEARTBEAT_PERIOD = 60
+    settings.CLUSTER_NODE_MISSED_HEARTBEAT_TOLERANCE = 2
+    return settings
+
+
+def control_node(hostname, last_seen):
+    return Instance.objects.create(hostname=hostname, node_type='control', node_state='ready', last_seen=last_seen)
+
+
+def lost_hostnames(clock):
+    """Run ctrl-0's heartbeat at clock() and return the peers it judges lost."""
+    # ctrl-0 is alive: its own last heartbeat was written by its own clock.
+    Instance.objects.filter(hostname='ctrl-0').update(last_seen=clock())
+    ctl = mock.MagicMock()
+    ctl.simple_command.return_value = {'KnownConnectionCosts': {'ctrl-0': {'ctrl-1': 1}}, 'Advertisements': []}
+    with (
+        mock.patch('awx.main.tasks.system.get_receptor_ctl', return_value=ctl),
+        mock.patch('awx.main.tasks.system.inspect_execution_and_hop_nodes'),
+        mock.patch.object(Instance, 'local_health_check'),
+    ):
+        _, _, lost, _ = _heartbeat_instance_management(clock)
+    return sorted(inst.hostname for inst in lost)
+
+
+# ── Clock skew between controllers ───────────────────────────────────────────
+
+
+@pytest.mark.django_db
+def test_peer_is_lost_only_after_the_grace_period(heartbeat_settings):
+    t = now()
+    control_node('ctrl-0', t)
+    control_node('ctrl-1', t)
+    clock = FakeClock(t)
+
+    clock.advance(GRACE)
+    assert lost_hostnames(clock) == []
+    clock.advance(1)
+    assert lost_hostnames(clock) == ['ctrl-1']
+
+
+@pytest.mark.django_db
+def test_judge_with_a_fast_clock_finds_a_healthy_peer_lost(heartbeat_settings):
+    # ctrl-0's clock runs five minutes fast, so its own last_seen is five minutes ahead.
+    # ctrl-1 heartbeated just now by the correct clock and is healthy.
+    t = now()
+    control_node('ctrl-0', t + timedelta(minutes=5))
+    control_node('ctrl-1', t)
+
+    assert lost_hostnames(FakeClock(t + timedelta(minutes=5))) == ['ctrl-1']
+
+
+@pytest.mark.django_db
+def test_dead_peer_with_a_fast_clock_is_found_late(heartbeat_settings):
+    # ctrl-1 ran five minutes fast and then died. Its last_seen is in the future, so a
+    # correct judge only finds it lost once real time passes last_seen plus the grace period.
+    t = now()
+    control_node('ctrl-0', t)
+    control_node('ctrl-1', t + timedelta(minutes=5))
+    clock = FakeClock(t)
+
+    clock.advance(5 * 60 + GRACE)
+    assert lost_hostnames(clock) == []
+    clock.advance(1)
+    assert lost_hostnames(clock) == ['ctrl-1']
+
+
+# ── Adoption deadline, end to end ─────────────────────────────────────────────
+
+
+@pytest.fixture
+def adoption_env(settings, tmp_path):
+    settings.AWX_ISOLATION_BASE_PATH = str(tmp_path)
+    settings.HADR_JOB_ADOPTION_TIMEOUT = TIMEOUT
+    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 orphan(me_inst):
+    """A running job whose last saved event was written at the returned time."""
+    started = now() - timedelta(hours=1)
+    job = Job.objects.create(controller_node=me_inst.hostname, status='running', work_unit_id=UNIT, started=started)
+    JobEvent.objects.create(job=job, counter=1, event='playbook_on_start', job_created=job.created)
+    last_event = now() - timedelta(minutes=10)
+    JobEvent.objects.filter(job=job).update(created=last_event)
+    return job, last_event
+
+
+@pytest.mark.django_db
+def test_orphan_within_the_deadline_is_adopted(adoption_env, orphan):
+    job, last_event = orphan
+    work = FakeReceptorWork({UNIT: FakeWorkUnit(make_events(10))})
+    clock = FakeClock(last_event + timedelta(seconds=TIMEOUT - 1))
+
+    adopt_job(job.id, open_work=lambda: work, clock=clock)
+
+    job.refresh_from_db()
+    assert job.status == 'successful'
+    assert job.finished == clock()
+    assert RecordingDispatcher.counters() == list(range(2, 11))
+    assert ('close',) in work.calls
+
+
+@pytest.mark.django_db
+def test_orphan_past_the_deadline_is_failed_without_contacting_receptor(adoption_env, orphan):
+    job, last_event = orphan
+    opened = mock.Mock()
+
+    adopt_job(job.id, open_work=opened, clock=FakeClock(last_event + timedelta(seconds=TIMEOUT + 1)))
+
+    job.refresh_from_db()
+    assert job.status == 'failed'
+    assert 'HADR_JOB_ADOPTION_TIMEOUT' in job.job_explanation
+    opened.assert_not_called()