Fix prefill crash

CHANGELOG.md

@@ -18,6 +18,7 @@
 * Job and worker IDs are no longer reused after the server is restarted from a pruned journal
 * Fixes invalid counters when journal restores job with multiple submits
 * Fixed server crash when restoring a journal containing a task that failed before it started
+* Fixed server crash when a lost worker held prefilled tasks together with a task of a higher priority
 
 
 ## v0.27.0

crates/tako/src/internal/server/reactor.rs

@@ -82,6 +82,17 @@ pub(crate) fn on_remove_worker(
     let mut retracted = Vec::new();
     match worker.assignment() {
         WorkerAssignment::Sn(sn) => {
+            // The prefilled tasks have to be processed first. Adding a ready task below may
+            // dispose a prefill with a lower priority, and the disposed tasks are retracted,
+            // which is not possible for tasks of the removed worker.
+            for task_id in &sn.prefilled_tasks {
+                let task = task_map.get_task_mut(*task_id);
+                task.increment_instance_id();
+                task.state = TaskRuntimeState::Waiting { unfinished_deps: 0 };
+                task_queues
+                    .get_mut(task.resource_rq_id)
+                    .move_prefilled_task_to_ready(*task_id);
+            }
             for task_id in &sn.assigned_tasks {
                 let task = task_map.get_task_mut(*task_id);
                 if task.is_sn_running() {
@@ -95,14 +106,6 @@ pub(crate) fn on_remove_worker(
                 task.increment_instance_id();
                 task_queues.add_ready_task(task, &mut retracted);
             }
-            for task_id in &sn.prefilled_tasks {
-                let task = task_map.get_task_mut(*task_id);
-                task.increment_instance_id();
-                task.state = TaskRuntimeState::Waiting { unfinished_deps: 0 };
-                task_queues
-                    .get_mut(task.resource_rq_id)
-                    .move_prefilled_task_to_ready(*task_id);
-            }
         }
         WorkerAssignment::Mn(mn) => {
             let task = task_map.get_task_mut(mn.task_id);

crates/tako/src/internal/tests/test_reactor.rs

@@ -885,6 +885,60 @@ fn test_prefill_worker_lost() {
     rt.sanity_check();
 }
 
+/// Returning a higher-priority task of the lost worker to the queue disposes the lower-priority
+/// prefill. The prefill of the lost worker has to be resolved without retracting it, while the
+/// prefill of the remaining workers has to be retracted as usual.
+#[test]
+fn test_prefill_worker_lost_with_higher_priority_task() {
+    let mut rt = TestEnv::new();
+    rt.set_scheduler_config(SchedulerConfig {
+        proactive_filling_reserve: 0,
+        proactive_filling_max: 1,
+        ..Default::default()
+    });
+    let workers = rt.new_workers(2, &WorkerBuilder::new(2));
+    let high = rt.new_task(&TaskBuilder::new().user_priority(10));
+    rt.schedule();
+    let lost = match rt.task(high).state {
+        TaskRuntimeState::Assigned { worker_id, .. } => worker_id,
+        _ => panic!("Invalid state"),
+    };
+    let other = *workers.iter().find(|w| **w != lost).unwrap();
+
+    let low = rt.new_tasks(5, &TaskBuilder::new());
+    rt.schedule();
+    let prefilled_on = |rt: &TestEnv, worker: WorkerId| {
+        low.iter()
+            .copied()
+            .find(|t| {
+                matches!(rt.task(*t).state,
+                    TaskRuntimeState::Prefilled { worker_id } if worker_id == worker)
+            })
+            .unwrap()
+    };
+    let prefilled_lost = prefilled_on(&rt, lost);
+    let prefilled_other = prefilled_on(&rt, other);
+
+    let mut comm = TestComm::new();
+    on_remove_worker(rt.core(), &mut comm, lost, LostWorkerReason::HeartbeatLost);
+    comm.check_need_scheduling();
+    match &comm.take_worker_msgs(other, 1)[0] {
+        ToWorkerMessage::RetractTasks(ts) => assert_eq!(ts.ids, vec![prefilled_other]),
+        _ => panic!("Invalid worker msg"),
+    }
+    comm.take_broadcasts(1);
+    comm.client.take_lost_workers();
+    comm.emptiness_check();
+
+    assert!(rt.task(high).is_waiting());
+    assert!(rt.task(prefilled_lost).is_waiting());
+    assert!(matches!(
+        rt.task(prefilled_other).state,
+        TaskRuntimeState::Retracting { worker_id } if worker_id == other
+    ));
+    rt.sanity_check();
+}
+
 #[test]
 fn test_prefill_started_on_same_worker() {
     let mut rt = TestEnv::new();