From f47408fdf4cc588d0b611faf5f42d459e5f49160 Mon Sep 17 00:00:00 2001 From: ChengyuZhu6 Date: Thu, 21 Mar 2024 11:51:39 +0800 Subject: [PATCH] agent:image: Refactor code to improve memory efficiency of image service Currently, `.lock().await.clone()` results in `Option` being duplicated in memory with each call to `singleton()`. Consequently, if kata-agent receives numerous image pulling requests simultaneously, it will lead to the allocation of multiple `Option` instances in memory, thereby consuming additional memory resources. In image.rs, we introduce two public functions: `merge_bundle_oci()` and `init_image_service()`. These functions will encapsulate the operations on `IMAGE_SERVICE`, ensuring that its internal details remain hidden from external modules such as `rpc.rs`. Fixes: #9225 -- part II Signed-off-by: Xynnn007 Signed-off-by: ChengyuZhu6 --- src/agent/Cargo.lock | 35 +++-- src/agent/src/image.rs | 145 +++++++++++--------- src/agent/src/rpc.rs | 11 +- src/agent/src/storage/image_pull_handler.rs | 5 +- 4 files changed, 109 insertions(+), 87 deletions(-) diff --git a/src/agent/Cargo.lock b/src/agent/Cargo.lock index fe02bd72b4..4c65509bd3 100644 --- a/src/agent/Cargo.lock +++ b/src/agent/Cargo.lock @@ -2282,9 +2282,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.151" +version = "0.2.153" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "302d7ab3130588088d277783b1e2d2e10c9e9e4a16dd9050e6ec93fb3e7048f4" +checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd" [[package]] name = "libloading" @@ -2293,7 +2293,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c2a198fb6b0eada2a8df47933734e6d35d350665a33a3593d7164fa52c75c19" dependencies = [ "cfg-if 1.0.0", - "windows-targets 0.48.0", + "windows-targets 0.52.4", ] [[package]] @@ -2805,9 +2805,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.15.0" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e82dad04139b71a90c080c8463fe0dc7902db5192d939bd0950f074d014339e1" +checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" [[package]] name = "opaque-debug" @@ -4640,9 +4640,9 @@ checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" [[package]] name = "sysinfo" -version = "0.29.11" +version = "0.30.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd727fc423c2060f6c92d9534cef765c65a6ed3f428a03d7def74a8c4348e666" +checksum = "0c385888ef380a852a16209afc8cfad22795dd8873d69c9a14d2e2088f118d18" dependencies = [ "cfg-if 1.0.0", "core-foundation-sys", @@ -4650,7 +4650,7 @@ dependencies = [ "ntapi", "once_cell", "rayon", - "winapi", + "windows", ] [[package]] @@ -5537,6 +5537,25 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e48a53791691ab099e5e2ad123536d0fff50652600abaf43bbf952894110d0be" +dependencies = [ + "windows-core", + "windows-targets 0.52.4", +] + +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ + "windows-targets 0.52.4", +] + [[package]] name = "windows-sys" version = "0.36.1" diff --git a/src/agent/src/image.rs b/src/agent/src/image.rs index 4ca37af70c..ce1f999bec 100644 --- a/src/agent/src/image.rs +++ b/src/agent/src/image.rs @@ -33,7 +33,7 @@ const K8S_CONTAINER_TYPE_KEYS: [&str; 2] = [ #[rustfmt::skip] lazy_static! { - pub static ref IMAGE_SERVICE: Mutex> = Mutex::new(None); + pub static ref IMAGE_SERVICE: Arc>> = Arc::new(Mutex::new(None)); } // Convenience function to obtain the scope logger. @@ -41,33 +41,21 @@ fn sl() -> slog::Logger { slog_scope::logger().new(o!("subsystem" => "image")) } -#[derive(Clone)] pub struct ImageService { - image_client: Arc>, - images: Arc>>, + image_client: ImageClient, + images: HashMap, } impl ImageService { pub fn new() -> Self { Self { - image_client: Arc::new(Mutex::new(ImageClient::new(PathBuf::from( - KATA_IMAGE_WORK_DIR, - )))), - images: Arc::new(Mutex::new(HashMap::new())), + image_client: ImageClient::new(PathBuf::from(KATA_IMAGE_WORK_DIR)), + images: HashMap::new(), } } - /// Get the singleton instance of image service. - pub async fn singleton() -> Result { - IMAGE_SERVICE - .lock() - .await - .clone() - .ok_or_else(|| anyhow!("image service is uninitialized")) - } - - async fn add_image(&self, image: String, cid: String) { - self.images.lock().await.insert(image, cid); + async fn add_image(&mut self, image: String, cid: String) { + self.images.insert(image, cid); } /// pause image is packaged in rootfs @@ -111,7 +99,7 @@ impl ImageService { /// # Returns /// - The image rootfs bundle path. (exp. /run/kata-containers/cb0b47276ea66ee9f44cc53afa94d7980b57a52c3f306f68cb034e58d9fbd3c6/images/rootfs) pub async fn pull_image( - &self, + &mut self, image: &str, cid: &str, image_metadata: &HashMap, @@ -145,8 +133,6 @@ impl ImageService { let res = self .image_client - .lock() - .await .pull_image(image, &bundle_path, &None, &None) .await; match res { @@ -170,51 +156,6 @@ impl ImageService { Ok(image_bundle_path.as_path().display().to_string()) } - /// When being passed an image name through a container annotation, merge its - /// corresponding bundle OCI specification into the passed container creation one. - pub async fn merge_bundle_oci(&self, container_oci: &mut oci::Spec) -> Result<()> { - if let Some(image_name) = container_oci.annotations.get(ANNO_K8S_IMAGE_NAME) { - let images = self.images.lock().await; - if let Some(container_id) = images.get(image_name) { - let image_oci_config_path = Path::new(CONTAINER_BASE) - .join(container_id) - .join(CONFIG_JSON); - debug!( - sl(), - "Image bundle config path: {:?}", image_oci_config_path - ); - - let image_oci = - oci::Spec::load(image_oci_config_path.to_str().ok_or_else(|| { - anyhow!( - "Invalid container image OCI config path {:?}", - image_oci_config_path - ) - })?) - .context("load image bundle")?; - - if let (Some(container_root), Some(image_root)) = - (container_oci.root.as_mut(), image_oci.root.as_ref()) - { - let root_path = Path::new(CONTAINER_BASE) - .join(container_id) - .join(image_root.path.clone()); - container_root.path = String::from(root_path.to_str().ok_or_else(|| { - anyhow!("Invalid container image root path {:?}", root_path) - })?); - } - - if let (Some(container_process), Some(image_process)) = - (container_oci.process.as_mut(), image_oci.process.as_ref()) - { - self.merge_oci_process(container_process, image_process); - } - } - } - - Ok(()) - } - /// Partially merge an OCI process specification into another one. fn merge_oci_process(&self, target: &mut oci::Process, source: &oci::Process) { // Override the target args only when the target args is empty and source.args is not empty @@ -261,12 +202,82 @@ pub async fn set_proxy_env_vars() { env::set_var("NO_PROXY", no_proxy); } } + match env::var("NO_PROXY") { Ok(val) => info!(sl(), "no_proxy is set to: {}", val), Err(e) => info!(sl(), "no_proxy is not set ({})", e), }; } +/// When being passed an image name through a container annotation, merge its +/// corresponding bundle OCI specification into the passed container creation one. +pub async fn merge_bundle_oci(container_oci: &mut oci::Spec) -> Result<()> { + let image_service = IMAGE_SERVICE.clone(); + let mut image_service = image_service.lock().await; + let image_service = image_service + .as_mut() + .expect("Image Service not initialized"); + if let Some(image_name) = container_oci.annotations.get(ANNO_K8S_IMAGE_NAME) { + if let Some(container_id) = image_service.images.get(image_name) { + let image_oci_config_path = Path::new(CONTAINER_BASE) + .join(container_id) + .join(CONFIG_JSON); + debug!( + sl(), + "Image bundle config path: {:?}", image_oci_config_path + ); + + let image_oci = oci::Spec::load(image_oci_config_path.to_str().ok_or_else(|| { + anyhow!( + "Invalid container image OCI config path {:?}", + image_oci_config_path + ) + })?) + .context("load image bundle")?; + + if let (Some(container_root), Some(image_root)) = + (container_oci.root.as_mut(), image_oci.root.as_ref()) + { + let root_path = Path::new(CONTAINER_BASE) + .join(container_id) + .join(image_root.path.clone()); + container_root.path = + String::from(root_path.to_str().ok_or_else(|| { + anyhow!("Invalid container image root path {:?}", root_path) + })?); + } + + if let (Some(container_process), Some(image_process)) = + (container_oci.process.as_mut(), image_oci.process.as_ref()) + { + image_service.merge_oci_process(container_process, image_process); + } + } + } + + Ok(()) +} + +/// Init the image service +pub async fn init_image_service() { + let image_service = ImageService::new(); + *IMAGE_SERVICE.lock().await = Some(image_service); +} + +pub async fn pull_image( + image: &str, + cid: &str, + image_metadata: &HashMap, +) -> Result { + let image_service = IMAGE_SERVICE.clone(); + let mut image_service = image_service.lock().await; + let image_service = image_service + .as_mut() + .expect("Image Service not initialized"); + + image_service.pull_image(image, cid, image_metadata).await +} + #[cfg(test)] mod tests { use super::ImageService; diff --git a/src/agent/src/rpc.rs b/src/agent/src/rpc.rs index 0cf1d45d86..38b6f161a9 100644 --- a/src/agent/src/rpc.rs +++ b/src/agent/src/rpc.rs @@ -205,10 +205,7 @@ impl AgentService { // In case of pulling image inside guest, we need to merge the image bundle OCI spec // into the container creation request OCI spec. #[cfg(feature = "guest-pull")] - { - let image_service = image::ImageService::singleton().await?; - image_service.merge_bundle_oci(&mut oci).await?; - } + image::merge_bundle_oci(&mut oci).await?; // Some devices need some extra processing (the ones invoked with // --device for instance), and that's what this call is doing. It @@ -1609,10 +1606,8 @@ pub async fn start( let hservice = health_ttrpc::create_health(Arc::new(health_service)); #[cfg(feature = "guest-pull")] - { - let image_service = image::ImageService::new(); - *image::IMAGE_SERVICE.lock().await = Some(image_service.clone()); - } + image::init_image_service().await; + let server = TtrpcServer::new() .bind(server_address)? .register_service(aservice) diff --git a/src/agent/src/storage/image_pull_handler.rs b/src/agent/src/storage/image_pull_handler.rs index e713198975..3cf546622c 100644 --- a/src/agent/src/storage/image_pull_handler.rs +++ b/src/agent/src/storage/image_pull_handler.rs @@ -49,10 +49,7 @@ impl StorageHandler for ImagePullHandler { .cid .clone() .ok_or_else(|| anyhow!("failed to get container id"))?; - let image_service = image::ImageService::singleton().await?; - let bundle_path = image_service - .pull_image(image_name, &cid, &image_pull_volume.metadata) - .await?; + let bundle_path = image::pull_image(image_name, &cid, &image_pull_volume.metadata).await?; storage.source = bundle_path; storage.options = vec!["bind".to_string(), "ro".to_string()];