Move counter to image client

image-rs/src/builder.rs

@@ -127,12 +127,15 @@ impl ClientBuilder {
 
         let meta_store = Arc::new(RwLock::new(meta_store));
 
+        let layer_index = ImageClient::get_layer_index(&self.config.work_dir.join("layers"));
+
         Ok(ImageClient {
             registry_auth,
             signature_validator,
             meta_store,
             snapshots,
             config: self.config,
+            layers_index: layer_index,
         })
     }
 }

image-rs/src/image.rs

@@ -10,7 +10,9 @@ use oci_client::Reference;
 use oci_spec::image::{ImageConfiguration, Os};
 use serde::{Deserialize, Serialize};
 use std::collections::{BTreeSet, HashMap};
+use std::fs;
 use std::path::{Path, PathBuf};
+use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 
 use tokio::sync::RwLock;
@@ -93,23 +95,16 @@ pub struct ImageClient {
 
     /// The config
     pub(crate) config: ImageConfig,
+
+    /// Next layer index
+    pub(crate) layers_index: AtomicUsize,
 }
 
 impl Default for ImageClient {
     // construct a default instance of `ImageClient`
     fn default() -> ImageClient {
         let work_dir = Path::new(DEFAULT_WORK_DIR);
-        let config = ImageConfig::try_from(work_dir.join(CONFIGURATION_FILE_NAME).as_path())
-            .unwrap_or_default();
-        let meta_store = MetaStore::try_from(work_dir.join(METAFILE).as_path()).unwrap_or_default();
-        let snapshots = Self::init_snapshots(&config.work_dir, &meta_store);
-        ImageClient {
-            registry_auth: None,
-            meta_store: Arc::new(RwLock::new(meta_store)),
-            snapshots,
-            signature_validator: None,
-            config,
-        }
+        ImageClient::new(work_dir.to_path_buf())
     }
 }
 
@@ -148,6 +143,32 @@ impl ImageClient {
         snapshots
     }
 
+    pub fn get_layer_index(data_dir: &Path) -> AtomicUsize {
+        let mut next = 0;
+        if data_dir.exists() {
+            match fs::read_dir(data_dir) {
+                Ok(paths) => {
+                    for path in paths {
+                        match path {
+                            Ok(entry_path) => match entry_path.file_name().to_str() {
+                                Some(name) => {
+                                    let n = name.parse::<usize>().unwrap_or(0);
+                                    if n >= next {
+                                        next = n + 1;
+                                    }
+                                }
+                                None => {}
+                            },
+                            Err(_) => {}
+                        }
+                    }
+                }
+                Err(_) => {}
+            }
+        }
+        AtomicUsize::new(next)
+    }
+
     /// Create an ImageClient instance with specific work directory.
     pub fn new(work_dir: PathBuf) -> Self {
         let config = ImageConfig::try_from(work_dir.join(CONFIGURATION_FILE_NAME).as_path())
@@ -161,6 +182,7 @@ impl ImageClient {
             registry_auth: None,
             signature_validator: None,
             config,
+            layers_index: Self::get_layer_index(&work_dir.join("layers")),
         }
     }
 
@@ -213,6 +235,7 @@ impl ImageClient {
             self.config.skip_proxy_ips.as_deref(),
             self.config.image_pull_proxy.as_deref(),
             self.config.extra_root_certificates.clone(),
+            &self.layers_index,
         )?;
         let (image_manifest, image_digest, image_config) = client.pull_manifest().await?;
 

image-rs/src/pull.rs

@@ -19,7 +19,6 @@ use crate::decrypt::Decryptor;
 use crate::image::LayerMeta;
 use crate::meta_store::MetaStore;
 use crate::stream::stream_processing;
-use std::fs;
 
 /// The PullClient connects to remote OCI registry, pulls the container image,
 /// and save the image layers under data_dir and return the layer meta info.
@@ -40,7 +39,7 @@ pub struct PullClient<'a> {
     pub max_concurrent_download: usize,
 
     /// Next layer index
-    pub layers_index: AtomicUsize,
+    pub layers_index: &'a AtomicUsize,
 }
 
 impl<'a> PullClient<'a> {
@@ -54,6 +53,7 @@ impl<'a> PullClient<'a> {
         no_proxy: Option<&str>,
         https_proxy: Option<&str>,
         extra_root_certificates: Vec<String>,
+        layers_index: &'a AtomicUsize,
     ) -> Result<PullClient<'a>> {
         let mut client_config = ClientConfig::default();
         if let Some(no_proxy) = no_proxy {
@@ -74,24 +74,6 @@ impl<'a> PullClient<'a> {
         client_config.extra_root_certificates.extend(certs);
         let client = Client::try_from(client_config)?;
 
-        let mut next = 0;
-        if data_dir.exists() {
-            let paths = fs::read_dir(data_dir)?;
-            for path in paths {
-                let entry_path = path?.path();
-                let name = entry_path
-                    .file_name()
-                    .unwrap_or_default()
-                    .to_str()
-                    .ok_or(anyhow!("Invalid unicode in path: {:?}", entry_path))?;
-                let n = name.parse::<usize>().unwrap_or(0);
-                if n >= next {
-                    next = n + 1;
-                }
-            }
-        }
-        let layers_index = AtomicUsize::new(next);
-
         Ok(PullClient {
             client,
             auth,
@@ -281,6 +263,7 @@ mod tests {
             None,
             None,
             vec![],
+            &AtomicUsize::new(0),
         )
         .unwrap();
         let (image_manifest, _image_digest, image_config) = client.pull_manifest().await.unwrap();
@@ -332,6 +315,7 @@ mod tests {
                 None,
                 None,
                 vec![],
+                &AtomicUsize::new(0),
             )
             .unwrap();
             let (image_manifest, _image_digest, image_config) =
@@ -371,6 +355,7 @@ mod tests {
                 None,
                 None,
                 vec![],
+                &AtomicUsize::new(0),
             )
             .unwrap();
             let (image_manifest, _image_digest, image_config) =
@@ -439,6 +424,7 @@ mod tests {
             None,
             None,
             vec![],
+            &AtomicUsize::new(0),
         )
         .unwrap();
 
@@ -531,6 +517,7 @@ mod tests {
                 None,
                 None,
                 vec![],
+                &AtomicUsize::new(0),
             )
             .unwrap();
             let (image_manifest, _image_digest, image_config) =