1use std::collections::{HashMap, HashSet};2use std::sync::Arc;34use anyhow::{Context, Result, bail};5use futures::stream::{FuturesUnordered, StreamExt};6use tokio::sync::{Semaphore, broadcast};7use tracing::{debug, info, instrument, warn};89use crate::drv::DrvGraph;1011#[derive(Clone, Debug)]12pub enum BuildEvent {13 SubstitutePrepassStarted {14 paths: usize,15 },16 SubstitutePrepassFinished {17 satisfied: usize,18 },19 DrvStarted {20 drv_path: String,21 name: String,22 wanted: Vec<String>,23 },24 DrvSkipped {25 drv_path: String,26 name: String,27 },28 DrvFinished {29 drv_path: String,30 name: String,31 },32 DrvFailed {33 drv_path: String,34 name: String,35 error: String,36 },37 DrvCancelled {38 drv_path: String,39 name: String,40 failed_dep: String,41 },42}4344pub struct Scheduler {45 parallelism: usize,46 events: broadcast::Sender<BuildEvent>,47}4849impl Scheduler {50 pub fn new(parallelism: usize) -> Self {51 let parallelism = parallelism.max(1);52 let (events, _) = broadcast::channel(1024);53 Self {54 parallelism,55 events,56 }57 }5859 pub fn subscribe(&self) -> broadcast::Receiver<BuildEvent> {60 self.events.subscribe()61 }6263 #[instrument(name = "scheduler", skip(self, graph), fields(root = %graph.root, nodes = graph.nodes.len()))]64 pub async fn run(&self, graph: Arc<DrvGraph>, root_outputs: Vec<String>) -> Result<()> {65 let wanted = graph.wanted_outputs(&root_outputs);6667 self.substitute_prepass(&graph, &wanted).await?;68 self.build_topo(&graph, wanted).await69 }7071 async fn substitute_prepass(72 &self,73 graph: &DrvGraph,74 wanted: &HashMap<String, Vec<String>>,75 ) -> Result<()> {76 let paths = collect_substitute_paths(graph, wanted);77 if paths.is_empty() {78 return Ok(());79 }80 let _ = self81 .events82 .send(BuildEvent::SubstitutePrepassStarted { paths: paths.len() });83 debug!("substitute pre-pass: {} paths", paths.len());8485 let satisfied = tokio::task::spawn_blocking(move || crate::substitute_paths(&paths))86 .await87 .expect("substitute pre-pass task should not panic")?;8889 let _ = self.events.send(BuildEvent::SubstitutePrepassFinished {90 satisfied: satisfied.len(),91 });92 Ok(())93 }9495 async fn build_topo(96 &self,97 graph: &Arc<DrvGraph>,98 wanted: HashMap<String, Vec<String>>,99 ) -> Result<()> {100 let mut indeg: HashMap<String, usize> = graph101 .nodes102 .iter()103 .map(|(k, n)| (k.clone(), n.input_drvs.len()))104 .collect();105 let mut dependents: HashMap<String, Vec<String>> = HashMap::new();106 for (path, node) in &graph.nodes {107 for dep in node.input_drvs.keys() {108 dependents109 .entry(dep.clone())110 .or_default()111 .push(path.clone());112 }113 }114115 let sem = Arc::new(Semaphore::new(self.parallelism));116 let mut ready: Vec<String> = indeg117 .iter()118 .filter(|(_, d)| **d == 0)119 .map(|(k, _)| k.clone())120 .collect();121 let mut in_flight = FuturesUnordered::new();122 let mut failed: HashMap<String, String> = HashMap::new();123 124 let mut tainted: HashMap<String, String> = HashMap::new();125126 loop {127 let batch: Vec<String> = std::mem::take(&mut ready);128 for path in batch {129 if let Some(failed_dep) = tainted.get(&path) {130 let name = graph131 .nodes132 .get(&path)133 .map(|n| n.name.clone())134 .unwrap_or_default();135 let _ = self.events.send(BuildEvent::DrvCancelled {136 drv_path: path.clone(),137 name,138 failed_dep: failed_dep.clone(),139 });140 propagate_done(&dependents, &mut indeg, &mut ready, &path);141 continue;142 }143144 let sem = sem.clone();145 let events = self.events.clone();146 let graph = graph.clone();147 let wanted_here = wanted.get(&path).cloned().unwrap_or_default();148 in_flight.push(tokio::spawn(async move {149 let _permit = sem.acquire_owned().await.expect("semaphore not closed");150 let node = graph151 .nodes152 .get(&path)153 .expect("ready node must be in graph")154 .clone();155 let name = node.name.clone();156157 let all_valid = !wanted_here.is_empty()158 && wanted_here.iter().all(|o| {159 node.outputs160 .get(o)161 .map(|p| crate::is_valid_path(p).unwrap_or(false))162 .unwrap_or(false)163 });164 if all_valid {165 let _ = events.send(BuildEvent::DrvSkipped {166 drv_path: path.clone(),167 name: name.clone(),168 });169 return (path, name, Ok::<(), anyhow::Error>(()));170 }171172 let _ = events.send(BuildEvent::DrvStarted {173 drv_path: path.clone(),174 name: name.clone(),175 wanted: wanted_here.clone(),176 });177178 let path_for_build = path.clone();179 let res = tokio::task::spawn_blocking(move || {180 crate::build_drv_outputs(&path_for_build, &wanted_here)181 })182 .await183 .expect("build task should not panic");184185 match res {186 Ok(_) => {187 let _ = events.send(BuildEvent::DrvFinished {188 drv_path: path.clone(),189 name: name.clone(),190 });191 (path, name, Ok(()))192 }193 Err(e) => {194 let msg = format!("{e:#}");195 let _ = events.send(BuildEvent::DrvFailed {196 drv_path: path.clone(),197 name: name.clone(),198 error: msg,199 });200 (path, name, Err(e))201 }202 }203 }));204 }205206 let Some(joined) = in_flight.next().await else {207 break;208 };209 let (finished, _name, res) = match joined {210 Ok(t) => t,211 Err(e) => bail!("scheduler task panicked: {e}"),212 };213 match res {214 Ok(()) => {215 propagate_done(&dependents, &mut indeg, &mut ready, &finished);216 }217 Err(e) => {218 failed.insert(finished.clone(), format!("{e:#}"));219 mark_tainted(&dependents, &finished, &mut tainted);220 propagate_done(&dependents, &mut indeg, &mut ready, &finished);221 }222 }223 }224225 let stuck: Vec<_> = indeg226 .iter()227 .filter(|(_, d)| **d != 0)228 .map(|(k, _)| k.as_str())229 .collect();230 if !stuck.is_empty() {231 warn!(232 "scheduler finished with {} nodes still pending (loop?)",233 stuck.len()234 );235 }236237 if failed.is_empty() {238 info!("scheduler completed");239 Ok(())240 } else {241 let mut report = format!("{} drv(s) failed to build:", failed.len());242 let mut sorted: Vec<_> = failed.iter().collect();243 sorted.sort_by(|a, b| a.0.cmp(b.0));244 for (path, err) in sorted {245 let name = graph246 .nodes247 .get(path)248 .map(|n| n.name.as_str())249 .unwrap_or("?");250 let chain = path_to_root(graph, path);251 report.push_str(&format!(252 "\n\n {name} ({path}):\n {err}\n needed by: {}",253 chain.join(" => "),254 ));255 }256 Err(anyhow::anyhow!(report))257 }258 }259}260261fn propagate_done(262 dependents: &HashMap<String, Vec<String>>,263 indeg: &mut HashMap<String, usize>,264 ready: &mut Vec<String>,265 finished: &str,266) {267 if let Some(deps) = dependents.get(finished) {268 for d in deps {269 let entry = indeg.get_mut(d).expect("dependent must have indeg");270 *entry = entry.saturating_sub(1);271 if *entry == 0 {272 ready.push(d.clone());273 }274 }275 }276}277278fn mark_tainted(279 dependents: &HashMap<String, Vec<String>>,280 failed: &str,281 tainted: &mut HashMap<String, String>,282) {283 let mut queue: Vec<String> = dependents.get(failed).cloned().unwrap_or_default();284 while let Some(node) = queue.pop() {285 if tainted286 .entry(node.clone())287 .or_insert_with(|| failed.to_owned())288 == failed289 {290 if let Some(deps) = dependents.get(&node) {291 for d in deps {292 if !tainted.contains_key(d) {293 queue.push(d.clone());294 }295 }296 }297 }298 }299}300301fn path_to_root(graph: &DrvGraph, from: &str) -> Vec<String> {302 let mut dependents: HashMap<&str, Vec<&str>> = HashMap::new();303 for (path, node) in &graph.nodes {304 for dep in node.input_drvs.keys() {305 dependents306 .entry(dep.as_str())307 .or_default()308 .push(path.as_str());309 }310 }311312 let mut chain: Vec<String> = vec![node_name(graph, from)];313 let mut cur = from;314 let mut seen: HashSet<&str> = HashSet::new();315 seen.insert(cur);316 while cur != graph.root.as_str() {317 let Some(next) = dependents.get(cur).and_then(|v| v.first().copied()) else {318 break;319 };320 if !seen.insert(next) {321 break;322 }323 chain.push(node_name(graph, next));324 cur = next;325 }326 chain327}328329fn node_name(graph: &DrvGraph, path: &str) -> String {330 graph331 .nodes332 .get(path)333 .map(|n| n.name.clone())334 .unwrap_or_else(|| path.to_owned())335}336337fn collect_substitute_paths(338 graph: &DrvGraph,339 wanted: &HashMap<String, Vec<String>>,340) -> Vec<String> {341 let mut paths: HashSet<String> = HashSet::new();342 for node in graph.nodes.values() {343 for src in &node.input_srcs {344 paths.insert(src.clone());345 }346 }347 for (path, outs) in wanted {348 let Some(node) = graph.nodes.get(path) else {349 continue;350 };351 for o in outs {352 if let Some(p) = node.outputs.get(o) {353 paths.insert(p.clone());354 }355 }356 }357 let mut v: Vec<_> = paths.into_iter().collect();358 v.sort();359 v360}361362363364365pub fn build_graph_sync(graph: Arc<DrvGraph>, root_outputs: Vec<String>) -> Result<()> {366 let parallelism = std::thread::available_parallelism()367 .map(|p| p.get())368 .unwrap_or(4);369 let scheduler = Scheduler::new(parallelism);370 crate::await_in_nix(async move { scheduler.run(graph, root_outputs).await })371 .context("scheduler run")372}