1use std::collections::{HashMap, HashSet};2use std::mem;3use std::sync::Arc;45use anyhow::{Context, Result, bail};6use camino::{Utf8Path, Utf8PathBuf};7use futures::stream::{FuturesUnordered, StreamExt};8use tokio::sync::{Semaphore, broadcast};9use tokio::task::spawn_blocking;10use tracing::{debug, info, instrument, warn};1112use crate::drv::DrvGraph;13use crate::{Store, eval_store};1415#[derive(Clone, Debug)]16pub enum BuildEvent {17 SubstitutePrepassStarted {18 paths: usize,19 },20 SubstitutePrepassFinished {21 satisfied: usize,22 },23 DrvStarted {24 drv_path: Utf8PathBuf,25 name: String,26 wanted: Vec<String>,27 },28 DrvSkipped {29 drv_path: Utf8PathBuf,30 name: String,31 },32 DrvFinished {33 drv_path: Utf8PathBuf,34 name: String,35 },36 DrvFailed {37 drv_path: Utf8PathBuf,38 name: String,39 error: String,40 },41 DrvCancelled {42 drv_path: Utf8PathBuf,43 name: String,44 failed_dep: Utf8PathBuf,45 },46}4748pub struct Scheduler {49 store: Arc<Store>,50 parallelism: usize,51 events: broadcast::Sender<BuildEvent>,52}5354impl Scheduler {55 pub fn new(parallelism: usize) -> Self {56 let parallelism = parallelism.max(1);57 let (events, _) = broadcast::channel(1024);58 Self {59 store: eval_store(),60 parallelism,61 events,62 }63 }6465 pub fn subscribe(&self) -> broadcast::Receiver<BuildEvent> {66 self.events.subscribe()67 }6869 #[instrument(name = "scheduler", skip(self, graph), fields(root = %graph.root, nodes = graph.nodes.len()))]70 pub async fn run(&self, graph: Arc<DrvGraph>, root_outputs: Vec<String>) -> Result<()> {71 let wanted = graph.wanted_outputs(&root_outputs);7273 self.substitute_prepass(&graph, &wanted).await?;74 self.build_topo(&graph, wanted).await75 }7677 async fn substitute_prepass(78 &self,79 graph: &DrvGraph,80 wanted: &HashMap<Utf8PathBuf, Vec<String>>,81 ) -> Result<()> {82 let paths = collect_substitute_paths(graph, wanted);83 if paths.is_empty() {84 return Ok(());85 }86 let _ = self87 .events88 .send(BuildEvent::SubstitutePrepassStarted { paths: paths.len() });89 debug!("substitute pre-pass: {} paths", paths.len());9091 let store = self.store.clone();92 let satisfied = spawn_blocking(move || store.substitute_paths(&paths))93 .await94 .expect("substitute pre-pass task should not panic")?;9596 let _ = self.events.send(BuildEvent::SubstitutePrepassFinished {97 satisfied: satisfied.len(),98 });99 Ok(())100 }101102 async fn build_topo(103 &self,104 graph: &Arc<DrvGraph>,105 wanted: HashMap<Utf8PathBuf, Vec<String>>,106 ) -> Result<()> {107 let mut indeg: HashMap<Utf8PathBuf, usize> = graph108 .nodes109 .iter()110 .map(|(k, n)| (k.clone(), n.input_drvs.len()))111 .collect();112 let mut dependents: HashMap<Utf8PathBuf, Vec<Utf8PathBuf>> = HashMap::new();113 for (path, node) in &graph.nodes {114 for dep in node.input_drvs.keys() {115 dependents116 .entry(dep.clone())117 .or_default()118 .push(path.clone());119 }120 }121122 let sem = Arc::new(Semaphore::new(self.parallelism));123 let mut ready: Vec<Utf8PathBuf> = indeg124 .iter()125 .filter(|(_, d)| **d == 0)126 .map(|(k, _)| k.clone())127 .collect();128 let mut in_flight = FuturesUnordered::new();129 let mut failed: HashMap<Utf8PathBuf, String> = HashMap::new();130 131 let mut tainted: HashMap<Utf8PathBuf, Utf8PathBuf> = HashMap::new();132133 loop {134 let batch: Vec<Utf8PathBuf> = mem::take(&mut ready);135 for path in batch {136 if let Some(failed_dep) = tainted.get(&path) {137 let name = graph138 .nodes139 .get(&path)140 .map(|n| n.name.clone())141 .unwrap_or_default();142 let _ = self.events.send(BuildEvent::DrvCancelled {143 drv_path: path.clone(),144 name,145 failed_dep: failed_dep.clone(),146 });147 propagate_done(&dependents, &mut indeg, &mut ready, &path);148 continue;149 }150151 let sem = sem.clone();152 let events = self.events.clone();153 let graph = graph.clone();154 let wanted_here = wanted.get(&path).cloned().unwrap_or_default();155 let store = self.store.clone();156 in_flight.push(tokio::spawn(async move {157 let _permit = sem.acquire_owned().await.expect("semaphore not closed");158 let node = graph159 .nodes160 .get(&path)161 .expect("ready node must be in graph")162 .clone();163 let name = node.name.clone();164165 let all_valid = !wanted_here.is_empty()166 && wanted_here.iter().all(|o| {167 node.outputs168 .get(o)169 .map(|p| store.is_valid_path(p))170 .unwrap_or(false)171 });172 if all_valid {173 let _ = events.send(BuildEvent::DrvSkipped {174 drv_path: path.clone(),175 name: name.clone(),176 });177 return (path, name, Ok::<(), anyhow::Error>(()));178 }179180 let _ = events.send(BuildEvent::DrvStarted {181 drv_path: path.clone(),182 name: name.clone(),183 wanted: wanted_here.clone(),184 });185186 let path_for_build = path.clone();187 let store = store.clone();188 let res = spawn_blocking(move || {189 store.build_drv_outputs(&path_for_build, &wanted_here)190 })191 .await192 .expect("build task should not panic");193194 match res {195 Ok(_) => {196 let _ = events.send(BuildEvent::DrvFinished {197 drv_path: path.clone(),198 name: name.clone(),199 });200 (path, name, Ok(()))201 }202 Err(e) => {203 let msg = format!("{e:#}");204 let _ = events.send(BuildEvent::DrvFailed {205 drv_path: path.clone(),206 name: name.clone(),207 error: msg,208 });209 (path, name, Err(e))210 }211 }212 }));213 }214215 let Some(joined) = in_flight.next().await else {216 break;217 };218 let (finished, _name, res) = match joined {219 Ok(t) => t,220 Err(e) => bail!("scheduler task panicked: {e}"),221 };222 match res {223 Ok(()) => {224 propagate_done(&dependents, &mut indeg, &mut ready, &finished);225 }226 Err(e) => {227 failed.insert(finished.clone(), format!("{e:#}"));228 mark_tainted(&dependents, &finished, &mut tainted);229 propagate_done(&dependents, &mut indeg, &mut ready, &finished);230 }231 }232 }233234 let stuck: Vec<_> = indeg235 .iter()236 .filter(|(_, d)| **d != 0)237 .map(|(k, _)| k.as_str())238 .collect();239 if !stuck.is_empty() {240 warn!(241 "scheduler finished with {} nodes still pending (loop?)",242 stuck.len()243 );244 }245246 if failed.is_empty() {247 info!("scheduler completed");248 Ok(())249 } else {250 let mut report = format!("{} drv(s) failed to build:", failed.len());251 let mut sorted: Vec<_> = failed.iter().collect();252 sorted.sort_by(|a, b| a.0.cmp(b.0));253 for (path, err) in sorted {254 let name = graph255 .nodes256 .get(path)257 .map(|n| n.name.as_str())258 .unwrap_or("?");259 let chain = path_to_root(graph, path);260 report.push_str(&format!(261 "\n\n {name} ({path}):\n {err}\n needed by: {}",262 chain.join(" => "),263 ));264 }265 Err(anyhow::anyhow!(report))266 }267 }268}269270fn propagate_done(271 dependents: &HashMap<Utf8PathBuf, Vec<Utf8PathBuf>>,272 indeg: &mut HashMap<Utf8PathBuf, usize>,273 ready: &mut Vec<Utf8PathBuf>,274 finished: &Utf8Path,275) {276 if let Some(deps) = dependents.get(finished) {277 for d in deps {278 let entry = indeg.get_mut(d).expect("dependent must have indeg");279 *entry = entry.saturating_sub(1);280 if *entry == 0 {281 ready.push(d.clone());282 }283 }284 }285}286287fn mark_tainted(288 dependents: &HashMap<Utf8PathBuf, Vec<Utf8PathBuf>>,289 failed: &Utf8Path,290 tainted: &mut HashMap<Utf8PathBuf, Utf8PathBuf>,291) {292 let mut queue: Vec<Utf8PathBuf> = dependents.get(failed).cloned().unwrap_or_default();293 while let Some(node) = queue.pop() {294 if tainted295 .entry(node.clone())296 .or_insert_with(|| failed.to_owned())297 == failed298 {299 if let Some(deps) = dependents.get(&node) {300 for d in deps {301 if !tainted.contains_key(d) {302 queue.push(d.clone());303 }304 }305 }306 }307 }308}309310fn path_to_root(graph: &DrvGraph, from: &Utf8Path) -> Vec<String> {311 let mut dependents: HashMap<&Utf8Path, Vec<&Utf8Path>> = HashMap::new();312 for (path, node) in &graph.nodes {313 for dep in node.input_drvs.keys() {314 dependents.entry(dep).or_default().push(path);315 }316 }317318 let mut chain: Vec<String> = vec![node_name(graph, from)];319 let mut cur = from;320 let mut seen: HashSet<&Utf8Path> = HashSet::new();321 seen.insert(cur);322 while cur != graph.root.as_str() {323 let Some(next) = dependents.get(cur).and_then(|v| v.first().copied()) else {324 break;325 };326 if !seen.insert(next) {327 break;328 }329 chain.push(node_name(graph, next));330 cur = next;331 }332 chain333}334335fn node_name(graph: &DrvGraph, path: &Utf8Path) -> String {336 graph337 .nodes338 .get(path)339 .map(|n| n.name.clone())340 .unwrap_or_else(|| path.to_string())341}342343fn collect_substitute_paths(344 graph: &DrvGraph,345 wanted: &HashMap<Utf8PathBuf, Vec<String>>,346) -> Vec<Utf8PathBuf> {347 let mut paths: HashSet<Utf8PathBuf> = HashSet::new();348 for node in graph.nodes.values() {349 for src in &node.input_srcs {350 paths.insert(src.clone());351 }352 }353 for (path, outs) in wanted {354 let Some(node) = graph.nodes.get(path) else {355 continue;356 };357 for o in outs {358 if let Some(p) = node.outputs.get(o) {359 paths.insert(p.clone());360 }361 }362 }363 let mut v: Vec<_> = paths.into_iter().collect();364 v.sort();365 v366}367368369370371pub fn build_graph_sync(graph: Arc<DrvGraph>, root_outputs: Vec<String>) -> Result<()> {372 let parallelism = std::thread::available_parallelism()373 .map(|p| p.get())374 .unwrap_or(4);375 let scheduler = Scheduler::new(parallelism);376 crate::await_in_nix(async move { scheduler.run(graph, root_outputs).await })377 .context("scheduler run")378}