use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use a3s_box_core::error::{BoxError, Result};
use a3s_box_core::{ImageStoreBackend, StoredImage};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
static PUT_SEQ: AtomicU64 = AtomicU64::new(0);
#[derive(Debug, Default, Serialize, Deserialize)]
struct StoreIndex {
images: Vec<StoredImage>,
}
pub struct ImageStore {
store_dir: PathBuf,
index: Arc<RwLock<HashMap<String, StoredImage>>>,
max_size_bytes: u64,
}
impl ImageStore {
pub fn new(store_dir: &Path, max_size_bytes: u64) -> Result<Self> {
std::fs::create_dir_all(store_dir).map_err(|e| {
BoxError::OciImageError(format!(
"Failed to create image store directory {}: {}",
store_dir.display(),
e
))
})?;
let tmp_dir = store_dir.join("tmp");
if tmp_dir.is_dir() {
if let Err(e) = std::fs::remove_dir_all(&tmp_dir) {
tracing::debug!(path = %tmp_dir.display(), error = %e, "Failed to sweep image store tmp dir");
}
}
let mut store = Self {
store_dir: store_dir.to_path_buf(),
index: Arc::new(RwLock::new(HashMap::new())),
max_size_bytes,
};
store.load_index()?;
Ok(store)
}
pub async fn get(&self, reference: &str) -> Option<StoredImage> {
let mut index = self.index.write().await;
if let Some(image) = index.get_mut(reference) {
image.last_used = Utc::now();
let updated = image.clone();
drop(index);
if let Err(e) = self.save_index_inner().await {
tracing::warn!(error = %e, "Failed to persist image store index (last_used may be stale)");
}
Some(updated)
} else {
None
}
}
pub async fn get_by_digest(&self, digest: &str) -> Option<StoredImage> {
let mut index = self.index.write().await;
let found = index.values_mut().find(|img| img.digest == digest);
if let Some(image) = found {
image.last_used = Utc::now();
let updated = image.clone();
drop(index);
if let Err(e) = self.save_index_inner().await {
tracing::warn!(error = %e, "Failed to persist image store index (last_used may be stale)");
}
Some(updated)
} else {
None
}
}
pub async fn resolve(&self, image: &str) -> Option<StoredImage> {
if let Some(found) = self.get(image).await {
return Some(found);
}
let digest_part = image.rsplit_once('@').map_or(image, |(_, digest)| digest);
if let Some(found) = self.get_by_digest(digest_part).await {
return Some(found);
}
match super::ImageReference::parse(image) {
Ok(parsed) => self.get(&parsed.full_reference()).await,
Err(_) => None,
}
}
pub async fn put(
&self,
reference: &str,
digest: &str,
source_dir: &Path,
) -> Result<StoredImage> {
let digest_hex = digest.strip_prefix("sha256:").unwrap_or(digest);
let target_dir = self.store_dir.join("sha256").join(digest_hex);
if !target_dir.exists() {
let seq = PUT_SEQ.fetch_add(1, Ordering::Relaxed);
let staging = self.store_dir.join("sha256").join(format!(
".staging-{}-{}-{}",
digest_hex,
std::process::id(),
seq
));
let _ = std::fs::remove_dir_all(&staging);
copy_dir_recursive(source_dir, &staging).map_err(|e| {
let _ = std::fs::remove_dir_all(&staging);
BoxError::OciImageError(format!("Failed to copy image to store: {}", e))
})?;
if let Err(e) = std::fs::rename(&staging, &target_dir) {
let _ = std::fs::remove_dir_all(&staging);
if !target_dir.exists() {
return Err(BoxError::OciImageError(format!(
"Failed to publish image to store: {}",
e
)));
}
}
}
let size_bytes = dir_size(&target_dir);
let now = Utc::now();
let stored = StoredImage {
reference: reference.to_string(),
digest: digest.to_string(),
size_bytes,
pulled_at: now,
last_used: now,
path: target_dir,
};
self.with_index_lock(|index| {
if let Some(old) = index.get(reference).cloned() {
if old.digest != digest {
let still_referenced = index
.iter()
.any(|(key, img)| key.as_str() != reference && img.digest == old.digest);
if !still_referenced && !index.contains_key(&old.digest) {
let mut dangling = old.clone();
dangling.reference = old.digest.clone();
index.insert(old.digest.clone(), dangling);
}
}
}
index.insert(reference.to_string(), stored.clone());
Ok(())
})
.await?;
Ok(stored)
}
pub async fn remove(&self, image: &str) -> Result<()> {
self.with_index_lock(|index| {
let keys: Vec<String> = if index.contains_key(image) {
vec![image.to_string()]
} else {
index
.values()
.filter(|img| img.digest == image)
.map(|img| img.reference.clone())
.collect()
};
if keys.is_empty() {
return Err(BoxError::OciImageError(format!(
"Image not found: {}",
image
)));
}
let removed: Vec<StoredImage> = keys.iter().filter_map(|k| index.remove(k)).collect();
for img in removed {
let digest_still_used = index.values().any(|other| other.digest == img.digest);
if !digest_still_used && img.path.exists() {
std::fs::remove_dir_all(&img.path).map_err(|e| {
BoxError::OciImageError(format!(
"Failed to remove image directory {}: {}",
img.path.display(),
e
))
})?;
}
}
Ok(())
})
.await
}
pub async fn list(&self) -> Vec<StoredImage> {
let index = self.index.read().await;
index.values().cloned().collect()
}
pub async fn evict(&self) -> Result<Vec<String>> {
let mut evicted = Vec::new();
let mut total = self.total_size().await;
while total > self.max_size_bytes {
let lru_ref = {
let index = self.index.read().await;
index
.values()
.min_by_key(|img| img.last_used)
.map(|img| img.reference.clone())
};
match lru_ref {
Some(reference) => {
self.remove(&reference).await?;
evicted.push(reference);
total = self.total_size().await;
}
None => break,
}
}
Ok(evicted)
}
pub async fn total_size(&self) -> u64 {
let index = self.index.read().await;
index.values().map(|img| img.size_bytes).sum()
}
fn load_index(&mut self) -> Result<()> {
self.index = Arc::new(RwLock::new(self.read_index_from_disk()?));
Ok(())
}
fn read_index_from_disk(&self) -> Result<HashMap<String, StoredImage>> {
let index_path = self.store_dir.join("index.json");
if !index_path.exists() {
return Ok(HashMap::new());
}
let data = std::fs::read_to_string(&index_path).map_err(|e| {
BoxError::OciImageError(format!(
"Failed to read image store index {}: {}",
index_path.display(),
e
))
})?;
#[derive(serde::Deserialize)]
struct RawIndex {
#[serde(default)]
images: Vec<serde_json::Value>,
}
let raw: RawIndex = match serde_json::from_str(&data) {
Ok(raw) => raw,
Err(err) => {
let preserved = crate::store_io::quarantine_label(&index_path);
tracing::warn!(
"image store index {} is corrupt ({err}); preserved a copy at \
{preserved} and started from an empty catalog (re-pulled images \
will repopulate it)",
index_path.display(),
);
return Ok(HashMap::new());
}
};
let mut index = HashMap::new();
let mut skipped = 0usize;
for value in raw.images {
match serde_json::from_value::<StoredImage>(value) {
Ok(image) => {
if image.path.exists() {
index.insert(image.reference.clone(), image);
}
}
Err(err) => {
skipped += 1;
tracing::warn!("skipping unreadable image index entry ({err})");
}
}
}
if skipped > 0 {
let preserved = crate::store_io::quarantine_copy(&index_path)
.map(|p| p.display().to_string())
.unwrap_or_else(|| "<backup failed>".to_string());
tracing::warn!(
"{skipped} image index entr{} skipped as unreadable; preserved a copy at \
{preserved}; affected images will be re-pulled on demand",
if skipped == 1 { "y" } else { "ies" },
);
}
Ok(index)
}
async fn with_index_lock<F, R>(&self, f: F) -> Result<R>
where
F: FnOnce(&mut HashMap<String, StoredImage>) -> Result<R>,
{
let index_path = self.store_dir.join("index.json");
let _lock = {
let p = index_path.clone();
tokio::task::spawn_blocking(move || crate::file_lock::FileLock::acquire(&p))
.await
.map_err(|e| BoxError::OciImageError(format!("index lock task failed: {e}")))?
.map_err(|e| {
BoxError::OciImageError(format!(
"failed to lock image index {}: {e}",
index_path.display()
))
})?
};
let fresh = self.read_index_from_disk()?;
let result = {
let mut idx = self.index.write().await;
*idx = fresh;
f(&mut idx)?
};
self.save_index_inner().await?;
Ok(result)
}
async fn save_index_inner(&self) -> Result<()> {
let index = self.index.read().await;
let store_index = StoreIndex {
images: index.values().cloned().collect(),
};
drop(index);
let data = serde_json::to_string_pretty(&store_index)?;
let index_path = self.store_dir.join("index.json");
let tmp_path = self.store_dir.join("index.json.tmp");
tokio::fs::write(&tmp_path, data).await.map_err(|e| {
BoxError::OciImageError(format!(
"Failed to write image store index {}: {}",
tmp_path.display(),
e
))
})?;
tokio::fs::rename(&tmp_path, &index_path)
.await
.map_err(|e| {
BoxError::OciImageError(format!(
"Failed to commit image store index {}: {}",
index_path.display(),
e
))
})?;
Ok(())
}
pub fn store_dir(&self) -> &Path {
&self.store_dir
}
}
#[async_trait::async_trait]
impl ImageStoreBackend for ImageStore {
async fn get(&self, reference: &str) -> Option<StoredImage> {
self.get(reference).await
}
async fn get_by_digest(&self, digest: &str) -> Option<StoredImage> {
self.get_by_digest(digest).await
}
async fn put(&self, reference: &str, digest: &str, source_dir: &Path) -> Result<StoredImage> {
self.put(reference, digest, source_dir).await
}
async fn remove(&self, reference: &str) -> Result<()> {
self.remove(reference).await
}
async fn list(&self) -> Vec<StoredImage> {
self.list().await
}
async fn evict(&self) -> Result<Vec<String>> {
self.evict().await
}
async fn total_size(&self) -> u64 {
self.total_size().await
}
}
fn copy_dir_recursive(src: &Path, dst: &Path) -> std::io::Result<()> {
std::fs::create_dir_all(dst)?;
for entry in std::fs::read_dir(src)? {
let entry = entry?;
let src_path = entry.path();
let dst_path = dst.join(entry.file_name());
if src_path.is_dir() {
copy_dir_recursive(&src_path, &dst_path)?;
} else {
std::fs::copy(&src_path, &dst_path)?;
}
}
Ok(())
}
fn dir_size(path: &Path) -> u64 {
let mut total = 0;
if let Ok(entries) = std::fs::read_dir(path) {
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
total += dir_size(&path);
} else if let Ok(meta) = path.metadata() {
total += meta.len();
}
}
}
total
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn create_test_oci_layout(dir: &Path) {
std::fs::create_dir_all(dir.join("blobs/sha256")).unwrap();
std::fs::write(dir.join("oci-layout"), r#"{"imageLayoutVersion":"1.0.0"}"#).unwrap();
std::fs::write(dir.join("index.json"), r#"{"manifests":[]}"#).unwrap();
std::fs::write(dir.join("blobs/sha256/testblob"), "x".repeat(1024)).unwrap();
}
#[tokio::test]
async fn test_new_creates_directory() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("images");
let store = ImageStore::new(&store_dir, 1024 * 1024).unwrap();
assert!(store_dir.exists());
assert_eq!(store.total_size().await, 0);
}
#[tokio::test]
async fn test_put_and_get() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
let stored = store
.put("nginx:latest", "sha256:abc123", &source_dir)
.await
.unwrap();
assert_eq!(stored.reference, "nginx:latest");
assert_eq!(stored.digest, "sha256:abc123");
assert!(stored.size_bytes > 0);
assert!(stored.path.exists());
let fetched = store.get("nginx:latest").await.unwrap();
assert_eq!(fetched.digest, "sha256:abc123");
let fetched = store.get_by_digest("sha256:abc123").await.unwrap();
assert_eq!(fetched.reference, "nginx:latest");
}
#[tokio::test]
async fn test_get_nonexistent() {
let tmp = TempDir::new().unwrap();
let store = ImageStore::new(tmp.path(), 1024 * 1024).unwrap();
assert!(store.get("nonexistent").await.is_none());
}
#[tokio::test]
async fn test_remove() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
store
.put("nginx:latest", "sha256:abc123", &source_dir)
.await
.unwrap();
store.remove("nginx:latest").await.unwrap();
assert!(store.get("nginx:latest").await.is_none());
}
#[tokio::test]
async fn test_retag_keeps_displaced_image_as_dangling() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
store
.put("app:latest", "sha256:old", &source_dir)
.await
.unwrap();
store
.put("app:latest", "sha256:new", &source_dir)
.await
.unwrap();
assert_eq!(store.get("app:latest").await.unwrap().digest, "sha256:new");
let dangling = store.get("sha256:old").await.unwrap();
assert_eq!(dangling.digest, "sha256:old");
assert_eq!(store.list().await.len(), 2);
}
#[tokio::test]
async fn test_reput_same_digest_does_not_create_dangling() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
store
.put("app:latest", "sha256:same", &source_dir)
.await
.unwrap();
store
.put("app:latest", "sha256:same", &source_dir)
.await
.unwrap();
assert_eq!(store.list().await.len(), 1);
}
#[tokio::test]
async fn test_remove_by_digest() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
let stored = store
.put("gcr.io/test/img:test", "sha256:deadbeef", &source_dir)
.await
.unwrap();
let path = stored.path.clone();
store.remove("sha256:deadbeef").await.unwrap();
assert!(store.get("gcr.io/test/img:test").await.is_none());
assert!(store.get_by_digest("sha256:deadbeef").await.is_none());
assert!(!path.exists(), "on-disk layout should be deleted");
}
#[tokio::test]
async fn test_remove_by_digest_removes_all_tags() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
store
.put("img:v1", "sha256:shared", &source_dir)
.await
.unwrap();
let stored = store
.put("img:latest", "sha256:shared", &source_dir)
.await
.unwrap();
let path = stored.path.clone();
store.remove("sha256:shared").await.unwrap();
assert!(store.get("img:v1").await.is_none());
assert!(store.get("img:latest").await.is_none());
assert!(!path.exists(), "shared layout should be deleted");
}
#[tokio::test]
async fn test_resolve_by_name_digest_and_normalized() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
store
.put(
"gcr.io/x/test-image-predefined-group:latest",
"sha256:grp",
&source_dir,
)
.await
.unwrap();
assert!(store
.resolve("gcr.io/x/test-image-predefined-group:latest")
.await
.is_some());
assert_eq!(
store
.resolve("gcr.io/x/test-image-predefined-group")
.await
.map(|i| i.digest),
Some("sha256:grp".to_string())
);
assert!(store.resolve("sha256:grp").await.is_some());
assert!(store
.resolve("gcr.io/x/test-image-predefined-group@sha256:grp")
.await
.is_some());
assert!(store.resolve("nope:latest").await.is_none());
}
#[tokio::test]
async fn test_remove_nonexistent() {
let tmp = TempDir::new().unwrap();
let store = ImageStore::new(tmp.path(), 1024 * 1024).unwrap();
assert!(store.remove("nonexistent").await.is_err());
}
#[tokio::test]
async fn test_list() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
store
.put("nginx:latest", "sha256:aaa", &source_dir)
.await
.unwrap();
store
.put("alpine:3.18", "sha256:bbb", &source_dir)
.await
.unwrap();
let images = store.list().await;
assert_eq!(images.len(), 2);
}
#[tokio::test]
async fn test_total_size() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
store
.put("nginx:latest", "sha256:aaa", &source_dir)
.await
.unwrap();
assert!(store.total_size().await > 0);
}
#[tokio::test]
async fn test_lru_eviction() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let store = ImageStore::new(&store_dir, 100).unwrap();
store
.put("old:v1", "sha256:old1", &source_dir)
.await
.unwrap();
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
store
.put("new:v2", "sha256:new2", &source_dir)
.await
.unwrap();
store.get("new:v2").await;
let evicted = store.evict().await.unwrap();
assert!(!evicted.is_empty());
assert!(evicted.contains(&"old:v1".to_string()));
}
#[tokio::test]
async fn test_index_persistence() {
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
{
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
store
.put("nginx:latest", "sha256:persist", &source_dir)
.await
.unwrap();
}
{
let store = ImageStore::new(&store_dir, 10 * 1024 * 1024).unwrap();
let image = store.get("nginx:latest").await;
assert!(image.is_some());
assert_eq!(image.unwrap().digest, "sha256:persist");
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_cross_instance_puts_persist_both() {
use std::collections::HashSet;
use std::sync::Arc;
let tmp = TempDir::new().unwrap();
let store_dir = tmp.path().join("store");
let source_dir = tmp.path().join("source");
create_test_oci_layout(&source_dir);
let s1 = Arc::new(ImageStore::new(&store_dir, u64::MAX).unwrap());
let s2 = Arc::new(ImageStore::new(&store_dir, u64::MAX).unwrap());
let (src1, src2) = (source_dir.clone(), source_dir.clone());
let h1 = {
let s1 = Arc::clone(&s1);
tokio::spawn(async move { s1.put("img:a", "sha256:aaaa", &src1).await.unwrap() })
};
let h2 = {
let s2 = Arc::clone(&s2);
tokio::spawn(async move { s2.put("img:b", "sha256:bbbb", &src2).await.unwrap() })
};
h1.await.unwrap();
h2.await.unwrap();
let s3 = ImageStore::new(&store_dir, u64::MAX).unwrap();
let refs: HashSet<String> = s3.list().await.into_iter().map(|i| i.reference).collect();
assert!(refs.contains("img:a"), "img:a lost: {refs:?}");
assert!(refs.contains("img:b"), "img:b lost: {refs:?}");
}
#[tokio::test]
async fn corrupt_index_is_quarantined_not_fatal() {
let tmp = tempfile::tempdir().unwrap();
let store_dir = tmp.path().join("images");
std::fs::create_dir_all(&store_dir).unwrap();
std::fs::write(store_dir.join("index.json"), "{ not valid json").unwrap();
let store = ImageStore::new(&store_dir, u64::MAX)
.expect("corrupt index.json must not brick the image store");
assert!(
store.list().await.is_empty(),
"store must start from an empty catalog after quarantine"
);
let quarantined = std::fs::read_dir(&store_dir)
.unwrap()
.filter_map(|e| e.ok())
.any(|e| {
e.file_name()
.to_string_lossy()
.contains("index.json.corrupt-")
});
assert!(
quarantined,
"corrupt index.json must be quarantined to a sibling"
);
}
}