use std::collections::{HashMap, HashSet};
use crate::compose::types::{ComposeFile, Service};
use crate::engine::build::libpod_pull_policy;
use super::parallel::join_bounded;
use super::Engine;
impl Engine {
pub(super) async fn prefetch_images(
&self,
file: &ComposeFile,
enabled: &HashSet<String>,
target_set: &Option<HashSet<String>>,
) {
let mut by_image: HashMap<&str, &Service> = HashMap::new();
for (name, service) in &file.services {
if !enabled.contains(name) {
continue;
}
if let Some(set) = target_set {
if !set.contains(name) {
continue;
}
}
if service.build.is_some() && !self.no_build {
continue;
}
let Some(image) = service.image.as_deref() else {
continue;
};
let raw_policy = self
.pull_policy_override
.as_deref()
.or(service.pull_policy.as_deref());
if libpod_pull_policy(raw_policy).unwrap_or("missing") == "never" {
continue;
}
by_image.entry(image).or_insert(service);
}
let futs = by_image.into_values().map(|service| async move {
let image = service.image.as_deref().unwrap_or_default();
let raw_policy = self
.pull_policy_override
.as_deref()
.or(service.pull_policy.as_deref());
let policy = libpod_pull_policy(raw_policy).unwrap_or("missing");
if policy == "missing" && self.image_present(image).await {
if service.platform.is_none() {
if let Ok(mut seen) = self.images_seen_present.lock() {
seen.insert(image.to_string());
}
}
return;
}
if let Err(e) = self.pull_image(service).await {
tracing::debug!("prefetch miss for {image}: {e}");
}
});
join_bounded(futs).await;
}
}
#[cfg(test)]
mod tests {
#[cfg(unix)]
use std::collections::HashSet;
#[cfg(unix)]
use crate::engine::fake_podman;
#[cfg(unix)]
use crate::engine::Engine;
#[cfg(unix)]
fn engine_with(client: crate::libpod::Client, project: &str) -> Engine {
Engine::with_base_dir(client, project.into(), std::env::temp_dir())
}
#[tokio::test]
#[cfg(unix)]
async fn prefetch_dedupes_shared_image_and_skips_never_and_build_services() {
let fake = fake_podman::start(|method, target| {
if method == "POST" && target.contains("/images/pull") {
(200, String::new())
} else {
(404, r#"{"message":"not found"}"#.to_string())
}
});
let e = engine_with(fake.client(), "proj");
let file = crate::parse_str(
"services:\n a:\n image: shared\n b:\n image: shared\n c:\n image: skip-me\n pull_policy: never\n d:\n image: build-me\n build:\n context: .\n",
)
.unwrap();
let enabled: HashSet<String> = file.services.keys().cloned().collect();
e.prefetch_images(&file, &enabled, &None).await;
let seen = fake.requests.lock().unwrap();
let shared_pulls = seen
.iter()
.filter(|r| r.contains("/images/pull") && r.contains("reference=shared"))
.count();
assert_eq!(
shared_pulls, 1,
"two services sharing one image must pull it once: {seen:?}"
);
assert!(
!seen.iter().any(|r| r.contains("skip-me")),
"a never-policy service must not be prefetched: {seen:?}"
);
assert!(
!seen.iter().any(|r| r.contains("build-me")),
"a service with a build: section must not be prefetched: {seen:?}"
);
}
#[tokio::test]
#[cfg(unix)]
async fn prefetch_skips_services_outside_the_target_set() {
let fake = fake_podman::start(|method, target| {
if method == "POST" && target.contains("/images/pull") {
(200, String::new())
} else {
(404, r#"{"message":"not found"}"#.to_string())
}
});
let e = engine_with(fake.client(), "proj");
let file =
crate::parse_str("services:\n web:\n image: img-web\n db:\n image: img-db\n")
.unwrap();
let enabled: HashSet<String> = file.services.keys().cloned().collect();
let target_set: Option<HashSet<String>> = Some(["web".to_string()].into_iter().collect());
e.prefetch_images(&file, &enabled, &target_set).await;
let seen = fake.requests.lock().unwrap();
assert!(
seen.iter()
.any(|r| r.contains("/images/pull") && r.contains("reference=img-web")),
"the targeted service's image must be prefetched: {seen:?}"
);
assert!(
!seen.iter().any(|r| r.contains("img-db")),
"a service outside the target set must not be prefetched: {seen:?}"
);
}
}