From cf8402824f9b832443a9e4c3a8fcf3086c0a0636 Mon Sep 17 00:00:00 2001 From: Silenio Quarti Date: Wed, 12 Feb 2025 11:47:55 -0500 Subject: [PATCH] Move counter to image client --- image-rs/src/builder.rs | 3 +++ image-rs/src/image.rs | 38 +++++++++++++++++++++++++++----------- image-rs/src/pull.rs | 32 ++++++++++++-------------------- 3 files changed, 42 insertions(+), 31 deletions(-) diff --git a/image-rs/src/builder.rs b/image-rs/src/builder.rs index 50b7ac25e..f1cd76b14 100644 --- a/image-rs/src/builder.rs +++ b/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, }) } } diff --git a/image-rs/src/image.rs b/image-rs/src/image.rs index f5a80e07a..7022144ef 100644 --- a/image-rs/src/image.rs +++ b/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,25 @@ impl ImageClient { snapshots } + pub fn get_layer_index(data_dir: &Path) -> AtomicUsize { + let mut next = 0; + if data_dir.exists() { + if let Ok(paths) = fs::read_dir(data_dir) { + for path in paths { + if let Ok(entry_path) = path { + if let Some(name) = entry_path.file_name().to_str() { + let n = name.parse::().unwrap_or(0); + if n >= next { + next = n + 1; + } + } + } + } + } + } + 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 +175,7 @@ impl ImageClient { registry_auth: None, signature_validator: None, config, + layers_index: Self::get_layer_index(&work_dir.join("layers")), } } @@ -213,6 +228,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?; diff --git a/image-rs/src/pull.rs b/image-rs/src/pull.rs index 2dedf9fb2..c7ed211b8 100644 --- a/image-rs/src/pull.rs +++ b/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, + layers_index: &'a AtomicUsize, ) -> Result> { 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::().unwrap_or(0); - if n >= next { - next = n + 1; - } - } - } - let layers_index = AtomicUsize::new(next); - Ok(PullClient { client, auth, @@ -273,6 +255,7 @@ mod tests { "nginx@sha256:9700d098d545f9d2ee0660dfb155fe64f4447720a0a763a93f2cf08997227279"; let tempdir = tempfile::tempdir().unwrap(); let image = Reference::try_from(image_url.to_string()).expect("create reference failed"); + let layers_index = AtomicUsize::new(0); let mut client = PullClient::new( image, tempdir.path(), @@ -281,6 +264,7 @@ mod tests { None, None, vec![], + &layers_index, ) .unwrap(); let (image_manifest, _image_digest, image_config) = client.pull_manifest().await.unwrap(); @@ -324,6 +308,7 @@ mod tests { let tempdir = tempfile::tempdir().unwrap(); let image = Reference::try_from(image_url.to_string()).expect("create reference failed"); + let layers_index = AtomicUsize::new(0); let mut client = PullClient::new( image, tempdir.path(), @@ -332,6 +317,7 @@ mod tests { None, None, vec![], + &layers_index, ) .unwrap(); let (image_manifest, _image_digest, image_config) = @@ -363,6 +349,7 @@ mod tests { let tempdir = tempfile::tempdir().unwrap(); let image = Reference::try_from(image_url.to_string()).expect("create reference failed"); + let layers_index = AtomicUsize::new(0); let mut client = PullClient::new( image, tempdir.path(), @@ -371,6 +358,7 @@ mod tests { None, None, vec![], + &layers_index, ) .unwrap(); let (image_manifest, _image_digest, image_config) = @@ -431,6 +419,7 @@ mod tests { }; let tempdir = tempfile::tempdir().unwrap(); + let layers_index = AtomicUsize::new(0); let mut client = PullClient::new( oci_image, tempdir.path(), @@ -439,6 +428,7 @@ mod tests { None, None, vec![], + &layers_index, ) .unwrap(); @@ -523,6 +513,7 @@ mod tests { for image_url in nydus_images.iter() { let tempdir = tempfile::tempdir().unwrap(); let image = Reference::try_from(*image_url).expect("create reference failed"); + let layers_index = AtomicUsize::new(0); let mut client = PullClient::new( image, tempdir.path(), @@ -531,6 +522,7 @@ mod tests { None, None, vec![], + &layers_index, ) .unwrap(); let (image_manifest, _image_digest, image_config) =