iris/worker: protect freshly-submitted tasks from StartTasks→PollTasks race (#5043)

## Summary

Fixes #5041.

When the iris controller dispatches a task via `StartTasks` and polls
the worker for state via `PollTasks` before its own view of
`expected_tasks` has caught up, the worker treated the just-submitted
task as "unexpected" and killed it. That kill rolled up the workers-pool
job to `JOB_STATE_KILLED` and cascaded the surviving tasks with
`error="Job was terminated."`, surfacing in zephyr as the misleading
`"Worker job terminated permanently… Workers likely crashed"` abort.

`handle_heartbeat` already guarded against this race by passing
`extra_expected_keys` for the tasks it had just submitted in that RPC.
`handle_poll_tasks` did not — the PollTasks path has no "tasks_to_run"
field because StartTasks is a separate RPC — so freshly-submitted tasks
had no protection.

### Approach

- Track recent submissions on the worker: `submit_task` now records
`(task_id, attempt_id) -> monotonic_time` in `self._recent_submissions`.
- `_reconcile_expected_tasks` treats keys within a 30s grace window as
expected.
- Stale entries are pruned on each reconciliation so the dict stays
bounded.
- `_reset_worker_state` clears the tracking alongside `self._tasks`.

This fixes both `handle_poll_tasks` and the more general case where
heartbeat-submitted tasks still need race protection on the following
tick. The bespoke `extra_keys` set in `handle_heartbeat` is gone since
`submit_task` now populates `_recent_submissions` for all entry points.

Author: ravwojdyla-agent

lib/iris/src/iris/cluster/worker/worker.py

@@ -125,6 +125,13 @@ def worker_config_from_proto(
 class Worker:
     """Unified worker managing all components and lifecycle."""
 
+    # Grace period during which a freshly-submitted task is treated as
+    # "expected" by reconciliation, even if it hasn't yet appeared in the
+    # controller's expected_tasks list. Protects against the StartTasks →
+    # PollTasks race where the controller polls before its internal view
+    # catches up with the task it just assigned.
+    _RECENT_SUBMISSION_GRACE_SECONDS = 30.0
+
     def __init__(
         self,
         config: WorkerConfig,
@@ -176,6 +183,11 @@ def __init__(
         # Task state: maps (task_id, attempt_id) -> TaskAttempt.
         # Preserves all attempts so logs for historical attempts remain accessible.
         self._tasks: dict[tuple[str, int], TaskAttempt] = {}
+        # Freshly-submitted tasks -> monotonic submission time. Used by
+        # reconciliation to grant a grace period before a task becomes
+        # eligible for "unexpected, kill" if the controller hasn't yet
+        # listed it in expected_tasks. See _RECENT_SUBMISSION_GRACE_SECONDS.
+        self._recent_submissions: dict[tuple[str, int], float] = {}
         self._lock = threading.Lock()
 
         self._host_metrics = HostMetricsCollector(disk_path=str(self._cache_dir))
@@ -586,6 +598,7 @@ def _reset_worker_state(self) -> None:
         # Clear task tracking
         with self._lock:
             self._tasks.clear()
+            self._recent_submissions.clear()
 
         # Replace the task thread container so new tasks get a fresh group.
         self._task_threads = self._threads.create_child("tasks")
@@ -701,6 +714,7 @@ def submit_task(self, request: job_pb2.RunTaskRequest) -> str:
 
         with self._lock:
             self._tasks[key] = attempt
+            self._recent_submissions[key] = time.monotonic()
 
         # Start execution in a monitored non-daemon thread. When stop() is called,
         # the on_stop callback kills the container so attempt.run() exits promptly.
@@ -791,19 +805,33 @@ def _missing_task_status(task_id: str, expected_attempt_id: int) -> job_pb2.Work
             finished_at=timestamp_to_proto(Timestamp.now()),
         )
 
+    def _prune_and_get_recent_submission_keys(self) -> set[tuple[str, int]]:
+        """Return keys submitted within the grace window, pruning stale entries.
+
+        Caller must hold ``self._lock``. Stale entries (older than the grace
+        window) are removed from ``self._recent_submissions`` so the dict
+        doesn't grow unbounded.
+        """
+        now = time.monotonic()
+        cutoff = now - self._RECENT_SUBMISSION_GRACE_SECONDS
+        stale = [key for key, ts in self._recent_submissions.items() if ts < cutoff]
+        for key in stale:
+            del self._recent_submissions[key]
+        return set(self._recent_submissions)
+
     def _reconcile_expected_tasks(
         self,
         expected_entries,
-        extra_expected_keys: set[tuple[str, int]] | None = None,
     ) -> tuple[list[job_pb2.WorkerTaskStatus], list[tuple[str, int]]]:
         """Build status entries for expected tasks; collect non-terminal local tasks
         not in the expected set as targets to kill.
 
         Caller must hold ``self._lock``.
 
-        ``extra_expected_keys`` keeps freshly-submitted tasks (e.g. ``tasks_to_run``
-        on the legacy heartbeat) from being killed when they aren't yet in the
-        controller's expected set.
+        Freshly-submitted tasks (``self._recent_submissions``) are protected
+        from reconciliation kills via the grace window, which covers the
+        StartTasks → PollTasks race where the controller polls before its
+        internal view catches up with a task it just assigned.
         """
         tasks: list[job_pb2.WorkerTaskStatus] = []
         expected_keys: set[tuple[str, int]] = set()
@@ -817,8 +845,7 @@ def _reconcile_expected_tasks(
                 tasks.append(self._missing_task_status(task_id, expected_attempt_id))
             else:
                 tasks.append(self._encode_task_status(task, task_id))
-        if extra_expected_keys:
-            expected_keys |= extra_expected_keys
+        expected_keys |= self._prune_and_get_recent_submission_keys()
         tasks_to_kill: list[tuple[str, int]] = []
         for key, task in self._tasks.items():
             if key not in expected_keys and task.status not in self._TERMINAL_STATES:
@@ -851,8 +878,12 @@ def handle_heartbeat(self, request: job_pb2.HeartbeatRequest) -> job_pb2.Heartbe
            found on worker"). This happens when the worker has reset its state
            (_tasks.clear() in _reset_worker_state) between heartbeats — from
            the controller's perspective this is equivalent to a worker restart.
-        4. Kill unexpected tasks — any task in self._tasks that is NOT in
-           expected_tasks or tasks_to_run is killed (controller no longer wants it)
+        4. Kill unexpected tasks — any non-terminal task in self._tasks that is
+           NOT in expected_tasks and is not within the recent-submission grace
+           window is killed (controller no longer wants it). The grace window
+           keeps tasks just submitted via StartTasks or this heartbeat's
+           tasks_to_run from being killed when the controller hasn't yet
+           listed them in expected_tasks.
 
         The ordering guarantee between steps 1 and 3 is critical: a task that
         appears in both tasks_to_run and expected_tasks (which is always the case
@@ -891,12 +922,11 @@ def handle_heartbeat(self, request: job_pb2.HeartbeatRequest) -> job_pb2.Heartbe
                         logger.warning("Heartbeat: failed to kill task %s: %s", task_id, e)
 
             with slow_log(logger, "heartbeat reconciliation", threshold_ms=200):
-                # tasks_to_run was just submitted above; carry those keys so a
-                # newly-assigned task isn't killed if the controller hasn't yet
-                # listed it in expected_tasks.
-                extra_keys = {(r.task_id, r.attempt_id) for r in request.tasks_to_run}
+                # tasks_to_run was just submitted above; those keys live in
+                # self._recent_submissions and are protected from the race by
+                # _reconcile_expected_tasks' grace-window logic.
                 with self._lock:
-                    tasks, tasks_to_kill = self._reconcile_expected_tasks(request.expected_tasks, extra_keys)
+                    tasks, tasks_to_kill = self._reconcile_expected_tasks(request.expected_tasks)
 
                 # Kill removed tasks asynchronously outside lock to avoid deadlock
                 for task_id, attempt_id in tasks_to_kill:
@@ -959,7 +989,12 @@ def handle_stop_tasks(self, request: worker_pb2.Worker.StopTasksRequest) -> work
         return worker_pb2.Worker.StopTasksResponse()
 
     def handle_poll_tasks(self, request: worker_pb2.Worker.PollTasksRequest) -> worker_pb2.Worker.PollTasksResponse:
-        """Report status of expected tasks and kill unexpected tasks."""
+        """Report status of expected tasks and kill unexpected tasks.
+
+        Freshly-submitted tasks (via StartTasks) are protected from the
+        StartTasks → PollTasks race by the recent-submission grace window
+        applied in _reconcile_expected_tasks.
+        """
         with self._lock:
             tasks, tasks_to_kill = self._reconcile_expected_tasks(request.expected_tasks)
         for task_id, attempt_id in tasks_to_kill:

lib/iris/tests/cluster/worker/test_worker.py

@@ -27,6 +27,7 @@
 from iris.cluster.worker.service import WorkerServiceImpl
 from iris.cluster.worker.worker import Worker, WorkerConfig
 from iris.rpc import job_pb2
+from iris.rpc import worker_pb2
 from rigging.timing import Duration
 from iris.cluster.worker.worker_types import LogLine
 from tests.cluster.worker.conftest import (
@@ -479,6 +480,11 @@ def slow_stop(force=False):
     task = mock_worker.get_task(task_id_wire)
     wait_for_condition(lambda: task.status == job_pb2.TASK_STATE_RUNNING)
 
+    # Clear recent-submissions tracking to simulate the task having been
+    # around long enough for the grace window to have elapsed; this test
+    # exercises reconciliation-driven kill, not grace-window protection.
+    mock_worker._recent_submissions.clear()
+
     # Send heartbeat with empty expected_tasks -- the worker should kill
     # the running task because it's no longer expected
     heartbeat_req = job_pb2.HeartbeatRequest(expected_tasks=[])
@@ -494,6 +500,77 @@ def slow_stop(force=False):
     assert task.status == job_pb2.TASK_STATE_KILLED
 
 
+def test_poll_tasks_grace_window_protects_freshly_submitted_task(mock_worker, mock_runtime):
+    """PollTasks must not kill a task submitted moments before the controller polls.
+
+    Reproduces the StartTasks → PollTasks race from iris #5041: the controller
+    dispatches a task via StartTasks but polls before its own expected_tasks view
+    includes the new task. Without the grace window, the worker would read the
+    task as "unexpected" and kill it, cascading the whole pool to KILLED.
+    """
+    mock_handle = create_mock_container_handle(status_sequence=[ContainerStatus(phase=ContainerPhase.RUNNING)] * 1000)
+    mock_runtime.create_container = Mock(return_value=mock_handle)
+
+    task_id_wire = JobName.root("test-user", "poll-race").task(0).to_wire()
+    request = create_run_task_request(task_id=task_id_wire)
+    mock_worker.submit_task(request)
+
+    task = mock_worker.get_task(task_id_wire)
+    wait_for_condition(lambda: task.status == job_pb2.TASK_STATE_RUNNING)
+
+    # Controller polls with the just-submitted task missing from expected_tasks
+    # (race: controller hasn't reconciled its own StartTasks response yet).
+    mock_worker.handle_poll_tasks(worker_pb2.Worker.PollTasksRequest(expected_tasks=[]))
+
+    # The task must not have been marked for kill.
+    assert task.should_stop is False
+    assert task.status == job_pb2.TASK_STATE_RUNNING
+
+    # Clean up.
+    mock_worker.kill_task(task_id_wire)
+    task.thread.join(timeout=15.0)
+
+
+def test_poll_tasks_kills_task_outside_grace_window(mock_worker, mock_runtime):
+    """Once the grace window has elapsed, reconciliation resumes killing unexpected tasks."""
+    mock_handle = create_mock_container_handle(status_sequence=[ContainerStatus(phase=ContainerPhase.RUNNING)] * 1000)
+    mock_runtime.create_container = Mock(return_value=mock_handle)
+
+    task_id_wire = JobName.root("test-user", "poll-post-grace").task(0).to_wire()
+    request = create_run_task_request(task_id=task_id_wire)
+    mock_worker.submit_task(request)
+
+    task = mock_worker.get_task(task_id_wire)
+    wait_for_condition(lambda: task.status == job_pb2.TASK_STATE_RUNNING)
+
+    # Simulate grace window elapsing by clearing recent-submissions tracking.
+    mock_worker._recent_submissions.clear()
+
+    mock_worker.handle_poll_tasks(worker_pb2.Worker.PollTasksRequest(expected_tasks=[]))
+
+    assert task.should_stop is True
+    task.thread.join(timeout=15.0)
+    assert task.status == job_pb2.TASK_STATE_KILLED
+
+
+def test_recent_submissions_prune_removes_stale_entries(mock_worker):
+    """Stale recent-submission entries are pruned to keep the dict bounded."""
+    key_fresh = ("task-fresh", 0)
+    key_stale = ("task-stale", 0)
+    grace = mock_worker._RECENT_SUBMISSION_GRACE_SECONDS
+    now = time.monotonic()
+    # now - (grace + 1): clearly older than the window -> should be pruned
+    mock_worker._recent_submissions[key_stale] = now - (grace + 1)
+    mock_worker._recent_submissions[key_fresh] = now
+
+    with mock_worker._lock:
+        recent = mock_worker._prune_and_get_recent_submission_keys()
+
+    assert key_fresh in recent
+    assert key_stale not in recent
+    assert key_stale not in mock_worker._recent_submissions
+
+
 def test_kill_nonexistent_task(mock_worker):
     """Test killing a nonexistent task returns False."""
     result = mock_worker.kill_task(JobName.root("test-user", "nonexistent-task").task(0).to_wire())