use std::collections::BTreeMap;
use std::fs;
use std::io::Write;
use std::path::PathBuf;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use anyhow::{Context, Result, bail};
use qk_config::Workspace;
use qk_taskgraph::{Task, TaskGraph};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::paths::{self, Outputs};
use crate::store::{Artifact, mode, set_mode, symlink, validate_link};
use crate::{Cache, hash::digest_file};
pub const RESTORED_AT: std::time::SystemTime = std::time::UNIX_EPOCH;
const KEPT: usize = 8;
#[derive(Clone, Debug, Default)]
pub struct Warm {
pub identity: String,
pub outputs: bool,
pub paths: Vec<String>,
pub env: BTreeMap<String, String>,
pub directory: bool,
pub max_size: Option<u64>,
pub remote: bool,
pub preserve_mtimes: bool,
pub portable: bool,
pub key: Vec<KeyPart>,
pub restore_keys: Option<usize>,
pub survive: Vec<String>,
pub background: bool,
}
#[derive(Clone, Debug)]
pub enum KeyPart {
File(String),
Env(String),
}
pub fn config(workspace: &Workspace, task: &Task) -> Result<Option<Warm>> {
let Some(value) = task.definition.extra.get("qk:warm") else {
return Ok(None);
};
let object = value.as_object().context("qk:warm must be an object")?;
for key in object.keys() {
if !matches!(
key.as_str(),
"outputs"
| "paths"
| "env"
| "maxSize"
| "remote"
| "mtimes"
| "portable"
| "key"
| "restoreKeys"
| "group"
| "survive"
| "save"
) {
bail!("unknown qk:warm field {key:?}");
}
}
let flag = |name: &str, default: bool| -> Result<bool> {
object.get(name).map_or(Ok(default), |value| {
value
.as_bool()
.with_context(|| format!("qk:warm.{name} must be true or false"))
})
};
let group = object
.get("group")
.map(|value| {
value
.as_str()
.filter(|name| !name.is_empty())
.context("qk:warm.group must be a name")
})
.transpose()?;
if group.is_some() && (object.contains_key("outputs") || object.contains_key("paths")) {
bail!("a qk:warm.group shares {{warm}} alone; outputs and paths belong to one task");
}
let identity = group.map_or_else(|| task.id.clone(), |group| format!("group {group}"));
const SURVIVE_SHAPE: &str = "qk:warm.survive must be an array of target names";
let survive = object
.get("survive")
.map(|value| {
value
.as_array()
.context(SURVIVE_SHAPE)?
.iter()
.map(|target| {
target
.as_str()
.filter(|target| !target.is_empty())
.map(str::to_owned)
.context(SURVIVE_SHAPE)
})
.collect::<Result<Vec<_>>>()
})
.transpose()?
.unwrap_or_default();
let directory = directory(workspace, &identity);
let directory_text = directory.to_str().context("warm directory must be UTF-8")?;
let workspace_text = workspace
.root
.to_str()
.context("workspace root must be UTF-8")?;
let project_root = &workspace.projects[&task.project].root;
let project_directory = if project_root == "." {
workspace.root.clone()
} else {
workspace.root.join(project_root)
};
let project_text = project_directory
.to_str()
.context("project root must be UTF-8")?;
let expand = |text: &str| -> Result<String> {
if text.contains("{options.") || text.contains("{args.") {
bail!("dynamic warm environment values are not supported yet");
}
Ok(text
.replace("{workspaceRoot}", workspace_text)
.replace("{projectRoot}", project_text)
.replace("{warm}", directory_text))
};
let mut paths = Vec::new();
for path in object
.get("paths")
.map(|value| {
value
.as_array()
.context("qk:warm.paths must be an array of paths")
})
.transpose()?
.into_iter()
.flatten()
{
let path = path
.as_str()
.context("qk:warm.paths must be an array of paths")?;
if path.contains("{warm}") {
bail!("qk:warm.paths are workspace paths; {{warm}} is kept already");
}
let (excluded, path) = match path.strip_prefix('!') {
Some(path) => (true, path),
None => (false, path),
};
let path = paths::expand(workspace, &task.project, path)?
.trim_end_matches('/')
.to_owned();
paths::validate_path(&path)?;
paths.push(if excluded { format!("!{path}") } else { path });
}
let mut env = BTreeMap::new();
let mut uses_directory = false;
if let Some(values) = object.get("env") {
for (name, value) in values
.as_object()
.context("qk:warm.env must be an object")?
{
let value = value
.as_str()
.context("qk:warm.env values must be strings")?;
uses_directory |= value.contains("{warm}");
env.insert(name.clone(), expand(value)?);
}
}
let max_size = match object.get("maxSize") {
None => None,
Some(Value::String(text)) => Some(crate::parse_size(text)?),
Some(Value::Number(number)) => number.as_u64(),
Some(_) => bail!("qk:warm.maxSize must be a size such as \"2GB\""),
};
let preserve_mtimes = match object.get("mtimes").map(|value| value.as_str()) {
None | Some(Some("epoch")) => false,
Some(Some("preserve")) => true,
Some(_) => bail!("qk:warm.mtimes must be \"epoch\" or \"preserve\""),
};
const KEY_SHAPE: &str = "qk:warm.key must be an array of paths and {\"env\": name} objects";
let mut key = Vec::new();
for part in object
.get("key")
.map(|value| value.as_array().context(KEY_SHAPE))
.transpose()?
.into_iter()
.flatten()
{
key.push(match part {
Value::String(path) => {
let path = paths::expand(workspace, &task.project, path)?;
paths::validate_path(&path)?;
KeyPart::File(path)
}
Value::Object(object) if object.len() == 1 => KeyPart::Env(
object
.get("env")
.and_then(Value::as_str)
.context(KEY_SHAPE)?
.to_owned(),
),
_ => bail!(KEY_SHAPE),
});
}
let restore_keys = object
.get("restoreKeys")
.map(|value| {
value
.as_u64()
.and_then(|count| usize::try_from(count).ok())
.filter(|count| *count <= key.len())
.context("qk:warm.restoreKeys must be a count no larger than qk:warm.key")
})
.transpose()?;
let background = match object.get("save").map(Value::as_str) {
None | Some(Some("wait")) => false,
Some(Some("background")) => true,
Some(_) => bail!("qk:warm.save must be \"wait\" or \"background\""),
};
Ok(Some(Warm {
identity,
outputs: flag("outputs", false)?,
paths,
env,
directory: uses_directory,
max_size,
remote: flag("remote", true)?,
preserve_mtimes,
portable: flag("portable", false)?,
key,
restore_keys,
survive,
background,
}))
}
fn directory(workspace: &Workspace, identity: &str) -> PathBuf {
paths::worktree_state(&workspace.root)
.join("warm")
.join(&blake3::hash(identity.as_bytes()).to_hex()[..32])
}
impl Warm {
pub fn kept_paths(&self) -> Vec<String> {
self.paths
.iter()
.filter(|path| !path.starts_with('!'))
.cloned()
.collect()
}
}
#[derive(Default)]
pub struct Kept {
roots: Vec<PathBuf>,
}
impl Drop for Kept {
fn drop(&mut self) {
for root in self.roots.iter().rev() {
release(root);
}
}
}
#[derive(Default, Serialize, Deserialize)]
struct Stash {
#[serde(skip)]
root: PathBuf,
#[serde(skip)]
workspace: PathBuf,
moved: Vec<String>,
}
struct Active {
users: usize,
stash: Stash,
_lock: fs::File,
}
static ACTIVE: std::sync::Mutex<BTreeMap<PathBuf, Active>> = std::sync::Mutex::new(BTreeMap::new());
impl Stash {
fn manifest(root: &std::path::Path) -> PathBuf {
root.join("moved.json")
}
fn save_manifest(&self) -> Result<()> {
use std::io::Write;
let mut file = tempfile::NamedTempFile::new_in(&self.root)?;
file.write_all(&serde_json::to_vec(self)?)?;
file.as_file().sync_all()?;
file.persist(Self::manifest(&self.root))?;
Ok(())
}
fn put_back(&self) -> bool {
let mut complete = true;
for (index, path) in self.moved.iter().enumerate().rev() {
let kept = self.root.join(index.to_string());
match fs::symlink_metadata(&kept) {
Ok(_) => {}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
Err(error) => {
complete = false;
qk_executor::status!(
"qk: could not inspect saved {path}, kept for the next run ({error})"
);
continue;
}
}
let original = self.workspace.join(path);
let result = paths::safe_parents(&self.workspace, path).and_then(|()| {
remove(&original)?;
if let Some(parent) = original.parent() {
fs::create_dir_all(parent)?;
}
fs::rename(&kept, &original)?;
Ok(())
});
if let Err(error) = result {
complete = false;
qk_executor::status!(
"qk: could not put {path} back, kept for the next run ({error:#})"
);
}
}
if complete {
let _ = fs::remove_dir_all(&self.root);
}
complete
}
}
fn remove(path: &std::path::Path) -> std::io::Result<()> {
match fs::symlink_metadata(path) {
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(error),
Ok(metadata) if metadata.is_dir() => fs::remove_dir_all(path),
Ok(_) => fs::remove_file(path),
}
}
fn release(root: &std::path::Path) {
let mut active = ACTIVE
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let Some(entry) = active.get_mut(root) else {
return;
};
entry.users -= 1;
if entry.users > 0 {
return;
}
let entry = active.remove(root).expect("present above");
drop(active);
entry.stash.put_back();
}
fn names(owner: &Task, entry: &str, task: &Task) -> bool {
match entry.split_once(':') {
Some((project, target)) => task.project == project && task.target == target,
None => task.project == owner.project && task.target == entry,
}
}
fn depends_on(graph: &TaskGraph, from: &str, on: &str) -> bool {
let mut seen = std::collections::BTreeSet::new();
let mut next = vec![from];
while let Some(id) = next.pop() {
for dependency in graph
.tasks
.get(id)
.into_iter()
.flat_map(|task| &task.dependencies)
{
if dependency == on {
return true;
}
if seen.insert(dependency.as_str()) {
next.push(dependency);
}
}
}
false
}
pub fn keep_across(workspace: &Workspace, graph: &TaskGraph, task: &Task) -> Kept {
let mut kept = Kept::default();
for owner in graph.tasks.values() {
let Ok(Some(warm)) = config(workspace, owner) else {
continue;
};
if !warm.survive.iter().any(|entry| names(owner, entry, task))
|| !depends_on(graph, &owner.id, &task.id)
{
continue;
}
match acquire(workspace, owner, &warm) {
Ok(root) => kept.roots.push(root),
Err(error) => qk_executor::status!(
"qk: {}: warm paths not kept across {} ({error:#})",
owner.id,
task.id
),
}
}
kept
}
fn acquire(workspace: &Workspace, owner: &Task, warm: &Warm) -> Result<PathBuf> {
let state = paths::worktree_state(&workspace.root).join("warm-kept");
let root = state.join(hash32(&owner.id));
fs::create_dir_all(&state)?;
let lock = fs::OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(state.join(format!("{}.lock", hash32(&owner.id))))?;
let mut active = loop {
let mut active = ACTIVE
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(entry) = active.get_mut(&root) {
entry.users += 1;
return Ok(root);
}
match fs2::FileExt::try_lock_exclusive(&lock) {
Ok(()) => break active,
Err(error) if error.raw_os_error() == fs2::lock_contended_error().raw_os_error() => {
drop(active);
std::thread::sleep(Duration::from_millis(20));
}
Err(error) => return Err(error.into()),
}
};
let left = match fs::read(Stash::manifest(&root)) {
Ok(bytes) => {
let mut left: Stash = serde_json::from_slice(&bytes)
.context("invalid warm stash manifest; saved paths retained for recovery")?;
left.root = root.clone();
left.workspace = workspace.root.clone();
Some(left)
}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
if root.is_dir() && fs::read_dir(&root)?.next().is_some() {
bail!("missing warm stash manifest; saved paths retained for recovery");
}
None
}
Err(error) => return Err(error).context("cannot read warm stash; saved paths retained"),
};
let retained = left.filter(|left| !left.put_back());
let stash = stash(workspace, &root, warm, retained)?;
active.insert(
root.clone(),
Active {
users: 1,
stash,
_lock: lock,
},
);
Ok(root)
}
fn stash(
workspace: &Workspace,
root: &std::path::Path,
warm: &Warm,
retained: Option<Stash>,
) -> Result<Stash> {
let mut stash = match retained {
Some(retained) => retained,
None => {
let _ = fs::remove_dir_all(root);
Stash {
root: root.to_owned(),
workspace: workspace.root.clone(),
moved: Vec::new(),
}
}
};
let held: Vec<String> = stash
.moved
.iter()
.enumerate()
.filter_map(
|(index, path)| match fs::symlink_metadata(root.join(index.to_string())) {
Ok(_) => Some(Ok(path.clone())),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
Err(error) => Some(Err(error).context("cannot inspect retained warm path")),
},
)
.collect::<Result<_>>()?;
let overlaps = |path: &str| {
held.iter().any(|held| {
held == path
|| path.starts_with(&format!("{held}/"))
|| held.starts_with(&format!("{path}/"))
})
};
let mut kept = std::collections::BTreeSet::new();
for path in warm.kept_paths() {
if let Ok(found) = Outputs::from_paths(std::slice::from_ref(&path))
.and_then(|outputs| outputs.paths(&workspace.root))
{
kept.extend(found);
}
}
let topmost: Vec<String> = kept
.iter()
.filter(|path| {
!std::path::Path::new(path.as_str())
.ancestors()
.skip(1)
.any(|ancestor| {
ancestor
.to_str()
.is_some_and(|ancestor| kept.contains(ancestor))
})
})
.filter(|path| !overlaps(path))
.cloned()
.collect();
if topmost.is_empty() {
return Ok(stash);
}
fs::create_dir_all(root)?;
for path in topmost {
let moved = paths::safe_parents(&workspace.root, &path).and_then(|()| {
stash.moved.push(path.clone());
stash.save_manifest()?;
fs::rename(
workspace.root.join(&path),
root.join((stash.moved.len() - 1).to_string()),
)
.inspect_err(|_| {
stash.moved.pop();
})?;
Ok(())
});
if let Err(error) = moved {
stash.save_manifest()?;
stash.put_back();
return Err(error);
}
}
Ok(stash)
}
pub fn check_overlaps(workspace: &Workspace, graph: &TaskGraph) -> Result<()> {
let mut claimed: Vec<(String, &str)> = Vec::new();
for (id, task) in &graph.tasks {
let Some(warm) = config(workspace, task)? else {
continue;
};
for path in warm.kept_paths() {
if let Some((_, other)) = claimed.iter().find(|(known, _)| {
known == &path
|| known.starts_with(&format!("{path}/"))
|| path.starts_with(&format!("{known}/"))
}) {
bail!("{id} and {other} both keep {path} as warm state");
}
claimed.push((path, id));
}
}
Ok(())
}
#[derive(Clone, Default, Serialize, Deserialize)]
pub(crate) struct Group {
artifacts: BTreeMap<String, Artifact>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
stamps: BTreeMap<String, Vec<i64>>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
mtimes: BTreeMap<String, i64>,
}
impl Group {
fn size(&self, cache: &Cache) -> u64 {
self.artifacts
.values()
.filter_map(|artifact| match artifact {
Artifact::File { blob, .. } => {
fs::metadata(cache.root.join("blobs").join(blob)).ok()
}
_ => None,
})
.map(|metadata| metadata.len())
.sum()
}
}
#[derive(Clone, Default, Serialize, Deserialize)]
pub(crate) struct Record {
version: u32,
task: String,
worktree: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
commit: Option<String>,
saved: u64,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
key: Vec<String>,
groups: BTreeMap<String, Group>,
}
const RECORD_VERSION: u32 = 2;
#[derive(Default, Serialize, Deserialize)]
struct RestoredFiles {
groups: BTreeMap<String, BTreeMap<String, (Vec<i64>, Artifact)>>,
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Restored {
pub source: String,
pub groups: Vec<String>,
pub files: usize,
pub bytes: u64,
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WarmReport {
pub restored: Option<Restored>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub present: Vec<String>,
pub save_ms: Option<u64>,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub background: bool,
}
#[derive(Clone)]
struct Location {
name: &'static str,
base: PathBuf,
outputs: Option<Outputs>,
}
impl Location {
fn paths(&self) -> Result<std::collections::BTreeSet<String>> {
if !self.base.exists() {
return Ok(Default::default());
}
match &self.outputs {
Some(outputs) => outputs.paths(&self.base),
None => {
let mut paths = std::collections::BTreeSet::new();
for entry in walkdir::WalkDir::new(&self.base)
.min_depth(1)
.follow_links(false)
{
let entry = entry?;
paths.insert(paths::relative(&self.base, entry.path())?);
}
Ok(paths)
}
}
}
fn matches(&self, path: &str) -> bool {
self.outputs
.as_ref()
.is_none_or(|outputs| outputs.matches(path))
}
}
fn locations(workspace: &Workspace, task: &Task, warm: &Warm) -> Result<Vec<Location>> {
let mut locations = Vec::new();
if warm.outputs {
locations.push(Location {
name: "outputs",
base: workspace.root.clone(),
outputs: Some(Outputs::new(workspace, task)?),
});
}
if !warm.paths.is_empty() {
locations.push(Location {
name: "paths",
base: workspace.root.clone(),
outputs: Some(Outputs::from_paths(&warm.paths)?),
});
}
if warm.directory {
locations.push(Location {
name: "directory",
base: directory(workspace, &warm.identity),
outputs: None,
});
}
Ok(locations)
}
fn current_branch(workspace: &Workspace) -> Option<String> {
for name in ["GITHUB_HEAD_REF", "GITHUB_REF_NAME"] {
if let Ok(value) = std::env::var(name)
&& !value.is_empty()
{
return Some(value);
}
}
let output = paths::git(&workspace.root, &["rev-parse", "--abbrev-ref", "HEAD"]).ok()?;
let branch = String::from_utf8(output.stdout).ok()?.trim().to_owned();
(output.status.success() && !branch.is_empty() && branch != "HEAD").then_some(branch)
}
fn branches(workspace: &Workspace) -> Vec<String> {
let default = workspace
.config
.default_base
.clone()
.unwrap_or_else(|| "main".into());
let default = default
.strip_prefix("origin/")
.unwrap_or(&default)
.to_owned();
let mut branches: Vec<String> = current_branch(workspace).into_iter().collect();
if !branches.contains(&default) {
branches.push(default);
}
branches
}
fn stamp(metadata: &fs::Metadata) -> Vec<i64> {
let modified = metadata
.modified()
.ok()
.and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
.map_or(-1, |duration| duration.as_nanos() as i64);
#[allow(unused_mut)]
let mut stamp = vec![metadata.len() as i64, modified, i64::from(mode(metadata))];
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
stamp.extend([
metadata.ctime(),
metadata.ctime_nsec(),
metadata.ino() as i64,
]);
}
stamp
}
fn worktree(workspace: &Workspace) -> String {
fs::canonicalize(&workspace.root)
.unwrap_or_else(|_| workspace.root.clone())
.to_string_lossy()
.into_owned()
}
fn current_commit(workspace: &Workspace) -> Option<String> {
let output = paths::git(&workspace.root, &["rev-parse", "HEAD"]).ok()?;
let commit = String::from_utf8(output.stdout).ok()?.trim().to_owned();
(output.status.success() && !commit.is_empty()).then_some(commit)
}
fn key_digests(
workspace: &Workspace,
warm: &Warm,
prepared: &qk_executor::PreparedTask,
) -> Result<Vec<String>> {
warm.key
.iter()
.map(|part| {
Ok(match part {
KeyPart::File(path) => {
paths::safe_parents(&workspace.root, path)?;
let absolute = workspace.root.join(path);
if absolute.is_file() {
format!("file:{}", digest_file(&absolute)?)
} else {
"file:absent".to_owned()
}
}
KeyPart::Env(name) => match prepared
.execution
.get(std::ffi::OsStr::new(name))
.or_else(|| prepared.env.get(std::ffi::OsStr::new(name)))
{
Some(value) => format!(
"env:{}",
blake3::hash(value.as_encoded_bytes()).to_hex()[..32].to_owned()
),
None => "env:unset".to_owned(),
},
})
})
.collect()
}
fn key_match(warm: &Warm, current: &[String], saved: &[String]) -> Option<u8> {
if saved == current {
return Some(0);
}
let count = warm.restore_keys?;
(saved.len() == current.len() && saved[..count] == current[..count]).then_some(1)
}
fn behind(workspace: &Workspace, commit: &str) -> Option<u64> {
let ancestor = paths::git(
&workspace.root,
&["merge-base", "--is-ancestor", commit, "HEAD"],
)
.ok()?;
if !ancestor.status.success() {
return None;
}
let output = paths::git(
&workspace.root,
&["rev-list", "--count", &format!("{commit}..HEAD")],
)
.ok()?;
String::from_utf8(output.stdout).ok()?.trim().parse().ok()
}
fn hash32(text: &str) -> String {
blake3::hash(text.as_bytes()).to_hex()[..32].to_owned()
}
fn nanos(time: SystemTime) -> Option<i64> {
time.duration_since(UNIX_EPOCH)
.ok()
.and_then(|duration| i64::try_from(duration.as_nanos()).ok())
}
fn now_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |duration| duration.as_millis() as u64)
}
fn restored_files_path(workspace: &Workspace, warm: &Warm) -> PathBuf {
paths::worktree_state(&workspace.root)
.join("warm-restored")
.join(format!("{}.json", hash32(&warm.identity)))
}
fn read_restored_files(workspace: &Workspace, warm: &Warm) -> RestoredFiles {
fs::read(restored_files_path(workspace, warm))
.ok()
.and_then(|bytes| serde_json::from_slice(&bytes).ok())
.unwrap_or_default()
}
impl Cache {
fn warm_prefix(&self, identity: &str) -> String {
hash32(&format!(
"{identity}\0{}\0{}",
std::env::consts::OS,
std::env::consts::ARCH
))
}
fn lock_warm(&self, warm: &Warm) -> Result<fs::File> {
let directory = self.root.join("locks");
fs::create_dir_all(&directory)?;
let file = fs::OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(directory.join(format!("warm-{}", self.warm_prefix(&warm.identity))))?;
fs2::FileExt::lock_exclusive(&file)?;
Ok(file)
}
fn warm_record(&self, record: &Record) -> PathBuf {
self.root.join("warm").join(format!(
"{}-{}.json",
self.warm_prefix(&record.task),
hash32(&format!("{}\0{}", record.worktree, record.key.join("\0")))
))
}
fn warm_records(&self, warm: &Warm) -> Vec<(PathBuf, Record)> {
let prefix = format!("{}-", self.warm_prefix(&warm.identity));
let mut records: Vec<(PathBuf, Record)> = fs::read_dir(self.root.join("warm"))
.into_iter()
.flatten()
.flatten()
.filter(|entry| {
entry
.file_name()
.to_str()
.is_some_and(|name| name.starts_with(&prefix) && name.ends_with(".json"))
})
.filter_map(|entry| {
let record =
serde_json::from_slice::<Record>(&fs::read(entry.path()).ok()?).ok()?;
(record.version == RECORD_VERSION && record.task == warm.identity)
.then(|| (entry.path(), record))
})
.collect();
records.sort_by_key(|(_, record)| std::cmp::Reverse(record.saved));
records
}
fn choose_warm(
&self,
workspace: &Workspace,
warm: &Warm,
current: &[String],
) -> Option<(Record, String)> {
let own = worktree(workspace);
let mut candidates: Vec<(u8, bool, Record)> = self
.warm_records(warm)
.into_iter()
.filter_map(|(_, record)| {
let quality = key_match(warm, current, &record.key)?;
let other = record.worktree != own;
(!other || warm.portable).then_some((quality, other, record))
})
.collect();
candidates.sort_by_key(|(quality, other, _)| (*quality, *other));
let rank = candidates
.first()
.map(|(quality, other, _)| (*quality, *other))?;
let tied = candidates
.iter()
.take_while(|(quality, other, _)| (*quality, *other) == rank)
.count();
let index = if tied > 1 {
(0..tied)
.min_by_key(|&index| {
let behind = candidates[index]
.2
.commit
.as_deref()
.and_then(|commit| behind(workspace, commit));
(behind.is_none(), behind, index)
})
.unwrap_or(0)
} else {
0
};
let (_, other, record) = candidates.swap_remove(index);
let source = if other {
format!("worktree {}", record.worktree)
} else {
"local".to_owned()
};
Some((record, source))
}
fn store_warm(&self, record: &Record) -> Result<()> {
let path = self.warm_record(record);
let directory = path.parent().context("warm records have a directory")?;
fs::create_dir_all(directory)?;
let mut file = tempfile::NamedTempFile::new_in(self.root.join("tmp"))?;
{
let mut writer = std::io::BufWriter::new(file.as_file_mut());
serde_json::to_writer(&mut writer, record)?;
writer.flush()?;
}
file.persist(path)?;
Ok(())
}
pub(crate) fn restore_warm(
&self,
workspace: &Workspace,
task: &Task,
warm: &Warm,
prepared: &qk_executor::PreparedTask,
) -> Result<(Restored, Vec<String>)> {
let mut restored = Restored::default();
let _lock = self.lock_warm(warm)?;
let mut present = Vec::new();
let mut missing = Vec::new();
for location in locations(workspace, task, warm)? {
if location.paths()?.is_empty() {
missing.push(location);
} else {
present.push(location.name.to_owned());
}
}
if missing.is_empty() {
return Ok((restored, present));
}
let current = key_digests(workspace, warm, prepared)?;
let (record, source) = match self.choose_warm(workspace, warm, ¤t) {
Some(chosen) => chosen,
None => match self.remote_warm(workspace, task, warm, ¤t) {
Some((record, branch)) => (record, format!("remote {branch}")),
None => return Ok((restored, present)),
},
};
let keep_mtimes = warm.preserve_mtimes && source == "local";
restored.source = source;
let mut noted = read_restored_files(workspace, warm);
for location in missing {
let Some(group) = record.groups.get(location.name) else {
continue;
};
fs::create_dir_all(&location.base)?;
let mut files = BTreeMap::new();
for (path, artifact) in &group.artifacts {
paths::safe_parents(&location.base, path)?;
if !location.matches(path) {
bail!("warm state for {} holds {path}, outside its paths", task.id);
}
let destination = location.base.join(path);
if let Some(parent) = destination.parent() {
fs::create_dir_all(parent)?;
}
match artifact {
Artifact::File { blob, mode } => {
let source = self.root.join("blobs").join(blob);
if digest_file(&source)? != *blob {
bail!("warm state for {} has a corrupt blob", task.id);
}
fs::copy(&source, &destination)?;
let modified = keep_mtimes
.then(|| group.mtimes.get(path))
.flatten()
.and_then(|nanos| u64::try_from(*nanos).ok())
.map_or(RESTORED_AT, |nanos| {
UNIX_EPOCH + Duration::from_nanos(nanos)
});
fs::File::options()
.write(true)
.open(&destination)?
.set_modified(modified)?;
set_mode(&destination, *mode)?;
restored.files += 1;
restored.bytes += fs::metadata(&destination)?.len();
files.insert(
path.clone(),
(
stamp(&fs::symlink_metadata(&destination)?),
artifact.clone(),
),
);
}
Artifact::Directory { mode } => {
fs::create_dir_all(&destination)?;
set_mode(&destination, *mode)?;
}
Artifact::Symlink { target, directory } => {
validate_link(path, target)?;
symlink(target, &destination, *directory)?;
}
}
}
noted.groups.insert(location.name.to_owned(), files);
restored.groups.push(location.name.to_owned());
}
if !restored.groups.is_empty() {
let path = restored_files_path(workspace, warm);
fs::create_dir_all(path.parent().context("restore notes have a directory")?)?;
fs::write(path, serde_json::to_vec(¬ed)?)?;
}
Ok((restored, present))
}
fn remote_warm(
&self,
workspace: &Workspace,
task: &Task,
warm: &Warm,
current: &[String],
) -> Option<(Record, String)> {
let remote = self
.remote
.as_ref()
.filter(|_| warm.remote && warm.portable)?;
for branch in branches(workspace) {
match remote.fetch_warm(&self.root, &warm.identity, &branch) {
Ok(Some(bytes)) => {
let Some(record) =
serde_json::from_slice::<Record>(&bytes)
.ok()
.filter(|record| {
record.version == RECORD_VERSION && record.task == warm.identity
})
else {
continue;
};
if key_match(warm, current, &record.key).is_none() {
continue;
}
if self.store_warm(&record).is_err() {
return None;
}
return Some((record, branch));
}
Ok(None) => {}
Err(error) => {
qk_executor::status!(
"qk: {}: remote warm state unavailable ({error:#})",
task.id
);
continue;
}
}
}
None
}
pub(crate) fn save_warm(
&self,
workspace: &Workspace,
task: &Task,
warm: &Warm,
prepared: &qk_executor::PreparedTask,
) -> Result<()> {
self.save_prepared(&Save::prepare(workspace, task, warm, prepared)?)
}
pub(crate) fn save_warm_in_background(
&self,
workspace: &Workspace,
task: &Task,
warm: &Warm,
prepared: &qk_executor::PreparedTask,
) -> Result<()> {
let save = Save::prepare(workspace, task, warm, prepared)?;
let cache = Cache {
remote: self.remote.clone(),
..Cache::new(self.root.clone())
};
let handle = std::thread::spawn(move || {
if let Err(error) = cache.save_prepared(&save) {
qk_executor::status!("qk: {}: warm state not saved ({error:#})", save.task);
}
});
self.saves.lock().unwrap().push(handle);
Ok(())
}
fn save_prepared(&self, save: &Save) -> Result<()> {
let warm = &save.warm;
let _lock = self.lock_warm(warm)?;
let records = self.warm_records(warm);
let previous = records
.iter()
.filter(|(_, record)| record.worktree == save.worktree)
.min_by_key(|(_, record)| record.key != save.key)
.map(|(_, record)| record);
let noted: RestoredFiles = fs::read(&save.notes)
.ok()
.and_then(|bytes| serde_json::from_slice(&bytes).ok())
.unwrap_or_default();
let mut record = Record {
version: RECORD_VERSION,
task: warm.identity.clone(),
worktree: save.worktree.clone(),
commit: save.commit.clone(),
saved: now_ms(),
key: save.key.clone(),
groups: BTreeMap::new(),
};
let gone = |error: &std::io::Error| error.kind() == std::io::ErrorKind::NotFound;
for location in &save.locations {
if !location.base.exists() {
continue;
}
let known = previous.and_then(|previous| previous.groups.get(location.name));
let restored_files = noted.groups.get(location.name);
let mut group = Group::default();
let mut pending = Vec::new();
for path in location.paths()? {
paths::safe_parents(&location.base, &path)?;
let absolute = location.base.join(&path);
let metadata = match fs::symlink_metadata(&absolute) {
Err(error) if gone(&error) => continue,
result => result?,
};
let artifact = if metadata.file_type().is_symlink() {
let target = match fs::read_link(&absolute) {
Err(error) if gone(&error) => continue,
result => result?,
};
let target = target
.to_str()
.context("symlink target must be UTF-8")?
.to_owned();
validate_link(&path, &target)?;
Artifact::Symlink {
target,
directory: absolute.is_dir(),
}
} else if metadata.is_file() {
let now = stamp(&metadata);
let saved = known.and_then(|known| {
(known.stamps.get(&path) == Some(&now))
.then(|| known.artifacts.get(&path))
.flatten()
});
let restored = restored_files
.and_then(|files| files.get(&path))
.filter(|(stamp, _)| *stamp == now)
.map(|(_, artifact)| artifact);
let reused = saved.or(restored).filter(|artifact| match artifact {
Artifact::File { blob, .. } => self.root.join("blobs").join(blob).is_file(),
_ => false,
});
let Some(artifact) = reused.cloned() else {
pending.push((path, absolute, metadata));
continue;
};
group.stamps.insert(path.clone(), now);
if let Some(modified) = metadata.modified().ok().and_then(nanos) {
group.mtimes.insert(path.clone(), modified);
}
artifact
} else if metadata.is_dir() {
Artifact::Directory {
mode: mode(&metadata),
}
} else {
continue;
};
group.artifacts.insert(path, artifact);
}
let sources: Vec<_> = pending
.iter()
.map(|(_, absolute, _)| absolute.clone())
.collect();
for ((path, _, metadata), result) in pending.into_iter().zip(self.put_blobs(&sources)) {
let blob = match result {
Ok(blob) => blob,
Err(error) if error.downcast_ref::<std::io::Error>().is_some_and(gone) => {
continue;
}
Err(error) => return Err(error),
};
group.stamps.insert(path.clone(), stamp(&metadata));
if let Some(modified) = metadata.modified().ok().and_then(nanos) {
group.mtimes.insert(path.clone(), modified);
}
group.artifacts.insert(
path,
Artifact::File {
blob,
mode: mode(&metadata),
},
);
}
if let Some(limit) = warm.max_size
&& group.size(self) > limit
{
qk_executor::status!(
"qk: {}: warm {} is over its maxSize and was not saved",
save.task,
location.name
);
continue;
}
record.groups.insert(location.name.to_owned(), group);
}
self.store_warm(&record)?;
let _ = fs::remove_file(&save.notes);
let saved = self.warm_record(&record);
for (path, _) in records
.into_iter()
.filter(|(path, _)| *path != saved)
.skip(KEPT - 1)
{
let _ = fs::remove_file(path);
}
if warm.remote
&& warm.portable
&& let Some(remote) = &self.remote
&& let Some(branch) = &save.branch
{
for group in record.groups.values_mut() {
group.stamps.clear();
}
remote.upload_warm(
&self.root,
&warm.identity,
branch,
serde_json::to_vec(&record)?,
);
}
Ok(())
}
}
struct Save {
task: String,
warm: Warm,
locations: Vec<Location>,
worktree: String,
commit: Option<String>,
branch: Option<String>,
key: Vec<String>,
notes: PathBuf,
}
impl Save {
fn prepare(
workspace: &Workspace,
task: &Task,
warm: &Warm,
prepared: &qk_executor::PreparedTask,
) -> Result<Self> {
Ok(Self {
task: task.id.clone(),
warm: warm.clone(),
locations: locations(workspace, task, warm)?,
worktree: worktree(workspace),
commit: current_commit(workspace),
branch: current_branch(workspace),
key: key_digests(workspace, warm, prepared)?,
notes: restored_files_path(workspace, warm),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use qk_taskgraph::Request;
use serde_json::json;
fn fixture() -> (tempfile::TempDir, Workspace, TaskGraph) {
let temp = tempfile::tempdir().unwrap();
fs::write(temp.path().join("nx.json"), "{}").unwrap();
fs::write(
temp.path().join("project.json"),
json!({
"name": "app", "targets": {
"one": {"command": "true", "qk:warm": {"paths": ["scratch-one"]}},
"two": {"command": "true", "qk:warm": {"paths": ["scratch-two"]}}
}
})
.to_string(),
)
.unwrap();
for path in ["scratch-one", "scratch-two"] {
fs::create_dir(temp.path().join(path)).unwrap();
fs::write(temp.path().join(path).join("state"), "saved").unwrap();
}
let workspace = Workspace::load(temp.path()).unwrap();
let graph = TaskGraph::build(
&workspace,
&[
Request::parse("app:one").unwrap(),
Request::parse("app:two").unwrap(),
],
)
.unwrap();
(temp, workspace, graph)
}
#[test]
fn waiting_for_an_owner_does_not_block_other_releases() {
let (_temp, workspace, graph) = fixture();
let one = &graph.tasks["app:one"];
let two = &graph.tasks["app:two"];
let first = acquire(&workspace, one, &config(&workspace, one).unwrap().unwrap()).unwrap();
let state = paths::worktree_state(&workspace.root).join("warm-kept");
let lock = fs::OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(state.join(format!("{}.lock", hash32(&two.id))))
.unwrap();
fs2::FileExt::lock_exclusive(&lock).unwrap();
let (released, received) = std::sync::mpsc::channel();
let completed = std::thread::scope(|scope| {
scope.spawn(|| {
let second =
acquire(&workspace, two, &config(&workspace, two).unwrap().unwrap()).unwrap();
release(&second);
});
std::thread::sleep(Duration::from_millis(100));
scope.spawn(|| {
release(&first);
released.send(()).unwrap();
});
let completed = received.recv_timeout(Duration::from_secs(2)).is_ok();
fs2::FileExt::unlock(&lock).unwrap();
completed
});
assert!(
completed,
"a blocked owner lock prevented an unrelated release"
);
assert_eq!(
fs::read_to_string(workspace.root.join("scratch-one/state")).unwrap(),
"saved"
);
}
#[test]
fn unreadable_ownership_preserves_saved_paths() {
let (_temp, workspace, graph) = fixture();
let owner = &graph.tasks["app:one"];
let warm = config(&workspace, owner).unwrap().unwrap();
let root = acquire(&workspace, owner, &warm).unwrap();
let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
drop(entry); fs::write(Stash::manifest(&root), "{").unwrap();
assert!(acquire(&workspace, owner, &warm).is_err());
assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
fs::remove_file(Stash::manifest(&root)).unwrap();
assert!(acquire(&workspace, owner, &warm).is_err());
assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
fs::create_dir(Stash::manifest(&root)).unwrap();
assert!(acquire(&workspace, owner, &warm).is_err());
assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
}
#[cfg(unix)]
#[test]
fn inspection_errors_keep_recovery_incomplete() {
let (_temp, workspace, graph) = fixture();
let owner = &graph.tasks["app:one"];
let warm = config(&workspace, owner).unwrap().unwrap();
let root = acquire(&workspace, owner, &warm).unwrap();
let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
let saved = root.with_extension("saved");
fs::rename(&root, &saved).unwrap();
fs::write(&root, "blocked").unwrap(); assert!(!entry.stash.put_back());
let retained = Stash {
root: root.clone(),
workspace: workspace.root.clone(),
moved: entry.stash.moved.clone(),
};
assert!(stash(&workspace, &root, &warm, Some(retained)).is_err());
assert_eq!(fs::read_to_string(saved.join("0/state")).unwrap(), "saved");
assert_eq!(fs::read_to_string(&root).unwrap(), "blocked");
}
}