Merge pull request #45 from AI45Lab/fix/box-oci-concurrent-copy

fix(runtime): atomic OCI store/build-cache writes (stage + rename)

Author: ZhiXiao-Lin

src/runtime/src/oci/build/cache.rs

@@ -15,12 +15,16 @@
 //! NOT fail the build.
 
 use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicU64, Ordering};
 
 use a3s_box_core::dirs_home;
 use serde::{Deserialize, Serialize};
 
 use super::layer::{sha256_bytes, LayerInfo};
 
+/// Per-process counter for unique staging-file names in `store`.
+static STORE_SEQ: AtomicU64 = AtomicU64::new(0);
+
 /// Default cap on the total size of cached layer blobs (2 GiB). Override with
 /// `A3S_BOX_BUILDCACHE_MAX_BYTES`. When the cap is exceeded after a store, the
 /// oldest blobs are evicted (FIFO) until the total is back under the cap; a
@@ -125,8 +129,27 @@ impl BuildCache {
     /// the `keys/<key>` record. Best-effort: I/O errors are ignored.
     pub(crate) fn store(&self, key: &str, layer: &LayerInfo, diff_id: &str) {
         let blob_path = self.dir.join("blobs").join(&layer.digest);
-        if !blob_path.exists() && std::fs::copy(&layer.path, &blob_path).is_err() {
-            return;
+        if !blob_path.exists() {
+            // Copy to a unique temp file then atomically rename into place, so a
+            // concurrent build (or a copy that fails partway) never publishes a
+            // half-written blob that a key record then points at. The blob is
+            // content-addressed, so a rename that overwrites a racing winner is
+            // harmless (identical bytes).
+            let seq = STORE_SEQ.fetch_add(1, Ordering::Relaxed);
+            let staging = self.dir.join("blobs").join(format!(
+                ".staging-{}-{}-{}",
+                layer.digest,
+                std::process::id(),
+                seq
+            ));
+            if std::fs::copy(&layer.path, &staging).is_err() {
+                let _ = std::fs::remove_file(&staging);
+                return;
+            }
+            if std::fs::rename(&staging, &blob_path).is_err() {
+                let _ = std::fs::remove_file(&staging);
+                return;
+            }
         }
 
         let record = KeyRecord {

src/runtime/src/oci/store.rs

@@ -6,6 +6,7 @@
 
 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};
@@ -14,6 +15,9 @@ use chrono::Utc;
 use serde::{Deserialize, Serialize};
 use tokio::sync::RwLock;
 
+/// Per-process counter for unique staging-dir names in `put`.
+static PUT_SEQ: AtomicU64 = AtomicU64::new(0);
+
 /// Persistent index stored as JSON on disk.
 #[derive(Debug, Default, Serialize, Deserialize)]
 struct StoreIndex {
@@ -130,11 +134,36 @@ impl ImageStore {
         let digest_hex = digest.strip_prefix("sha256:").unwrap_or(digest);
         let target_dir = self.store_dir.join("sha256").join(digest_hex);
 
-        // Copy source to target if not already present
+        // Copy source to target if not already present. Stage into a unique temp
+        // dir then atomically rename into place, so a concurrent put() for the
+        // same digest — or a copy that fails partway — can never leave a
+        // half-populated content-addressed dir that a later caller mistakes for a
+        // complete image (the bare check-then-copy raced on both counts).
         if !target_dir.exists() {
-            copy_dir_recursive(source_dir, &target_dir).map_err(|e| {
+            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);
+                // A concurrent put() may have populated target_dir first (rename
+                // onto a non-empty dir fails); that's fine — only propagate if the
+                // image still isn't there.
+                if !target_dir.exists() {
+                    return Err(BoxError::OciImageError(format!(
+                        "Failed to publish image to store: {}",
+                        e
+                    )));
+                }
+            }
         }
 
         let size_bytes = dir_size(&target_dir);