diff --git a/CHANGELOG.md b/CHANGELOG.md index af068b77f..eee064965 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/crates/tako/src/internal/server/reactor.rs b/crates/tako/src/internal/server/reactor.rs index 7e34c0843..0ade5071f 100644 --- a/crates/tako/src/internal/server/reactor.rs +++ b/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); diff --git a/crates/tako/src/internal/tests/test_reactor.rs b/crates/tako/src/internal/tests/test_reactor.rs index edd9da0bb..ccbc92985 100644 --- a/crates/tako/src/internal/tests/test_reactor.rs +++ b/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();