1use std::collections::BTreeMap;
17use std::fs;
18use std::path::PathBuf;
19use std::time::{Duration, SystemTime, UNIX_EPOCH};
20
21use anyhow::{Context, Result, bail};
22use qk_config::Workspace;
23use qk_taskgraph::{Task, TaskGraph};
24use serde::{Deserialize, Serialize};
25use serde_json::Value;
26
27use crate::paths::{self, Outputs};
28use crate::store::{Artifact, mode, set_mode, symlink, validate_link};
29use crate::{Cache, hash::digest_file};
30
31pub const RESTORED_AT: std::time::SystemTime = std::time::UNIX_EPOCH;
36
37const KEPT: usize = 8;
39
40#[derive(Clone, Debug, Default)]
42pub struct Warm {
43 pub identity: String,
45 pub outputs: bool,
46 pub paths: Vec<String>,
48 pub env: BTreeMap<String, String>,
50 pub directory: bool,
52 pub max_size: Option<u64>,
53 pub remote: bool,
54 pub preserve_mtimes: bool,
57 pub portable: bool,
59 pub key: Vec<KeyPart>,
61 pub restore_keys: Option<usize>,
64 pub survive: Vec<String>,
67 pub background: bool,
70}
71
72#[derive(Clone, Debug)]
74pub enum KeyPart {
75 File(String),
77 Env(String),
79}
80
81pub fn config(workspace: &Workspace, task: &Task) -> Result<Option<Warm>> {
83 let Some(value) = task.definition.extra.get("qk:warm") else {
84 return Ok(None);
85 };
86 let object = value.as_object().context("qk:warm must be an object")?;
87 for key in object.keys() {
88 if !matches!(
89 key.as_str(),
90 "outputs"
91 | "paths"
92 | "env"
93 | "maxSize"
94 | "remote"
95 | "mtimes"
96 | "portable"
97 | "key"
98 | "restoreKeys"
99 | "group"
100 | "survive"
101 | "save"
102 ) {
103 bail!("unknown qk:warm field {key:?}");
104 }
105 }
106 let flag = |name: &str, default: bool| -> Result<bool> {
107 object.get(name).map_or(Ok(default), |value| {
108 value
109 .as_bool()
110 .with_context(|| format!("qk:warm.{name} must be true or false"))
111 })
112 };
113 let group = object
114 .get("group")
115 .map(|value| {
116 value
117 .as_str()
118 .filter(|name| !name.is_empty())
119 .context("qk:warm.group must be a name")
120 })
121 .transpose()?;
122 if group.is_some() && (object.contains_key("outputs") || object.contains_key("paths")) {
123 bail!("a qk:warm.group shares {{warm}} alone; outputs and paths belong to one task");
124 }
125 let identity = group.map_or_else(|| task.id.clone(), |group| format!("group {group}"));
126 const SURVIVE_SHAPE: &str = "qk:warm.survive must be an array of target names";
127 let survive = object
128 .get("survive")
129 .map(|value| {
130 value
131 .as_array()
132 .context(SURVIVE_SHAPE)?
133 .iter()
134 .map(|target| {
135 target
136 .as_str()
137 .filter(|target| !target.is_empty())
138 .map(str::to_owned)
139 .context(SURVIVE_SHAPE)
140 })
141 .collect::<Result<Vec<_>>>()
142 })
143 .transpose()?
144 .unwrap_or_default();
145 let directory = directory(workspace, &identity);
146 let directory_text = directory.to_str().context("warm directory must be UTF-8")?;
147 let expand = |text: &str| -> Result<String> {
148 let text = paths::expand(workspace, &task.project, text)?;
149 Ok(text.replace("{warm}", directory_text))
150 };
151 let mut paths = Vec::new();
152 for path in object
153 .get("paths")
154 .map(|value| {
155 value
156 .as_array()
157 .context("qk:warm.paths must be an array of paths")
158 })
159 .transpose()?
160 .into_iter()
161 .flatten()
162 {
163 let path = path
164 .as_str()
165 .context("qk:warm.paths must be an array of paths")?;
166 if path.contains("{warm}") {
167 bail!("qk:warm.paths are workspace paths; {{warm}} is kept already");
168 }
169 let (excluded, path) = match path.strip_prefix('!') {
170 Some(path) => (true, path),
171 None => (false, path),
172 };
173 let path = paths::expand(workspace, &task.project, path)?
174 .trim_end_matches('/')
175 .to_owned();
176 paths::validate_path(&path)?;
177 paths.push(if excluded { format!("!{path}") } else { path });
178 }
179 let mut env = BTreeMap::new();
180 let mut uses_directory = false;
181 if let Some(values) = object.get("env") {
182 for (name, value) in values
183 .as_object()
184 .context("qk:warm.env must be an object")?
185 {
186 let value = value
187 .as_str()
188 .context("qk:warm.env values must be strings")?;
189 uses_directory |= value.contains("{warm}");
190 env.insert(name.clone(), expand(value)?);
191 }
192 }
193 let max_size = match object.get("maxSize") {
194 None => None,
195 Some(Value::String(text)) => Some(crate::parse_size(text)?),
196 Some(Value::Number(number)) => number.as_u64(),
197 Some(_) => bail!("qk:warm.maxSize must be a size such as \"2GB\""),
198 };
199 let preserve_mtimes = match object.get("mtimes").map(|value| value.as_str()) {
200 None | Some(Some("epoch")) => false,
201 Some(Some("preserve")) => true,
202 Some(_) => bail!("qk:warm.mtimes must be \"epoch\" or \"preserve\""),
203 };
204 const KEY_SHAPE: &str = "qk:warm.key must be an array of paths and {\"env\": name} objects";
205 let mut key = Vec::new();
206 for part in object
207 .get("key")
208 .map(|value| value.as_array().context(KEY_SHAPE))
209 .transpose()?
210 .into_iter()
211 .flatten()
212 {
213 key.push(match part {
214 Value::String(path) => {
215 let path = paths::expand(workspace, &task.project, path)?;
216 paths::validate_path(&path)?;
217 KeyPart::File(path)
218 }
219 Value::Object(object) if object.len() == 1 => KeyPart::Env(
220 object
221 .get("env")
222 .and_then(Value::as_str)
223 .context(KEY_SHAPE)?
224 .to_owned(),
225 ),
226 _ => bail!(KEY_SHAPE),
227 });
228 }
229 let restore_keys = object
230 .get("restoreKeys")
231 .map(|value| {
232 value
233 .as_u64()
234 .and_then(|count| usize::try_from(count).ok())
235 .filter(|count| *count <= key.len())
236 .context("qk:warm.restoreKeys must be a count no larger than qk:warm.key")
237 })
238 .transpose()?;
239 let background = match object.get("save").map(Value::as_str) {
240 None | Some(Some("wait")) => false,
241 Some(Some("background")) => true,
242 Some(_) => bail!("qk:warm.save must be \"wait\" or \"background\""),
243 };
244 Ok(Some(Warm {
245 identity,
246 outputs: flag("outputs", false)?,
247 paths,
248 env,
249 directory: uses_directory,
250 max_size,
251 remote: flag("remote", true)?,
252 preserve_mtimes,
253 portable: flag("portable", true)?,
254 key,
255 restore_keys,
256 survive,
257 background,
258 }))
259}
260
261fn directory(workspace: &Workspace, identity: &str) -> PathBuf {
263 paths::worktree_state(&workspace.root)
264 .join("warm")
265 .join(&blake3::hash(identity.as_bytes()).to_hex()[..32])
266}
267
268impl Warm {
269 pub fn kept_paths(&self) -> Vec<String> {
272 self.paths
273 .iter()
274 .filter(|path| !path.starts_with('!'))
275 .cloned()
276 .collect()
277 }
278}
279
280#[derive(Default)]
283pub struct Kept {
284 roots: Vec<PathBuf>,
286}
287
288impl Drop for Kept {
289 fn drop(&mut self) {
290 for root in self.roots.iter().rev() {
291 release(root);
292 }
293 }
294}
295
296#[derive(Default, Serialize, Deserialize)]
298struct Stash {
299 #[serde(skip)]
300 root: PathBuf,
301 #[serde(skip)]
302 workspace: PathBuf,
303 moved: Vec<String>,
305}
306
307struct Active {
310 users: usize,
311 stash: Stash,
312 _lock: fs::File,
313}
314
315static ACTIVE: std::sync::Mutex<BTreeMap<PathBuf, Active>> = std::sync::Mutex::new(BTreeMap::new());
318
319impl Stash {
320 fn manifest(root: &std::path::Path) -> PathBuf {
322 root.join("moved.json")
323 }
324
325 fn save_manifest(&self) -> Result<()> {
328 use std::io::Write;
329 let mut file = tempfile::NamedTempFile::new_in(&self.root)?;
330 file.write_all(&serde_json::to_vec(self)?)?;
331 file.as_file().sync_all()?;
332 file.persist(Self::manifest(&self.root))?;
333 Ok(())
334 }
335
336 fn put_back(&self) -> bool {
342 let mut complete = true;
343 for (index, path) in self.moved.iter().enumerate().rev() {
344 let kept = self.root.join(index.to_string());
345 match fs::symlink_metadata(&kept) {
346 Ok(_) => {}
347 Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
348 Err(error) => {
349 complete = false;
350 qk_executor::status!(
351 "qk: could not inspect saved {path}, kept for the next run ({error})"
352 );
353 continue;
354 }
355 }
356 let original = self.workspace.join(path);
357 let result = paths::safe_parents(&self.workspace, path).and_then(|()| {
358 remove(&original)?;
359 if let Some(parent) = original.parent() {
360 fs::create_dir_all(parent)?;
361 }
362 fs::rename(&kept, &original)?;
363 Ok(())
364 });
365 if let Err(error) = result {
366 complete = false;
367 qk_executor::status!(
368 "qk: could not put {path} back, kept for the next run ({error:#})"
369 );
370 }
371 }
372 if complete {
373 let _ = fs::remove_dir_all(&self.root);
374 }
375 complete
376 }
377}
378
379fn remove(path: &std::path::Path) -> std::io::Result<()> {
381 match fs::symlink_metadata(path) {
382 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
383 Err(error) => Err(error),
384 Ok(metadata) if metadata.is_dir() => fs::remove_dir_all(path),
385 Ok(_) => fs::remove_file(path),
386 }
387}
388
389fn release(root: &std::path::Path) {
392 let mut active = ACTIVE
393 .lock()
394 .unwrap_or_else(std::sync::PoisonError::into_inner);
395 let Some(entry) = active.get_mut(root) else {
396 return;
397 };
398 entry.users -= 1;
399 if entry.users > 0 {
400 return;
401 }
402 let entry = active.remove(root).expect("present above");
403 drop(active);
406 entry.stash.put_back();
407}
408
409fn names(owner: &Task, entry: &str, task: &Task) -> bool {
411 match entry.split_once(':') {
412 Some((project, target)) => task.project == project && task.target == target,
413 None => task.project == owner.project && task.target == entry,
414 }
415}
416
417fn depends_on(graph: &TaskGraph, from: &str, on: &str) -> bool {
419 let mut seen = std::collections::BTreeSet::new();
420 let mut next = vec![from];
421 while let Some(id) = next.pop() {
422 for dependency in graph
423 .tasks
424 .get(id)
425 .into_iter()
426 .flat_map(|task| &task.dependencies)
427 {
428 if dependency == on {
429 return true;
430 }
431 if seen.insert(dependency.as_str()) {
432 next.push(dependency);
433 }
434 }
435 }
436 false
437}
438
439pub fn keep_across(workspace: &Workspace, graph: &TaskGraph, task: &Task) -> Kept {
442 let mut kept = Kept::default();
443 for owner in graph.tasks.values() {
444 let Ok(Some(warm)) = config(workspace, owner) else {
445 continue;
446 };
447 if !warm.survive.iter().any(|entry| names(owner, entry, task))
448 || !depends_on(graph, &owner.id, &task.id)
449 {
450 continue;
451 }
452 match acquire(workspace, owner, &warm) {
453 Ok(root) => kept.roots.push(root),
454 Err(error) => qk_executor::status!(
455 "qk: {}: warm paths not kept across {} ({error:#})",
456 owner.id,
457 task.id
458 ),
459 }
460 }
461 kept
462}
463
464fn acquire(workspace: &Workspace, owner: &Task, warm: &Warm) -> Result<PathBuf> {
467 let state = paths::worktree_state(&workspace.root).join("warm-kept");
468 let root = state.join(hash32(&owner.id));
469 fs::create_dir_all(&state)?;
470 let lock = fs::OpenOptions::new()
471 .create(true)
472 .truncate(false)
473 .read(true)
474 .write(true)
475 .open(state.join(format!("{}.lock", hash32(&owner.id))))?;
476 let mut active = loop {
479 let mut active = ACTIVE
480 .lock()
481 .unwrap_or_else(std::sync::PoisonError::into_inner);
482 if let Some(entry) = active.get_mut(&root) {
483 entry.users += 1;
484 return Ok(root);
485 }
486 match fs2::FileExt::try_lock_exclusive(&lock) {
487 Ok(()) => break active,
488 Err(error) if error.raw_os_error() == fs2::lock_contended_error().raw_os_error() => {
489 drop(active);
490 std::thread::sleep(Duration::from_millis(20));
491 }
492 Err(error) => return Err(error.into()),
493 }
494 };
495 let left = match fs::read(Stash::manifest(&root)) {
502 Ok(bytes) => {
503 let mut left: Stash = serde_json::from_slice(&bytes)
504 .context("invalid warm stash manifest; saved paths retained for recovery")?;
505 left.root = root.clone();
506 left.workspace = workspace.root.clone();
507 Some(left)
508 }
509 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
510 if root.is_dir() && fs::read_dir(&root)?.next().is_some() {
511 bail!("missing warm stash manifest; saved paths retained for recovery");
512 }
513 None
514 }
515 Err(error) => return Err(error).context("cannot read warm stash; saved paths retained"),
516 };
517 let retained = left.filter(|left| !left.put_back());
518 let stash = stash(workspace, &root, warm, retained)?;
519 active.insert(
520 root.clone(),
521 Active {
522 users: 1,
523 stash,
524 _lock: lock,
525 },
526 );
527 Ok(root)
528}
529
530fn stash(
534 workspace: &Workspace,
535 root: &std::path::Path,
536 warm: &Warm,
537 retained: Option<Stash>,
538) -> Result<Stash> {
539 let mut stash = match retained {
540 Some(retained) => retained,
541 None => {
542 let _ = fs::remove_dir_all(root);
543 Stash {
544 root: root.to_owned(),
545 workspace: workspace.root.clone(),
546 moved: Vec::new(),
547 }
548 }
549 };
550 let held: Vec<String> = stash
551 .moved
552 .iter()
553 .enumerate()
554 .filter_map(
555 |(index, path)| match fs::symlink_metadata(root.join(index.to_string())) {
556 Ok(_) => Some(Ok(path.clone())),
557 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
558 Err(error) => Some(Err(error).context("cannot inspect retained warm path")),
559 },
560 )
561 .collect::<Result<_>>()?;
562 let overlaps = |path: &str| {
563 held.iter().any(|held| {
564 held == path
565 || path.starts_with(&format!("{held}/"))
566 || held.starts_with(&format!("{path}/"))
567 })
568 };
569 let mut kept = std::collections::BTreeSet::new();
572 for path in warm.kept_paths() {
573 if let Ok(found) = Outputs::from_paths(std::slice::from_ref(&path))
574 .and_then(|outputs| outputs.paths(&workspace.root))
575 {
576 kept.extend(found);
577 }
578 }
579 let topmost: Vec<String> = kept
581 .iter()
582 .filter(|path| {
583 !std::path::Path::new(path.as_str())
584 .ancestors()
585 .skip(1)
586 .any(|ancestor| {
587 ancestor
588 .to_str()
589 .is_some_and(|ancestor| kept.contains(ancestor))
590 })
591 })
592 .filter(|path| !overlaps(path))
593 .cloned()
594 .collect();
595 if topmost.is_empty() {
596 return Ok(stash);
597 }
598 fs::create_dir_all(root)?;
599 for path in topmost {
600 let moved = paths::safe_parents(&workspace.root, &path).and_then(|()| {
601 stash.moved.push(path.clone());
602 stash.save_manifest()?;
603 fs::rename(
606 workspace.root.join(&path),
607 root.join((stash.moved.len() - 1).to_string()),
608 )
609 .inspect_err(|_| {
610 stash.moved.pop();
611 })?;
612 Ok(())
613 });
614 if let Err(error) = moved {
615 stash.save_manifest()?;
618 stash.put_back();
619 return Err(error);
620 }
621 }
622 Ok(stash)
623}
624
625pub fn check_overlaps(workspace: &Workspace, graph: &TaskGraph) -> Result<()> {
628 let mut claimed: Vec<(String, &str)> = Vec::new();
629 for (id, task) in &graph.tasks {
630 let Some(warm) = config(workspace, task)? else {
631 continue;
632 };
633 for path in warm.kept_paths() {
634 if let Some((_, other)) = claimed.iter().find(|(known, _)| {
635 known == &path
636 || known.starts_with(&format!("{path}/"))
637 || path.starts_with(&format!("{known}/"))
638 }) {
639 bail!("{id} and {other} both keep {path} as warm state");
640 }
641 claimed.push((path, id));
642 }
643 }
644 Ok(())
645}
646
647#[derive(Clone, Default, Serialize, Deserialize)]
649pub(crate) struct Group {
650 artifacts: BTreeMap<String, Artifact>,
651 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
654 stamps: BTreeMap<String, Vec<i64>>,
655 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
658 mtimes: BTreeMap<String, i64>,
659}
660
661impl Group {
662 fn size(&self, cache: &Cache) -> u64 {
664 self.artifacts
665 .values()
666 .filter_map(|artifact| match artifact {
667 Artifact::File { blob, .. } => {
668 fs::metadata(cache.root.join("blobs").join(blob)).ok()
669 }
670 _ => None,
671 })
672 .map(|metadata| metadata.len())
673 .sum()
674 }
675}
676
677#[derive(Clone, Default, Serialize, Deserialize)]
679pub(crate) struct Record {
680 version: u32,
681 task: String,
682 worktree: String,
684 #[serde(default, skip_serializing_if = "Option::is_none")]
686 commit: Option<String>,
687 saved: u64,
689 #[serde(default, skip_serializing_if = "Vec::is_empty")]
691 key: Vec<String>,
692 groups: BTreeMap<String, Group>,
693}
694
695const RECORD_VERSION: u32 = 2;
696
697#[derive(Default, Serialize, Deserialize)]
700struct RestoredFiles {
701 groups: BTreeMap<String, BTreeMap<String, (Vec<i64>, Artifact)>>,
702}
703
704#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
706pub struct Restored {
707 pub source: String,
710 pub groups: Vec<String>,
711 pub files: usize,
712 pub bytes: u64,
713}
714
715#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
717#[serde(rename_all = "camelCase")]
718pub struct WarmReport {
719 pub restored: Option<Restored>,
721 #[serde(default, skip_serializing_if = "Vec::is_empty")]
723 pub present: Vec<String>,
724 pub save_ms: Option<u64>,
727 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
729 pub background: bool,
730}
731
732#[derive(Clone)]
735struct Location {
736 name: &'static str,
737 base: PathBuf,
738 outputs: Option<Outputs>,
740}
741
742impl Location {
743 fn paths(&self) -> Result<std::collections::BTreeSet<String>> {
745 if !self.base.exists() {
746 return Ok(Default::default());
747 }
748 match &self.outputs {
749 Some(outputs) => outputs.paths(&self.base),
750 None => {
751 let mut paths = std::collections::BTreeSet::new();
752 for entry in walkdir::WalkDir::new(&self.base)
753 .min_depth(1)
754 .follow_links(false)
755 {
756 let entry = entry?;
757 paths.insert(paths::relative(&self.base, entry.path())?);
758 }
759 Ok(paths)
760 }
761 }
762 }
763
764 fn matches(&self, path: &str) -> bool {
766 self.outputs
767 .as_ref()
768 .is_none_or(|outputs| outputs.matches(path))
769 }
770}
771
772fn locations(workspace: &Workspace, task: &Task, warm: &Warm) -> Result<Vec<Location>> {
774 let mut locations = Vec::new();
775 if warm.outputs {
776 locations.push(Location {
777 name: "outputs",
778 base: workspace.root.clone(),
779 outputs: Some(Outputs::new(workspace, task)?),
780 });
781 }
782 if !warm.paths.is_empty() {
783 locations.push(Location {
784 name: "paths",
785 base: workspace.root.clone(),
786 outputs: Some(Outputs::from_paths(&warm.paths)?),
787 });
788 }
789 if warm.directory {
790 locations.push(Location {
791 name: "directory",
792 base: directory(workspace, &warm.identity),
793 outputs: None,
794 });
795 }
796 Ok(locations)
797}
798
799fn current_branch(workspace: &Workspace) -> Option<String> {
802 for name in ["GITHUB_HEAD_REF", "GITHUB_REF_NAME"] {
803 if let Ok(value) = std::env::var(name)
804 && !value.is_empty()
805 {
806 return Some(value);
807 }
808 }
809 let output = paths::git(&workspace.root, &["rev-parse", "--abbrev-ref", "HEAD"]).ok()?;
810 let branch = String::from_utf8(output.stdout).ok()?.trim().to_owned();
811 (output.status.success() && !branch.is_empty() && branch != "HEAD").then_some(branch)
812}
813
814fn branches(workspace: &Workspace) -> Vec<String> {
817 let default = workspace
818 .config
819 .default_base
820 .clone()
821 .unwrap_or_else(|| "main".into());
822 let default = default
823 .strip_prefix("origin/")
824 .unwrap_or(&default)
825 .to_owned();
826 let mut branches: Vec<String> = current_branch(workspace).into_iter().collect();
827 if !branches.contains(&default) {
828 branches.push(default);
829 }
830 branches
831}
832
833fn stamp(metadata: &fs::Metadata) -> Vec<i64> {
835 let modified = metadata
836 .modified()
837 .ok()
838 .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
839 .map_or(-1, |duration| duration.as_nanos() as i64);
840 #[allow(unused_mut)]
841 let mut stamp = vec![metadata.len() as i64, modified, i64::from(mode(metadata))];
842 #[cfg(unix)]
843 {
844 use std::os::unix::fs::MetadataExt;
845 stamp.extend([
846 metadata.ctime(),
847 metadata.ctime_nsec(),
848 metadata.ino() as i64,
849 ]);
850 }
851 stamp
852}
853
854fn worktree(workspace: &Workspace) -> String {
856 fs::canonicalize(&workspace.root)
857 .unwrap_or_else(|_| workspace.root.clone())
858 .to_string_lossy()
859 .into_owned()
860}
861
862fn current_commit(workspace: &Workspace) -> Option<String> {
864 let output = paths::git(&workspace.root, &["rev-parse", "HEAD"]).ok()?;
865 let commit = String::from_utf8(output.stdout).ok()?.trim().to_owned();
866 (output.status.success() && !commit.is_empty()).then_some(commit)
867}
868
869fn key_digests(
872 workspace: &Workspace,
873 warm: &Warm,
874 prepared: &qk_executor::PreparedTask,
875) -> Result<Vec<String>> {
876 warm.key
877 .iter()
878 .map(|part| {
879 Ok(match part {
880 KeyPart::File(path) => {
881 paths::safe_parents(&workspace.root, path)?;
882 let absolute = workspace.root.join(path);
883 if absolute.is_file() {
884 format!("file:{}", digest_file(&absolute)?)
885 } else {
886 "file:absent".to_owned()
887 }
888 }
889 KeyPart::Env(name) => match prepared
892 .execution
893 .get(std::ffi::OsStr::new(name))
894 .or_else(|| prepared.env.get(std::ffi::OsStr::new(name)))
895 {
896 Some(value) => format!(
897 "env:{}",
898 blake3::hash(value.as_encoded_bytes()).to_hex()[..32].to_owned()
899 ),
900 None => "env:unset".to_owned(),
901 },
902 })
903 })
904 .collect()
905}
906
907fn key_match(warm: &Warm, current: &[String], saved: &[String]) -> Option<u8> {
911 if saved == current {
912 return Some(0);
913 }
914 let count = warm.restore_keys?;
915 (saved.len() == current.len() && saved[..count] == current[..count]).then_some(1)
916}
917
918fn behind(workspace: &Workspace, commit: &str) -> Option<u64> {
921 let ancestor = paths::git(
922 &workspace.root,
923 &["merge-base", "--is-ancestor", commit, "HEAD"],
924 )
925 .ok()?;
926 if !ancestor.status.success() {
927 return None;
928 }
929 let output = paths::git(
930 &workspace.root,
931 &["rev-list", "--count", &format!("{commit}..HEAD")],
932 )
933 .ok()?;
934 String::from_utf8(output.stdout).ok()?.trim().parse().ok()
935}
936
937fn hash32(text: &str) -> String {
939 blake3::hash(text.as_bytes()).to_hex()[..32].to_owned()
940}
941
942fn nanos(time: SystemTime) -> Option<i64> {
944 time.duration_since(UNIX_EPOCH)
945 .ok()
946 .and_then(|duration| i64::try_from(duration.as_nanos()).ok())
947}
948
949fn now_ms() -> u64 {
951 SystemTime::now()
952 .duration_since(UNIX_EPOCH)
953 .map_or(0, |duration| duration.as_millis() as u64)
954}
955
956fn restored_files_path(workspace: &Workspace, warm: &Warm) -> PathBuf {
958 paths::worktree_state(&workspace.root)
959 .join("warm-restored")
960 .join(format!("{}.json", hash32(&warm.identity)))
961}
962
963fn read_restored_files(workspace: &Workspace, warm: &Warm) -> RestoredFiles {
965 fs::read(restored_files_path(workspace, warm))
966 .ok()
967 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
968 .unwrap_or_default()
969}
970
971impl Cache {
972 fn warm_prefix(&self, identity: &str) -> String {
974 hash32(&format!(
975 "{identity}\0{}\0{}",
976 std::env::consts::OS,
977 std::env::consts::ARCH
978 ))
979 }
980
981 fn lock_warm(&self, warm: &Warm) -> Result<fs::File> {
984 let directory = self.root.join("locks");
985 fs::create_dir_all(&directory)?;
986 let file = fs::OpenOptions::new()
987 .create(true)
988 .truncate(false)
989 .read(true)
990 .write(true)
991 .open(directory.join(format!("warm-{}", self.warm_prefix(&warm.identity))))?;
992 fs2::FileExt::lock_exclusive(&file)?;
993 Ok(file)
994 }
995
996 fn warm_record(&self, record: &Record) -> PathBuf {
997 self.root.join("warm").join(format!(
998 "{}-{}.json",
999 self.warm_prefix(&record.task),
1000 hash32(&format!("{}\0{}", record.worktree, record.key.join("\0")))
1001 ))
1002 }
1003
1004 fn warm_records(&self, warm: &Warm) -> Vec<(PathBuf, Record)> {
1006 let prefix = format!("{}-", self.warm_prefix(&warm.identity));
1007 let mut records: Vec<(PathBuf, Record)> = fs::read_dir(self.root.join("warm"))
1008 .into_iter()
1009 .flatten()
1010 .flatten()
1011 .filter(|entry| {
1012 entry
1013 .file_name()
1014 .to_str()
1015 .is_some_and(|name| name.starts_with(&prefix) && name.ends_with(".json"))
1016 })
1017 .filter_map(|entry| {
1018 let record =
1019 serde_json::from_slice::<Record>(&fs::read(entry.path()).ok()?).ok()?;
1020 (record.version == RECORD_VERSION && record.task == warm.identity)
1021 .then(|| (entry.path(), record))
1022 })
1023 .collect();
1024 records.sort_by_key(|(_, record)| std::cmp::Reverse(record.saved));
1025 records
1026 }
1027
1028 fn choose_warm(
1034 &self,
1035 workspace: &Workspace,
1036 warm: &Warm,
1037 current: &[String],
1038 ) -> Option<(Record, String)> {
1039 let own = worktree(workspace);
1040 let mut candidates: Vec<(u8, bool, Record)> = self
1041 .warm_records(warm)
1042 .into_iter()
1043 .filter_map(|(_, record)| {
1044 let quality = key_match(warm, current, &record.key)?;
1045 let other = record.worktree != own;
1046 (!other || warm.portable).then_some((quality, other, record))
1047 })
1048 .collect();
1049 candidates.sort_by_key(|(quality, other, _)| (*quality, *other));
1051 let rank = candidates
1052 .first()
1053 .map(|(quality, other, _)| (*quality, *other))?;
1054 let tied = candidates
1055 .iter()
1056 .take_while(|(quality, other, _)| (*quality, *other) == rank)
1057 .count();
1058 let index = if tied > 1 {
1059 (0..tied)
1060 .min_by_key(|&index| {
1061 let behind = candidates[index]
1062 .2
1063 .commit
1064 .as_deref()
1065 .and_then(|commit| behind(workspace, commit));
1066 (behind.is_none(), behind, index)
1067 })
1068 .unwrap_or(0)
1069 } else {
1070 0
1071 };
1072 let (_, other, record) = candidates.swap_remove(index);
1073 let source = if other {
1074 format!("worktree {}", record.worktree)
1075 } else {
1076 "local".to_owned()
1077 };
1078 Some((record, source))
1079 }
1080
1081 fn store_warm(&self, record: &Record) -> Result<()> {
1082 let path = self.warm_record(record);
1083 let directory = path.parent().context("warm records have a directory")?;
1084 fs::create_dir_all(directory)?;
1085 let mut file = tempfile::NamedTempFile::new_in(self.root.join("tmp"))?;
1086 serde_json::to_writer(std::io::BufWriter::new(file.as_file_mut()), record)?;
1087 file.persist(path)?;
1088 Ok(())
1089 }
1090
1091 pub(crate) fn restore_warm(
1097 &self,
1098 workspace: &Workspace,
1099 task: &Task,
1100 warm: &Warm,
1101 prepared: &qk_executor::PreparedTask,
1102 ) -> Result<(Restored, Vec<String>)> {
1103 let mut restored = Restored::default();
1104 let _lock = self.lock_warm(warm)?;
1105 let mut present = Vec::new();
1107 let mut missing = Vec::new();
1108 for location in locations(workspace, task, warm)? {
1109 if location.paths()?.is_empty() {
1110 missing.push(location);
1111 } else {
1112 present.push(location.name.to_owned());
1113 }
1114 }
1115 if missing.is_empty() {
1116 return Ok((restored, present));
1117 }
1118 let current = key_digests(workspace, warm, prepared)?;
1119 let (record, source) = match self.choose_warm(workspace, warm, ¤t) {
1120 Some(chosen) => chosen,
1121 None => match self.remote_warm(workspace, task, warm, ¤t) {
1122 Some((record, branch)) => (record, format!("remote {branch}")),
1123 None => return Ok((restored, present)),
1124 },
1125 };
1126 let keep_mtimes = warm.preserve_mtimes && source == "local";
1127 restored.source = source;
1128 let mut noted = read_restored_files(workspace, warm);
1129 for location in missing {
1130 let Some(group) = record.groups.get(location.name) else {
1131 continue;
1132 };
1133 fs::create_dir_all(&location.base)?;
1134 let mut files = BTreeMap::new();
1135 for (path, artifact) in &group.artifacts {
1136 paths::safe_parents(&location.base, path)?;
1137 if !location.matches(path) {
1138 bail!("warm state for {} holds {path}, outside its paths", task.id);
1139 }
1140 let destination = location.base.join(path);
1141 if let Some(parent) = destination.parent() {
1142 fs::create_dir_all(parent)?;
1143 }
1144 match artifact {
1145 Artifact::File { blob, mode } => {
1146 let source = self.root.join("blobs").join(blob);
1147 if digest_file(&source)? != *blob {
1148 bail!("warm state for {} has a corrupt blob", task.id);
1149 }
1150 fs::copy(&source, &destination)?;
1151 let modified = keep_mtimes
1152 .then(|| group.mtimes.get(path))
1153 .flatten()
1154 .and_then(|nanos| u64::try_from(*nanos).ok())
1155 .map_or(RESTORED_AT, |nanos| {
1156 UNIX_EPOCH + Duration::from_nanos(nanos)
1157 });
1158 fs::File::options()
1159 .write(true)
1160 .open(&destination)?
1161 .set_modified(modified)?;
1162 set_mode(&destination, *mode)?;
1163 restored.files += 1;
1164 restored.bytes += fs::metadata(&destination)?.len();
1165 files.insert(
1166 path.clone(),
1167 (
1168 stamp(&fs::symlink_metadata(&destination)?),
1169 artifact.clone(),
1170 ),
1171 );
1172 }
1173 Artifact::Directory { mode } => {
1174 fs::create_dir_all(&destination)?;
1175 set_mode(&destination, *mode)?;
1176 }
1177 Artifact::Symlink { target, directory } => {
1178 validate_link(path, target)?;
1179 symlink(target, &destination, *directory)?;
1180 }
1181 }
1182 }
1183 noted.groups.insert(location.name.to_owned(), files);
1184 restored.groups.push(location.name.to_owned());
1185 }
1186 if !restored.groups.is_empty() {
1187 let path = restored_files_path(workspace, warm);
1188 fs::create_dir_all(path.parent().context("restore notes have a directory")?)?;
1189 fs::write(path, serde_json::to_vec(¬ed)?)?;
1190 }
1191 Ok((restored, present))
1192 }
1193
1194 fn remote_warm(
1197 &self,
1198 workspace: &Workspace,
1199 task: &Task,
1200 warm: &Warm,
1201 current: &[String],
1202 ) -> Option<(Record, String)> {
1203 let remote = self
1204 .remote
1205 .as_ref()
1206 .filter(|_| warm.remote && warm.portable)?;
1207 for branch in branches(workspace) {
1208 match remote.fetch_warm(&self.root, &warm.identity, &branch) {
1209 Ok(Some(bytes)) => {
1210 let Some(record) =
1211 serde_json::from_slice::<Record>(&bytes)
1212 .ok()
1213 .filter(|record| {
1214 record.version == RECORD_VERSION && record.task == warm.identity
1215 })
1216 else {
1217 continue;
1218 };
1219 if key_match(warm, current, &record.key).is_none() {
1222 continue;
1223 }
1224 if self.store_warm(&record).is_err() {
1225 return None;
1226 }
1227 return Some((record, branch));
1228 }
1229 Ok(None) => {}
1230 Err(error) => {
1231 qk_executor::status!(
1232 "qk: {}: remote warm state unavailable ({error:#})",
1233 task.id
1234 );
1235 continue;
1238 }
1239 }
1240 }
1241 None
1242 }
1243
1244 pub(crate) fn save_warm(
1247 &self,
1248 workspace: &Workspace,
1249 task: &Task,
1250 warm: &Warm,
1251 prepared: &qk_executor::PreparedTask,
1252 ) -> Result<()> {
1253 self.save_prepared(&Save::prepare(workspace, task, warm, prepared)?)
1254 }
1255
1256 pub(crate) fn save_warm_in_background(
1259 &self,
1260 workspace: &Workspace,
1261 task: &Task,
1262 warm: &Warm,
1263 prepared: &qk_executor::PreparedTask,
1264 ) -> Result<()> {
1265 let save = Save::prepare(workspace, task, warm, prepared)?;
1266 let cache = Cache {
1267 remote: self.remote.clone(),
1268 ..Cache::new(self.root.clone())
1269 };
1270 let handle = std::thread::spawn(move || {
1271 if let Err(error) = cache.save_prepared(&save) {
1272 qk_executor::status!("qk: {}: warm state not saved ({error:#})", save.task);
1273 }
1274 });
1275 self.saves.lock().unwrap().push(handle);
1276 Ok(())
1277 }
1278
1279 fn save_prepared(&self, save: &Save) -> Result<()> {
1284 let warm = &save.warm;
1285 let _lock = self.lock_warm(warm)?;
1286 let records = self.warm_records(warm);
1287 let previous = records
1289 .iter()
1290 .filter(|(_, record)| record.worktree == save.worktree)
1291 .min_by_key(|(_, record)| record.key != save.key)
1292 .map(|(_, record)| record);
1293 let noted: RestoredFiles = fs::read(&save.notes)
1294 .ok()
1295 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1296 .unwrap_or_default();
1297 let mut record = Record {
1298 version: RECORD_VERSION,
1299 task: warm.identity.clone(),
1300 worktree: save.worktree.clone(),
1301 commit: save.commit.clone(),
1302 saved: now_ms(),
1303 key: save.key.clone(),
1304 groups: BTreeMap::new(),
1305 };
1306 let gone = |error: &std::io::Error| error.kind() == std::io::ErrorKind::NotFound;
1307 for location in &save.locations {
1308 if !location.base.exists() {
1309 continue;
1310 }
1311 let known = previous.and_then(|previous| previous.groups.get(location.name));
1312 let restored_files = noted.groups.get(location.name);
1313 let mut group = Group::default();
1314 for path in location.paths()? {
1315 paths::safe_parents(&location.base, &path)?;
1316 let absolute = location.base.join(&path);
1317 let metadata = match fs::symlink_metadata(&absolute) {
1318 Err(error) if gone(&error) => continue,
1319 result => result?,
1320 };
1321 let artifact = if metadata.file_type().is_symlink() {
1322 let target = match fs::read_link(&absolute) {
1323 Err(error) if gone(&error) => continue,
1324 result => result?,
1325 };
1326 let target = target
1327 .to_str()
1328 .context("symlink target must be UTF-8")?
1329 .to_owned();
1330 validate_link(&path, &target)?;
1331 Artifact::Symlink {
1332 target,
1333 directory: absolute.is_dir(),
1334 }
1335 } else if metadata.is_file() {
1336 let now = stamp(&metadata);
1337 let saved = known.and_then(|known| {
1338 (known.stamps.get(&path) == Some(&now))
1339 .then(|| known.artifacts.get(&path))
1340 .flatten()
1341 });
1342 let restored = restored_files
1343 .and_then(|files| files.get(&path))
1344 .filter(|(stamp, _)| *stamp == now)
1345 .map(|(_, artifact)| artifact);
1346 let reused = saved.or(restored).filter(|artifact| match artifact {
1347 Artifact::File { blob, .. } => self.root.join("blobs").join(blob).is_file(),
1348 _ => false,
1349 });
1350 let artifact = match reused {
1351 Some(artifact) => artifact.clone(),
1352 None => match self.put_blob(&absolute) {
1353 Ok(blob) => Artifact::File {
1354 blob,
1355 mode: mode(&metadata),
1356 },
1357 Err(error)
1358 if error.downcast_ref::<std::io::Error>().is_some_and(gone) =>
1359 {
1360 continue;
1361 }
1362 Err(error) => return Err(error),
1363 },
1364 };
1365 group.stamps.insert(path.clone(), now);
1366 if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1367 group.mtimes.insert(path.clone(), modified);
1368 }
1369 artifact
1370 } else if metadata.is_dir() {
1371 Artifact::Directory {
1372 mode: mode(&metadata),
1373 }
1374 } else {
1375 continue;
1376 };
1377 group.artifacts.insert(path, artifact);
1378 }
1379 if let Some(limit) = warm.max_size
1380 && group.size(self) > limit
1381 {
1382 qk_executor::status!(
1383 "qk: {}: warm {} is over its maxSize and was not saved",
1384 save.task,
1385 location.name
1386 );
1387 continue;
1388 }
1389 record.groups.insert(location.name.to_owned(), group);
1390 }
1391 self.store_warm(&record)?;
1392 let _ = fs::remove_file(&save.notes);
1393 let saved = self.warm_record(&record);
1396 for (path, _) in records
1397 .into_iter()
1398 .filter(|(path, _)| *path != saved)
1399 .skip(KEPT - 1)
1400 {
1401 let _ = fs::remove_file(path);
1402 }
1403 if warm.remote
1405 && warm.portable
1406 && let Some(remote) = &self.remote
1407 && let Some(branch) = &save.branch
1408 {
1409 for group in record.groups.values_mut() {
1410 group.stamps.clear();
1411 }
1412 remote.upload_warm(
1413 &self.root,
1414 &warm.identity,
1415 branch,
1416 serde_json::to_vec(&record)?,
1417 );
1418 }
1419 Ok(())
1420 }
1421}
1422
1423struct Save {
1426 task: String,
1427 warm: Warm,
1428 locations: Vec<Location>,
1429 worktree: String,
1430 commit: Option<String>,
1431 branch: Option<String>,
1432 key: Vec<String>,
1433 notes: PathBuf,
1434}
1435
1436impl Save {
1437 fn prepare(
1439 workspace: &Workspace,
1440 task: &Task,
1441 warm: &Warm,
1442 prepared: &qk_executor::PreparedTask,
1443 ) -> Result<Self> {
1444 Ok(Self {
1445 task: task.id.clone(),
1446 warm: warm.clone(),
1447 locations: locations(workspace, task, warm)?,
1448 worktree: worktree(workspace),
1449 commit: current_commit(workspace),
1450 branch: current_branch(workspace),
1451 key: key_digests(workspace, warm, prepared)?,
1452 notes: restored_files_path(workspace, warm),
1453 })
1454 }
1455}
1456
1457#[cfg(test)]
1458mod tests {
1459 use super::*;
1460 use qk_taskgraph::Request;
1461 use serde_json::json;
1462
1463 fn fixture() -> (tempfile::TempDir, Workspace, TaskGraph) {
1465 let temp = tempfile::tempdir().unwrap();
1466 fs::write(temp.path().join("nx.json"), "{}").unwrap();
1467 fs::write(
1468 temp.path().join("project.json"),
1469 json!({
1470 "name": "app", "targets": {
1471 "one": {"command": "true", "qk:warm": {"paths": ["scratch-one"]}},
1472 "two": {"command": "true", "qk:warm": {"paths": ["scratch-two"]}}
1473 }
1474 })
1475 .to_string(),
1476 )
1477 .unwrap();
1478 for path in ["scratch-one", "scratch-two"] {
1479 fs::create_dir(temp.path().join(path)).unwrap();
1480 fs::write(temp.path().join(path).join("state"), "saved").unwrap();
1481 }
1482 let workspace = Workspace::load(temp.path()).unwrap();
1483 let graph = TaskGraph::build(
1484 &workspace,
1485 &[
1486 Request::parse("app:one").unwrap(),
1487 Request::parse("app:two").unwrap(),
1488 ],
1489 )
1490 .unwrap();
1491 (temp, workspace, graph)
1492 }
1493
1494 #[test]
1496 fn waiting_for_an_owner_does_not_block_other_releases() {
1497 let (_temp, workspace, graph) = fixture();
1498 let one = &graph.tasks["app:one"];
1499 let two = &graph.tasks["app:two"];
1500 let first = acquire(&workspace, one, &config(&workspace, one).unwrap().unwrap()).unwrap();
1501 let state = paths::worktree_state(&workspace.root).join("warm-kept");
1502 let lock = fs::OpenOptions::new()
1503 .create(true)
1504 .truncate(false)
1505 .read(true)
1506 .write(true)
1507 .open(state.join(format!("{}.lock", hash32(&two.id))))
1508 .unwrap();
1509 fs2::FileExt::lock_exclusive(&lock).unwrap();
1510 let (released, received) = std::sync::mpsc::channel();
1511 let completed = std::thread::scope(|scope| {
1512 scope.spawn(|| {
1513 let second =
1514 acquire(&workspace, two, &config(&workspace, two).unwrap().unwrap()).unwrap();
1515 release(&second);
1516 });
1517 std::thread::sleep(Duration::from_millis(100));
1518 scope.spawn(|| {
1519 release(&first);
1520 released.send(()).unwrap();
1521 });
1522 let completed = received.recv_timeout(Duration::from_secs(2)).is_ok();
1523 fs2::FileExt::unlock(&lock).unwrap();
1524 completed
1525 });
1526 assert!(
1527 completed,
1528 "a blocked owner lock prevented an unrelated release"
1529 );
1530 assert_eq!(
1531 fs::read_to_string(workspace.root.join("scratch-one/state")).unwrap(),
1532 "saved"
1533 );
1534 }
1535
1536 #[test]
1538 fn unreadable_ownership_preserves_saved_paths() {
1539 let (_temp, workspace, graph) = fixture();
1540 let owner = &graph.tasks["app:one"];
1541 let warm = config(&workspace, owner).unwrap().unwrap();
1542 let root = acquire(&workspace, owner, &warm).unwrap();
1543 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
1544 drop(entry); fs::write(Stash::manifest(&root), "{").unwrap();
1546 assert!(acquire(&workspace, owner, &warm).is_err());
1547 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
1548 fs::remove_file(Stash::manifest(&root)).unwrap();
1549 assert!(acquire(&workspace, owner, &warm).is_err());
1550 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
1551 fs::create_dir(Stash::manifest(&root)).unwrap();
1552 assert!(acquire(&workspace, owner, &warm).is_err());
1553 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
1554 }
1555
1556 #[cfg(unix)]
1558 #[test]
1559 fn inspection_errors_keep_recovery_incomplete() {
1560 let (_temp, workspace, graph) = fixture();
1561 let owner = &graph.tasks["app:one"];
1562 let warm = config(&workspace, owner).unwrap().unwrap();
1563 let root = acquire(&workspace, owner, &warm).unwrap();
1564 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
1565 let saved = root.with_extension("saved");
1566 fs::rename(&root, &saved).unwrap();
1567 fs::write(&root, "blocked").unwrap(); assert!(!entry.stash.put_back());
1569 let retained = Stash {
1570 root: root.clone(),
1571 workspace: workspace.root.clone(),
1572 moved: entry.stash.moved.clone(),
1573 };
1574 assert!(stash(&workspace, &root, &warm, Some(retained)).is_err());
1575 assert_eq!(fs::read_to_string(saved.join("0/state")).unwrap(), "saved");
1576 assert_eq!(fs::read_to_string(&root).unwrap(), "blocked");
1577 }
1578}