use std::fs::{self, File, OpenOptions, TryLockError};
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
use crate::install::event::InstallProgress;
use crate::install::plan::InstallPlan;
use crate::install::provider::InstallProviderId;
use crate::persistence::{self, StoreError};
const JOB_FILE: &str = "job.json";
const STATUS_FILE: &str = "status.json";
const EVENTS_FILE: &str = "events.jsonl";
const LOCK_FILE: &str = "lock";
const CONTROL_FILE: &str = "control";
const SLUG_LIMIT: usize = 40;
pub const START_GRACE_MS: i64 = 3_000;
const CREATE_ATTEMPTS: u32 = 64;
const LOCK_REOPENS: u32 = 3;
#[derive(Debug, thiserror::Error)]
pub enum PullError {
#[error("pull io error: {0}")]
Io(#[from] io::Error),
#[error("pull store error: {0}")]
Store(#[from] StoreError),
#[error("unreadable pull descriptor at {path}: {source}")]
Unreadable {
path: PathBuf,
#[source]
source: serde_json::Error,
},
#[error("no pull matches \"{0}\"")]
NotFound(String),
#[error("\"{query}\" matches {count} pulls")]
Ambiguous {
query: String,
count: usize,
},
#[error("could not claim a directory for a pull of {0}")]
Unclaimable(String),
#[error("{id} is {state}, not ended")]
NotEnded {
id: String,
state: PullState,
},
#[error("a worker still holds {0}")]
Held(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum PullState {
Queued,
Running,
Paused,
Done,
Failed,
Cancelled,
Interrupted,
}
impl PullState {
pub fn is_terminal(self) -> bool {
matches!(self, Self::Done | Self::Failed | Self::Cancelled)
}
pub fn is_resumable(self) -> bool {
matches!(self, Self::Paused | Self::Interrupted)
}
pub fn is_live(self) -> bool {
matches!(self, Self::Queued | Self::Running)
}
pub fn as_str(self) -> &'static str {
match self {
Self::Queued => "queued",
Self::Running => "running",
Self::Paused => "paused",
Self::Done => "done",
Self::Failed => "failed",
Self::Cancelled => "cancelled",
Self::Interrupted => "interrupted",
}
}
}
impl std::fmt::Display for PullState {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(self.as_str())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum PullControl {
Pause,
Cancel,
}
impl PullControl {
pub fn as_str(self) -> &'static str {
match self {
Self::Pause => "pause",
Self::Cancel => "cancel",
}
}
pub fn parse(text: &str) -> Option<Self> {
match text.trim() {
"pause" => Some(Self::Pause),
"cancel" => Some(Self::Cancel),
_ => None,
}
}
pub fn resulting_state(self) -> PullState {
match self {
Self::Pause => PullState::Paused,
Self::Cancel => PullState::Cancelled,
}
}
}
impl std::fmt::Display for PullControl {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PullJob {
pub id: String,
pub provider: InstallProviderId,
pub reference: String,
pub display_name: String,
pub destination: String,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub revision: Option<String>,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub total_bytes: Option<i64>,
pub created_at_ms: i64,
}
impl PullJob {
pub(crate) fn from_plan(plan: &InstallPlan, id: impl Into<String>, created_at_ms: i64) -> Self {
Self {
id: id.into(),
provider: plan.provider.clone(),
reference: plan.reference.clone(),
display_name: plan.display_name.clone(),
destination: plan.destination.clone(),
revision: plan.revision.clone(),
total_bytes: plan.total_bytes,
created_at_ms,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PullStatus {
pub state: PullState,
#[serde(default)]
pub progress: InstallProgress,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub status_line: Option<String>,
#[serde(default)]
pub attempt: u32,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub next_attempt_at_ms: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub message: Option<String>,
#[serde(skip_serializing_if = "Option::is_none", default)]
pub pid: Option<u32>,
pub updated_at_ms: i64,
}
impl PullStatus {
pub fn past_stopping(&self) -> bool {
self.state == PullState::Running && self.status_line.as_deref() == Some(REGISTERING_LINE)
}
pub fn queued(now: i64) -> Self {
Self {
state: PullState::Queued,
progress: InstallProgress::default(),
status_line: None,
attempt: 0,
next_attempt_at_ms: None,
message: None,
pid: None,
updated_at_ms: now,
}
}
}
pub const REGISTERING_LINE: &str = "registering";
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PullEvent {
pub at_ms: i64,
#[serde(flatten)]
pub kind: PullEventKind,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "event", rename_all = "lowercase")]
pub enum PullEventKind {
State {
state: PullState,
},
Status {
text: String,
},
Retry {
attempt: u32,
reason: String,
delay_ms: i64,
},
}
#[derive(Debug)]
pub struct PullLock {
file: File,
}
impl Drop for PullLock {
fn drop(&mut self) {
let _ = self.file.unlock();
}
}
#[derive(Debug, Clone)]
pub struct PullStore {
root: PathBuf,
}
impl PullStore {
pub fn new(root: impl Into<PathBuf>) -> Self {
Self { root: root.into() }
}
pub fn root(&self) -> &Path {
&self.root
}
pub fn create(&self, plan: &InstallPlan, now: i64) -> Result<PullJobDir, PullError> {
let base = format!("{now}-{}", reference_slug(&plan.reference));
fs::create_dir_all(&self.root)?;
for attempt in 0..CREATE_ATTEMPTS {
let id = match attempt {
0 => base.clone(),
_ => format!("{base}-{}", attempt + 1),
};
let path = self.root.join(&id);
match fs::create_dir(&path) {
Ok(()) => {
let job = PullJob::from_plan(plan, id, now);
persistence::write_json_atomic(&path.join(JOB_FILE), &job)?;
let handle = PullJobDir { path, job };
handle.write_status(&PullStatus::queued(now))?;
return Ok(handle);
}
Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
Err(error) => return Err(error.into()),
}
}
Err(PullError::Unclaimable(plan.reference.clone()))
}
pub fn open(&self, id: &str) -> Result<PullJobDir, PullError> {
PullJobDir::open(self.root.join(id))
}
pub fn list(&self) -> Vec<PullJobDir> {
self.jobs().unwrap_or_default()
}
pub fn jobs(&self) -> Result<Vec<PullJobDir>, PullError> {
let entries = match fs::read_dir(&self.root) {
Ok(entries) => entries,
Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(error) => return Err(error.into()),
};
let mut jobs: Vec<PullJobDir> = entries
.flatten()
.filter(|entry| entry.path().is_dir())
.filter_map(|entry| PullJobDir::open(entry.path()).ok())
.collect();
jobs.sort_by(|left, right| {
left.job
.created_at_ms
.cmp(&right.job.created_at_ms)
.then_with(|| left.job.id.cmp(&right.job.id))
});
Ok(jobs)
}
pub fn resolve(&self, query: &str) -> Result<PullJobDir, PullError> {
let query = query.trim();
if query.is_empty() {
return Err(PullError::NotFound(String::new()));
}
let jobs = self.jobs()?;
if let Some(job) = jobs.iter().find(|job| job.job.id == query) {
return Ok(job.clone());
}
let by_prefix = jobs.iter().filter(|job| job.job.id.starts_with(query));
if let Some(job) = single(by_prefix, query)? {
return Ok(job.clone());
}
let by_reference = jobs
.iter()
.filter(|job| job.job.reference.eq_ignore_ascii_case(query));
match single(by_reference, query)? {
Some(job) => Ok(job.clone()),
None => Err(PullError::NotFound(query.to_owned())),
}
}
pub fn under_way(
&self,
provider: &InstallProviderId,
reference: &str,
now_ms: i64,
) -> Option<PullJobDir> {
self.list()
.into_iter()
.rev()
.filter(|job| {
job.job().provider == *provider
&& job.job().reference.eq_ignore_ascii_case(reference)
})
.find(|job| !job.status().state.is_terminal() && !job.abandoned(now_ms, START_GRACE_MS))
}
pub fn sweep(&self, keep: usize, before_ms: i64) -> usize {
let mut ended: Vec<(i64, PullJobDir)> = self
.list()
.into_iter()
.filter(|job| !job.worker_alive())
.filter_map(|job| {
let status = job.stored_status();
status
.state
.is_terminal()
.then_some((status.updated_at_ms, job))
})
.collect();
ended.sort_by_key(|(touched_at, _)| std::cmp::Reverse(*touched_at));
let mut removed = 0;
for (touched_at, job) in ended.into_iter().skip(keep) {
if touched_at < before_ms && job.forget().is_ok() {
removed += 1;
}
}
removed
}
}
fn single<'a>(
found: impl Iterator<Item = &'a PullJobDir>,
query: &str,
) -> Result<Option<&'a PullJobDir>, PullError> {
let matches: Vec<&PullJobDir> = found.collect();
match matches.as_slice() {
[] => Ok(None),
[only] => Ok(Some(only)),
many => {
let going: Vec<&PullJobDir> = many
.iter()
.copied()
.filter(|job| !job.status().state.is_terminal())
.collect();
match going.as_slice() {
[only] => Ok(Some(only)),
_ => Err(PullError::Ambiguous {
query: query.to_owned(),
count: many.len(),
}),
}
}
}
}
#[derive(Debug, Clone)]
pub struct PullJobDir {
path: PathBuf,
job: PullJob,
}
impl PullJobDir {
pub fn open(path: impl Into<PathBuf>) -> Result<Self, PullError> {
let path = path.into();
let descriptor = path.join(JOB_FILE);
let bytes = match fs::read(&descriptor) {
Ok(bytes) => bytes,
Err(error) if error.kind() == io::ErrorKind::NotFound => {
return Err(PullError::NotFound(name_of(&path)));
}
Err(error) => return Err(error.into()),
};
match serde_json::from_slice(&bytes) {
Ok(job) => Ok(Self { path, job }),
Err(source) => Err(PullError::Unreadable {
path: descriptor,
source,
}),
}
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn id(&self) -> &str {
&self.job.id
}
pub fn job(&self) -> &PullJob {
&self.job
}
pub fn lock_path(&self) -> PathBuf {
self.path.join(LOCK_FILE)
}
pub fn claim(&self) -> Result<Option<PullLock>, PullError> {
take_lock(&self.lock_path())
}
pub fn worker_alive(&self) -> bool {
let file = match File::open(self.lock_path()) {
Ok(file) => file,
Err(error) if error.kind() == io::ErrorKind::NotFound => return false,
Err(_) => return true,
};
match file.try_lock_shared() {
Ok(()) => {
let _ = file.unlock();
false
}
Err(_) => true,
}
}
pub fn stored_status(&self) -> PullStatus {
fs::read(self.path.join(STATUS_FILE))
.ok()
.and_then(|bytes| serde_json::from_slice(&bytes).ok())
.unwrap_or_else(|| PullStatus::queued(self.job.created_at_ms))
}
pub fn status(&self) -> PullStatus {
let mut status = self.stored_status();
let expects_worker = match status.state {
PullState::Running => true,
PullState::Queued => status.pid.is_some(),
_ => false,
};
if expects_worker && !self.worker_alive() {
status.state = PullState::Interrupted;
}
status
}
pub fn abandoned(&self, now_ms: i64, grace_ms: i64) -> bool {
self.abandoned_by(&self.stored_status(), now_ms, grace_ms)
}
pub fn abandoned_by(&self, status: &PullStatus, now_ms: i64, grace_ms: i64) -> bool {
status.state == PullState::Queued
&& status.pid.is_none()
&& now_ms.saturating_sub(status.updated_at_ms) >= grace_ms
&& !self.worker_alive()
}
pub fn write_status(&self, status: &PullStatus) -> Result<(), PullError> {
self.guard_write(|| {
persistence::write_json_atomic(&self.path.join(STATUS_FILE), status)?;
Ok(())
})
}
pub fn update_status(
&self,
now: i64,
change: impl FnOnce(&mut PullStatus),
) -> Result<PullStatus, PullError> {
let mut status = self.stored_status();
change(&mut status);
status.updated_at_ms = now;
self.write_status(&status)?;
Ok(status)
}
pub fn append(&self, kind: PullEventKind, now: i64) -> Result<(), PullError> {
let event = PullEvent { at_ms: now, kind };
let mut line = serde_json::to_vec(&event).map_err(StoreError::Encode)?;
line.push(b'\n');
let path = self.path.join(EVENTS_FILE);
let unterminated = ends_mid_line(&path);
let mut file = OpenOptions::new().create(true).append(true).open(&path)?;
if unterminated {
file.write_all(b"\n")?;
}
file.write_all(&line)?;
Ok(())
}
pub fn events(&self) -> Vec<PullEvent> {
let Ok(bytes) = fs::read(self.path.join(EVENTS_FILE)) else {
return Vec::new();
};
bytes
.split(|byte| *byte == b'\n')
.filter_map(|line| serde_json::from_slice(line).ok())
.collect()
}
pub fn control(&self) -> Option<PullControl> {
fs::read_to_string(self.path.join(CONTROL_FILE))
.ok()
.as_deref()
.and_then(PullControl::parse)
}
pub fn request(&self, control: PullControl) -> Result<(), PullError> {
self.guard_write(|| {
persistence::write_atomic(&self.path.join(CONTROL_FILE), control.as_str().as_bytes())?;
Ok(())
})
}
pub fn clear_control(&self, honoured: PullControl) -> Result<(), PullError> {
if self.control() != Some(honoured) {
return Ok(());
}
match fs::remove_file(self.path.join(CONTROL_FILE)) {
Ok(()) => Ok(()),
Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(error.into()),
}
}
pub fn remove(&self) -> Result<(), PullError> {
fs::remove_dir_all(&self.path)?;
Ok(())
}
pub fn forget(&self) -> Result<(), PullError> {
let state = self.stored_status().state;
if !state.is_terminal() {
return Err(PullError::NotEnded {
id: self.id().to_owned(),
state,
});
}
if self.worker_alive() {
return Err(PullError::Held(self.id().to_owned()));
}
self.remove()
}
fn guard_write(&self, write: impl FnOnce() -> Result<(), PullError>) -> Result<(), PullError> {
self.require_job()?;
write()?;
match self.require_job() {
Ok(()) => Ok(()),
Err(error) => {
let _ = fs::remove_dir_all(&self.path);
Err(error)
}
}
}
fn require_job(&self) -> Result<(), PullError> {
match self.path.join(JOB_FILE).is_file() {
true => Ok(()),
false => Err(PullError::NotFound(self.job.id.clone())),
}
}
}
pub fn take_lock(path: &Path) -> Result<Option<PullLock>, PullError> {
for _ in 0..LOCK_REOPENS {
let file = OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(path)?;
match file.try_lock() {
Ok(()) if still_at(&file, path) => return Ok(Some(PullLock { file })),
Ok(()) => continue,
Err(TryLockError::WouldBlock) => return Ok(None),
Err(TryLockError::Error(error)) => return Err(error.into()),
}
}
Ok(None)
}
#[cfg(unix)]
fn still_at(file: &File, path: &Path) -> bool {
use std::os::unix::fs::MetadataExt;
match (file.metadata(), fs::metadata(path)) {
(Ok(held), Ok(at_path)) => held.ino() == at_path.ino() && held.dev() == at_path.dev(),
_ => false,
}
}
#[cfg(not(unix))]
fn still_at(_file: &File, _path: &Path) -> bool {
true
}
fn ends_mid_line(path: &Path) -> bool {
fs::read(path)
.ok()
.filter(|bytes| !bytes.is_empty())
.is_some_and(|bytes| bytes.last() != Some(&b'\n'))
}
fn name_of(path: &Path) -> String {
path.file_name()
.and_then(|name| name.to_str())
.unwrap_or("pull")
.to_owned()
}
fn reference_slug(reference: &str) -> String {
let mut slug = String::with_capacity(reference.len().min(SLUG_LIMIT));
for character in reference.chars() {
if character.is_ascii_alphanumeric() {
slug.push(character.to_ascii_lowercase());
} else if !slug.ends_with('-') {
slug.push('-');
}
if slug.len() >= SLUG_LIMIT {
break;
}
}
let trimmed = slug.trim_matches('-');
match trimmed.is_empty() {
true => "model".to_owned(),
false => trimmed.to_owned(),
}
}