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}4748495051const DEFAULT_HEAVY_DRVS: &[&str] = &[52 "composable-kernel",53 "composable_kernel",54 "ghc",55 "gfortran",56 "llvm",57 "clang",58];5960fn heavy_patterns_from_env() -> Vec<String> {61 let mut pats: Vec<String> = DEFAULT_HEAVY_DRVS.iter().map(|s| (*s).to_owned()).collect();62 if let Ok(extra) = std::env::var("FLEET_HEAVY_DRVS") {63 pats.extend(64 extra65 .split(',')66 .map(str::trim)67 .filter(|s| !s.is_empty())68 .map(str::to_owned),69 );70 }71 pats72}7374pub struct Scheduler {75 store: Arc<Store>,76 parallelism: usize,77 heavy_patterns: Vec<String>,78 events: broadcast::Sender<BuildEvent>,79}8081impl Scheduler {82 pub fn new(parallelism: usize) -> Self {83 let parallelism = parallelism.max(1);84 let (events, _) = broadcast::channel(1024);85 Self {86 store: eval_store(),87 parallelism,88 heavy_patterns: heavy_patterns_from_env(),89 events,90 }91 }9293 94 fn permits_for(&self, name: &str) -> u32 {95 if self.heavy_patterns.iter().any(|p| name.contains(p)) {96 self.parallelism as u3297 } else {98 199 }100 }101102 pub fn subscribe(&self) -> broadcast::Receiver<BuildEvent> {103 self.events.subscribe()104 }105106 #[instrument(name = "scheduler", skip(self, graph), fields(root = %graph.root, nodes = graph.nodes.len()))]107 pub async fn run(&self, graph: Arc<DrvGraph>, root_outputs: Vec<String>) -> Result<()> {108 let wanted = graph.wanted_outputs(&root_outputs);109110 self.substitute_prepass(&graph, &wanted).await?;111 self.build_topo(&graph, wanted).await112 }113114 async fn substitute_prepass(115 &self,116 graph: &DrvGraph,117 wanted: &HashMap<Utf8PathBuf, Vec<String>>,118 ) -> Result<()> {119 let paths = collect_substitute_paths(graph, wanted);120 if paths.is_empty() {121 return Ok(());122 }123 let _ = self124 .events125 .send(BuildEvent::SubstitutePrepassStarted { paths: paths.len() });126 debug!("substitute pre-pass: {} paths", paths.len());127128 let store = self.store.clone();129 let satisfied = spawn_blocking(move || store.substitute_paths(&paths))130 .await131 .expect("substitute pre-pass task should not panic")?;132133 let _ = self.events.send(BuildEvent::SubstitutePrepassFinished {134 satisfied: satisfied.len(),135 });136 Ok(())137 }138139 async fn build_topo(140 &self,141 graph: &Arc<DrvGraph>,142 wanted: HashMap<Utf8PathBuf, Vec<String>>,143 ) -> Result<()> {144 let mut indeg: HashMap<Utf8PathBuf, usize> = graph145 .nodes146 .iter()147 .map(|(k, n)| (k.clone(), n.input_drvs.len()))148 .collect();149 let mut dependents: HashMap<Utf8PathBuf, Vec<Utf8PathBuf>> = HashMap::new();150 for (path, node) in &graph.nodes {151 for dep in node.input_drvs.keys() {152 dependents153 .entry(dep.clone())154 .or_default()155 .push(path.clone());156 }157 }158159 let sem = Arc::new(Semaphore::new(self.parallelism));160 let mut ready: Vec<Utf8PathBuf> = indeg161 .iter()162 .filter(|(_, d)| **d == 0)163 .map(|(k, _)| k.clone())164 .collect();165 let mut in_flight = FuturesUnordered::new();166 let mut failed: HashMap<Utf8PathBuf, String> = HashMap::new();167 168 let mut tainted: HashMap<Utf8PathBuf, Utf8PathBuf> = HashMap::new();169170 loop {171 let batch: Vec<Utf8PathBuf> = mem::take(&mut ready);172 for path in batch {173 if let Some(failed_dep) = tainted.get(&path) {174 let name = graph175 .nodes176 .get(&path)177 .map(|n| n.name.clone())178 .unwrap_or_default();179 crate::logging::remove_drv_span(&path);180 let _ = self.events.send(BuildEvent::DrvCancelled {181 drv_path: path.clone(),182 name,183 failed_dep: failed_dep.clone(),184 });185 propagate_done(&dependents, &mut indeg, &mut ready, &path);186 continue;187 }188189 let sem = sem.clone();190 let events = self.events.clone();191 let graph = graph.clone();192 let wanted_here = wanted.get(&path).cloned().unwrap_or_default();193 let store = self.store.clone();194 let permits = graph195 .nodes196 .get(&path)197 .map(|n| self.permits_for(&n.name))198 .unwrap_or(1);199 in_flight.push(tokio::spawn(async move {200 if permits > 1 {201 debug!(permits, "heavy drv: waiting for exclusive build slot");202 }203 204 205 let _permit = sem206 .acquire_many_owned(permits)207 .await208 .expect("semaphore not closed");209 let node = graph210 .nodes211 .get(&path)212 .expect("ready node must be in graph")213 .clone();214 let name = node.name.clone();215216 let all_valid = !wanted_here.is_empty()217 && wanted_here.iter().all(|o| {218 node.outputs219 .get(o)220 .map(|p| store.is_valid_path(p))221 .unwrap_or(false)222 });223 if all_valid {224 let _ = events.send(BuildEvent::DrvSkipped {225 drv_path: path.clone(),226 name: name.clone(),227 });228 return (path, name, Ok::<(), anyhow::Error>(()));229 }230231 let _ = events.send(BuildEvent::DrvStarted {232 drv_path: path.clone(),233 name: name.clone(),234 wanted: wanted_here.clone(),235 });236237 let path_for_build = path.clone();238 let store = store.clone();239 let res = spawn_blocking(move || {240 store.build_drv_outputs(&path_for_build, &wanted_here)241 })242 .await243 .expect("build task should not panic");244245 match res {246 Ok(_) => {247 let _ = events.send(BuildEvent::DrvFinished {248 drv_path: path.clone(),249 name: name.clone(),250 });251 (path, name, Ok(()))252 }253 Err(e) => {254 let msg = format!("{e:#}");255 let _ = events.send(BuildEvent::DrvFailed {256 drv_path: path.clone(),257 name: name.clone(),258 error: msg,259 });260 (path, name, Err(e))261 }262 }263 }));264 }265266 let Some(joined) = in_flight.next().await else {267 break;268 };269 let (finished, _name, res) = match joined {270 Ok(t) => t,271 Err(e) => bail!("scheduler task panicked: {e}"),272 };273 match res {274 Ok(()) => {275 propagate_done(&dependents, &mut indeg, &mut ready, &finished);276 }277 Err(e) => {278 crate::logging::remove_drv_span(&finished);279 failed.insert(finished.clone(), format!("{e:#}"));280 mark_tainted(&dependents, &finished, &mut tainted);281 propagate_done(&dependents, &mut indeg, &mut ready, &finished);282 }283 }284 }285286 let stuck: Vec<_> = indeg287 .iter()288 .filter(|(_, d)| **d != 0)289 .map(|(k, _)| k.as_str())290 .collect();291 if !stuck.is_empty() {292 warn!(293 "scheduler finished with {} nodes still pending (loop?)",294 stuck.len()295 );296 }297298 if failed.is_empty() {299 info!("scheduler completed");300 Ok(())301 } else {302 let mut report = format!("{} drv(s) failed to build:", failed.len());303 let mut sorted: Vec<_> = failed.iter().collect();304 sorted.sort_by(|a, b| a.0.cmp(b.0));305 for (path, err) in sorted {306 let name = graph307 .nodes308 .get(path)309 .map(|n| n.name.as_str())310 .unwrap_or("?");311 let chain = path_to_root(graph, path);312 report.push_str(&format!(313 "\n\n {name} ({path}):\n {err}\n needed by: {}",314 chain.join(" => "),315 ));316 }317 Err(anyhow::anyhow!(report))318 }319 }320}321322fn propagate_done(323 dependents: &HashMap<Utf8PathBuf, Vec<Utf8PathBuf>>,324 indeg: &mut HashMap<Utf8PathBuf, usize>,325 ready: &mut Vec<Utf8PathBuf>,326 finished: &Utf8Path,327) {328 if let Some(deps) = dependents.get(finished) {329 for d in deps {330 let entry = indeg.get_mut(d).expect("dependent must have indeg");331 *entry = entry.saturating_sub(1);332 if *entry == 0 {333 ready.push(d.clone());334 }335 }336 }337}338339fn mark_tainted(340 dependents: &HashMap<Utf8PathBuf, Vec<Utf8PathBuf>>,341 failed: &Utf8Path,342 tainted: &mut HashMap<Utf8PathBuf, Utf8PathBuf>,343) {344 let mut queue: Vec<Utf8PathBuf> = dependents.get(failed).cloned().unwrap_or_default();345 while let Some(node) = queue.pop() {346 if tainted347 .entry(node.clone())348 .or_insert_with(|| failed.to_owned())349 == failed350 {351 if let Some(deps) = dependents.get(&node) {352 for d in deps {353 if !tainted.contains_key(d) {354 queue.push(d.clone());355 }356 }357 }358 }359 }360}361362fn path_to_root(graph: &DrvGraph, from: &Utf8Path) -> Vec<String> {363 let mut dependents: HashMap<&Utf8Path, Vec<&Utf8Path>> = HashMap::new();364 for (path, node) in &graph.nodes {365 for dep in node.input_drvs.keys() {366 dependents.entry(dep).or_default().push(path);367 }368 }369370 let mut chain: Vec<String> = vec![node_name(graph, from)];371 let mut cur = from;372 let mut seen: HashSet<&Utf8Path> = HashSet::new();373 seen.insert(cur);374 while cur != graph.root.as_str() {375 let Some(next) = dependents.get(cur).and_then(|v| v.first().copied()) else {376 break;377 };378 if !seen.insert(next) {379 break;380 }381 chain.push(node_name(graph, next));382 cur = next;383 }384 chain385}386387fn node_name(graph: &DrvGraph, path: &Utf8Path) -> String {388 graph389 .nodes390 .get(path)391 .map(|n| n.name.clone())392 .unwrap_or_else(|| path.to_string())393}394395fn collect_substitute_paths(396 graph: &DrvGraph,397 wanted: &HashMap<Utf8PathBuf, Vec<String>>,398) -> Vec<Utf8PathBuf> {399 let mut paths: HashSet<Utf8PathBuf> = HashSet::new();400 for node in graph.nodes.values() {401 for src in &node.input_srcs {402 paths.insert(src.clone());403 }404 }405 for (path, outs) in wanted {406 let Some(node) = graph.nodes.get(path) else {407 continue;408 };409 for o in outs {410 if let Some(p) = node.outputs.get(o) {411 paths.insert(p.clone());412 }413 }414 }415 let mut v: Vec<_> = paths.into_iter().collect();416 v.sort();417 v418}419420421422423424425fn max_jobs() -> usize {426 let cores = || {427 std::thread::available_parallelism()428 .map(|p| p.get())429 .unwrap_or(4)430 };431 match crate::get_setting(c"max-jobs") {432 Ok(v) if v.trim() == "auto" => cores(),433 Ok(v) => v.trim().parse().unwrap_or_else(|_| cores()),434 Err(_) => cores(),435 }436}437438pub fn build_graph_sync(graph: Arc<DrvGraph>, root_outputs: Vec<String>) -> Result<()> {439 let scheduler = Scheduler::new(max_jobs());440 crate::await_in_nix(async move { scheduler.run(graph, root_outputs).await })441 .context("scheduler run")442}