diff --git a/crates/tako/src/internal/scheduler/gap.rs b/crates/tako/src/internal/scheduler/gap.rs index 063e233f3..4a85ef995 100644 --- a/crates/tako/src/internal/scheduler/gap.rs +++ b/crates/tako/src/internal/scheduler/gap.rs @@ -244,4 +244,25 @@ mod tests { let t2 = rt.new_task_cpus(1); assert_eq!(compute_gap(&mut rt, t1, t2, w), 2); } + + #[test] + fn test_compute_gap_resource_the_worker_lacks() { + // A worker's resource vector spans only the resources it declared; a request naming a + // resource that other workers have must not index past its end (server panic, 2026-09). + let mut rt = TestEnv::new(); + rt.new_named_resource("foo"); + let w = rt.new_worker(&WorkerBuilder::new(4)); // cpus only + let t2 = rt.new_task_cpus(1); + let t1 = rt.new_task(&TaskBuilder::new().cpus(1).add_resource(1, 1)); + // nothing of it can run here, so nothing is held back: the whole worker stays for t2 + assert_eq!(compute_gap(&mut rt, t1, t2, w), 4); + let t1 = rt.new_task( + &TaskBuilder::new() + .cpus(1) + .add_resource(1, 1) + .next_variant() + .cpus(3), + ); + assert_eq!(compute_gap(&mut rt, t1, t2, w), 1); + } } diff --git a/crates/tako/src/internal/server/workerload.rs b/crates/tako/src/internal/server/workerload.rs index cdcbb0254..947b98944 100644 --- a/crates/tako/src/internal/server/workerload.rs +++ b/crates/tako/src/internal/server/workerload.rs @@ -153,39 +153,39 @@ impl WorkerResources { .sum::() } + /// Subtract `amount` (or everything when the request asks for all) of one resource. + /// The vector only spans the resources this worker declared (`from_description`), while a + /// request may name a resource that exists elsewhere in the cluster: for such a resource the + /// worker holds nothing, so there is nothing to subtract. `get` already treats it as zero; + /// indexing directly panicked the server in the scheduler's gap computation once a + /// prioritized multi-variant request met a worker without the variant's resource. + fn subtract(&mut self, resource_id: ResourceId, amount: Option) { + if let Some(slot) = self.n_resources.get_mut(resource_id.as_usize()) { + *slot = match amount { + Some(amount) => slot.saturating_sub(amount), + None => ResourceAmount::ZERO, + }; + } + } + pub fn remove(&mut self, rq: &ResourceRequest) { for entry in rq.entries() { - if let Some(amount) = entry.request.amount_or_none_if_all() { - self.n_resources[entry.resource_id] = - self.n_resources[entry.resource_id].saturating_sub(amount); - } else { - self.n_resources[entry.resource_id] = ResourceAmount::ZERO; - } + self.subtract(entry.resource_id, entry.request.amount_or_none_if_all()); } } pub fn remove_multiple(&mut self, rq: &ResourceRequest, n: u32) { for entry in rq.entries() { - if let Some(amount) = entry.request.amount_or_none_if_all() { - let a = amount.times(n); - self.n_resources[entry.resource_id] = - self.n_resources[entry.resource_id].saturating_sub(a); - } else { - self.n_resources[entry.resource_id] = ResourceAmount::ZERO; - } + let amount = entry.request.amount_or_none_if_all().map(|a| a.times(n)); + self.subtract(entry.resource_id, amount); } } pub fn remove_multiple_masked(&mut self, rq: &ResourceRequest, n: u32, r_id: ResourceId) { for entry in rq.entries() { if entry.resource_id == r_id { - if let Some(amount) = entry.request.amount_or_none_if_all() { - let a = amount.times(n); - self.n_resources[entry.resource_id] = - self.n_resources[entry.resource_id].saturating_sub(a); - } else { - self.n_resources[entry.resource_id] = ResourceAmount::ZERO; - } + let amount = entry.request.amount_or_none_if_all().map(|a| a.times(n)); + self.subtract(entry.resource_id, amount); return; } }