use std::collections::{BTreeMap, BTreeSet};
use std::time::SystemTime;
use tokio::task::JoinHandle;
use tracing::{info, warn};
use super::Engine;
use super::runners::backoff;
use crate::config::{Image, PUBLISHED_RUNNER};
use crate::docker::{self, DockerApi, TaggedImage};
use crate::github::GithubApi;
#[derive(Debug)]
pub(super) enum ImageState {
Present,
Pulling {
task: JoinHandle<Result<(), docker::Error>>,
failures: u32,
},
Failed {
failures: u32,
retry_at: SystemTime,
},
}
#[derive(Debug)]
pub(super) struct Images {
states: BTreeMap<String, ImageState>,
pub(super) stale_due: bool,
pub(super) removing: Option<Removal>,
refused: BTreeSet<String>,
}
#[derive(Debug)]
pub(super) struct Removal {
tag: String,
pub(super) task: JoinHandle<Result<(), docker::Error>>,
}
impl Images {
pub(super) fn new() -> Images {
Images {
states: BTreeMap::new(),
stale_due: true,
removing: None,
refused: BTreeSet::new(),
}
}
}
impl<G: GithubApi + 'static, D: DockerApi + 'static> Engine<G, D> {
pub(super) async fn pull_images(&mut self) {
let finished: Vec<String> = self
.images
.states
.iter()
.filter(|(_, state)| matches!(state, ImageState::Pulling { task, .. } if task.is_finished()))
.map(|(image, _)| image.clone())
.collect();
for image in finished {
let Some(ImageState::Pulling { task, failures }) = self.images.states.remove(&image)
else {
continue;
};
let state = match task.await {
Ok(Ok(())) => {
self.images.stale_due = true;
ImageState::Present
}
Ok(Err(error)) => failed(&image, failures, &error.to_string()),
Err(error) => failed(&image, failures, &error.to_string()),
};
self.images.states.insert(image, state);
}
let wanted: Vec<String> = self
.config
.profiles
.values()
.filter_map(|profile| match &profile.image {
Image::Pulled(reference) => Some(reference.clone()),
Image::Stage(_) => None,
})
.collect();
let now = SystemTime::now();
for image in wanted {
let failures = match self.images.states.get(&image) {
None => 0,
Some(ImageState::Failed { failures, retry_at }) if *retry_at <= now => *failures,
Some(_) => continue,
};
match self.docker.image_present(&image).await {
Ok(true) => {
self.images.states.insert(image, ImageState::Present);
}
Ok(false) => {
let docker = self.docker.clone();
let reference = image.clone();
let task = tokio::spawn(async move { docker.pull(&reference).await });
self.images
.states
.insert(image, ImageState::Pulling { task, failures });
}
Err(error) => {
let state = failed(&image, failures, &error.to_string());
self.images.states.insert(image, state);
}
}
}
self.remove_stale_images().await;
}
async fn remove_stale_images(&mut self) {
if let Some(removal) = &self.images.removing {
if !removal.task.is_finished() {
return;
}
let Some(Removal { tag, task }) = self.images.removing.take() else {
return;
};
match task.await {
Ok(Ok(())) => {
info!("removed {tag}, which no profile runs");
self.images
.states
.retain(|_, state| !matches!(state, ImageState::Present));
}
Ok(Err(error)) => {
warn!("cannot remove {tag}: {error}");
self.images.refused.insert(tag);
}
Err(error) => {
warn!("removing {tag} stopped: {error}");
self.images.refused.insert(tag);
}
}
}
if !self.images.stale_due {
return;
}
let Some(tag) = self.next_stale_tag().await else {
self.images.stale_due = false;
self.images.refused.clear();
return;
};
let docker = self.docker.clone();
let reference = tag.clone();
let task = tokio::spawn(async move { docker.remove_image(&reference).await });
self.images.removing = Some(Removal { tag, task });
}
async fn next_stale_tag(&self) -> Option<String> {
let wanted: Vec<String> = self
.config
.profiles
.values()
.chain(self.retired.values())
.map(|profile| profile.image.reference())
.filter(|reference| !reference.is_empty())
.chain(
self.images
.states
.iter()
.filter(|(_, state)| matches!(state, ImageState::Pulling { .. }))
.map(|(reference, _)| reference.clone()),
)
.collect();
let listed = async {
let images = self.docker.tagged_images().await?;
let mut kept = self.docker.container_images().await?;
for reference in &wanted {
kept.extend(self.docker.image_id(reference).await?);
}
Ok::<_, docker::Error>((images, kept))
};
let (images, kept) = match listed.await {
Ok(listed) => listed,
Err(error) => {
warn!("cannot list Docker's images and containers: {error}");
return None;
}
};
let wanted: Vec<Reference> = wanted
.iter()
.map(|reference| Reference::parse(reference))
.collect();
images
.iter()
.filter(|image| !kept.contains(&image.id))
.filter(|image| !wanted.iter().any(|reference| reference.names(image)))
.flat_map(|image| &image.tags)
.find(|tag| {
Reference::parse(tag).name == PUBLISHED_RUNNER
&& !self.images.refused.contains(*tag)
})
.cloned()
}
pub(super) fn awaits_image(&self, image: &Image) -> bool {
match image {
Image::Stage(_) => false,
Image::Pulled(reference) => {
!matches!(self.images.states.get(reference), Some(ImageState::Present))
}
}
}
pub(super) fn recheck_image(&mut self, image: &str) {
if matches!(self.images.states.get(image), Some(ImageState::Present)) {
self.images.states.remove(image);
}
}
}
#[derive(Debug, PartialEq, Eq)]
struct Reference {
name: String,
tag: Option<String>,
digest: Option<String>,
}
impl Reference {
fn parse(reference: &str) -> Reference {
let (rest, digest) = match reference.split_once('@') {
Some((rest, digest)) => (rest, Some(digest.to_owned())),
None => (reference, None),
};
let (name, tag) = match rest.rsplit_once(':') {
Some((name, tag)) if !tag.contains('/') => (name, Some(tag.to_owned())),
_ => (rest, None),
};
let tag = tag.or_else(|| digest.is_none().then(|| "latest".to_owned()));
let (registry, path) = match name.split_once('/') {
Some((first, path))
if first.contains(['.', ':'])
|| first == "localhost"
|| first.chars().any(|c| c.is_ascii_uppercase()) =>
{
(first, path)
}
_ => ("docker.io", name),
};
let registry = if registry == "index.docker.io" {
"docker.io"
} else {
registry
};
let name = if registry == "docker.io" && !path.contains('/') {
format!("docker.io/library/{path}")
} else {
format!("{registry}/{path}")
};
Reference { name, tag, digest }
}
fn tag_ref(&self) -> (&str, Option<&str>) {
(&self.name, self.tag.as_deref())
}
fn names(&self, image: &TaggedImage) -> bool {
let by_tag = self.tag.is_some()
&& image
.tags
.iter()
.any(|tag| Reference::parse(tag).tag_ref() == self.tag_ref());
let by_digest = self.digest.is_some()
&& image.digests.iter().any(|digest| {
let digest = Reference::parse(digest);
digest.name == self.name && digest.digest == self.digest
});
by_tag || by_digest
}
}
fn failed(image: &str, failures: u32, error: &str) -> ImageState {
let failures = failures + 1;
let wait = backoff(failures);
warn!(
"cannot pull {image}: {error}; trying again in {}s",
wait.as_secs()
);
ImageState::Failed {
failures,
retry_at: SystemTime::now() + wait,
}
}
#[cfg(test)]
mod tests {
use super::Reference;
#[test]
fn docker_spellings_of_one_image_name_it_alike() {
for (spelled, listed) in [
("index.docker.io/acme/runner:app", "acme/runner:app"),
("docker.io/acme/runner", "acme/runner:latest"),
("docker.io/library/alpine:3", "alpine:3"),
("index.docker.io/alpine", "alpine:latest"),
("localhost:5000/runner:app", "localhost:5000/runner:app"),
] {
assert_eq!(
Reference::parse(spelled),
Reference::parse(listed),
"{spelled}"
);
}
for (one, other) in [
("ghcr.io/acme/runner:app", "acme/runner:app"),
("localhost/runner:app", "runner:app"),
] {
assert_ne!(Reference::parse(one), Reference::parse(other), "{one}");
}
}
}