Merge pull request #9325 from ChengyuZhu6/image_service

agent:image: Refactor code to improve memory efficiency of image service

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"

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<Option<ImageService>> = Mutex::new(None);
+    pub static ref IMAGE_SERVICE: Arc<Mutex<Option<ImageService>>> = 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<Mutex<ImageClient>>,
+    image_client: ImageClient,
-    images: Arc<Mutex<HashMap<String, String>>>,
+    images: HashMap<String, String>,
 }
 
 impl ImageService {
     pub fn new() -> Self {
         Self {
-            image_client: Arc::new(Mutex::new(ImageClient::new(PathBuf::from(
+            image_client: ImageClient::new(PathBuf::from(KATA_IMAGE_WORK_DIR)),
-                KATA_IMAGE_WORK_DIR,
+            images: HashMap::new(),
-            )))),
-            images: Arc::new(Mutex::new(HashMap::new())),
         }
     }
 
-    /// Get the singleton instance of image service.
+    async fn add_image(&mut self, image: String, cid: String) {
-    pub async fn singleton() -> Result<ImageService> {
+        self.images.insert(image, cid);
-        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);
     }
 
     /// 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<String, String>,
@@ -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<String, String>,
+) -> Result<String> {
+    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;

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")]
-        {
+        image::merge_bundle_oci(&mut oci).await?;
-            let image_service = image::ImageService::singleton().await?;
-            image_service.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")]
-    {
+    image::init_image_service().await;
-        let image_service = image::ImageService::new();
+
-        *image::IMAGE_SERVICE.lock().await = Some(image_service.clone());
-    }
     let server = TtrpcServer::new()
         .bind(server_address)?
         .register_service(aservice)

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::pull_image(image_name, &cid, &image_pull_volume.metadata).await?;
-        let bundle_path = image_service
-            .pull_image(image_name, &cid, &image_pull_volume.metadata)
-            .await?;
 
         storage.source = bundle_path;
         storage.options = vec!["bind".to_string(), "ro".to_string()];