Optimize implementation of image-rs/stream

image-rs/Cargo.toml

@@ -17,7 +17,6 @@ base64 = "0.13.0"
 cfg-if = { version = "1.0.0", optional = true }
 dircpy = { version = "0.3.12", optional = true }
 flate2 = "1.0"
-flume = "0.10.14"
 fs_extra = { version = "1.2.0", optional = true }
 futures = { version = "0.3.28", optional = true }
 futures-util = "0.3"

image-rs/src/digest.rs

@@ -0,0 +1,43 @@
+// Copyright (c) 2022 Intel Corporation
+//
+// SPDX-License-Identifier: Apache-2.0
+
+use sha2::Digest;
+
+pub const DIGEST_SHA256_PREFIX: &str = "sha256:";
+pub const DIGEST_SHA512_PREFIX: &str = "sha512:";
+
+pub trait DigestHasher {
+    fn digest_update(&mut self, buf: &[u8]);
+    fn digest_finalize(self) -> String;
+}
+
+#[derive(Clone, Debug)]
+pub enum LayerDigestHasher {
+    Sha256(sha2::Sha256),
+    Sha512(sha2::Sha512),
+}
+
+impl DigestHasher for LayerDigestHasher {
+    fn digest_update(&mut self, buf: &[u8]) {
+        match self {
+            LayerDigestHasher::Sha256(hasher) => {
+                hasher.update(buf);
+            }
+            LayerDigestHasher::Sha512(hasher) => {
+                hasher.update(buf);
+            }
+        }
+    }
+
+    fn digest_finalize(self) -> String {
+        match self {
+            LayerDigestHasher::Sha256(hasher) => {
+                format!("{}{:x}", DIGEST_SHA256_PREFIX, hasher.finalize())
+            }
+            LayerDigestHasher::Sha512(hasher) => {
+                format!("{}{:x}", DIGEST_SHA512_PREFIX, hasher.finalize())
+            }
+        }
+    }
+}

image-rs/src/lib.rs

@@ -10,6 +10,7 @@ pub mod bundle;
 pub mod config;
 pub mod decoder;
 pub mod decrypt;
+pub mod digest;
 pub mod image;
 pub mod meta_store;
 #[cfg(feature = "nydus")]

image-rs/src/pull.rs

@@ -16,14 +16,12 @@
 
 use crate::decoder::Compression;
 use crate::decrypt::Decryptor;
+use crate::digest::{DIGEST_SHA256_PREFIX, DIGEST_SHA512_PREFIX};
 use crate::image::LayerMeta;
 use crate::meta_store::MetaStore;
 use crate::stream::stream_processing;
 use crate::unpack::unpack;
 
-const DIGEST_SHA256: &str = "sha256";
-const DIGEST_SHA512: &str = "sha512";
-
 const ERR_NO_DECRYPT_CFG: &str = "decrypt_config is None";
 const ERR_BAD_UNCOMPRESSED_DIGEST: &str = "unsupported uncompressed digest format";
 const ERR_BAD_COMPRESSED_DIGEST: &str = "unsupported compressed digest format";
@@ -164,16 +162,16 @@
         layer_meta.decoder = Compression::try_from(media_type_str)?;
 
         if layer_meta.decoder == Compression::Uncompressed {
-            let digest = if diff_id.starts_with(DIGEST_SHA256) {
+            let digest = if diff_id.starts_with(DIGEST_SHA256_PREFIX) {
                 format!(
-                    "{}:{:x}",
-                    DIGEST_SHA256,
+                    "{}{:x}",
+                    DIGEST_SHA256_PREFIX,
                     sha2::Sha256::digest(plaintext_layer.as_slice())
                 )
-            } else if diff_id.starts_with(DIGEST_SHA512) {
+            } else if diff_id.starts_with(DIGEST_SHA512_PREFIX) {
                 format!(
-                    "{}:{:x}",
-                    DIGEST_SHA512,
+                    "{}{:x}",
+                    DIGEST_SHA512_PREFIX,
                     sha2::Sha512::digest(plaintext_layer.as_slice())
                 )
             } else {
@@ -188,12 +186,12 @@
                 .decoder
                 .decompress(plaintext_layer.as_slice(), &mut out)?;
 
-            if diff_id.starts_with(DIGEST_SHA256) {
+            if diff_id.starts_with(DIGEST_SHA256_PREFIX) {
                 layer_meta.uncompressed_digest =
-                    format!("{DIGEST_SHA256}:{:x}", sha2::Sha256::digest(&out));
-            } else if diff_id.starts_with(DIGEST_SHA512) {
+                    format!("{DIGEST_SHA256_PREFIX}{:x}", sha2::Sha256::digest(&out));
+            } else if diff_id.starts_with(DIGEST_SHA512_PREFIX) {
                 layer_meta.uncompressed_digest =
-                    format!("{DIGEST_SHA512}:{:x}", sha2::Sha512::digest(&out));
+                    format!("{DIGEST_SHA512_PREFIX}{:x}", sha2::Sha512::digest(&out));
             } else {
                 bail!("{}: {:?}", ERR_BAD_COMPRESSED_DIGEST, diff_id);
             }
@@ -380,7 +378,7 @@
 
         for image_url in oci_images.iter() {
             let tempdir = tempfile::tempdir().unwrap();
             let image = Reference::try_from(image_url.clone()).expect("create reference failed");
             let mut client = PullClient::new(
                 image,
                 tempdir.path(),
@@ -413,7 +411,7 @@
 
         for image_url in oci_images.iter() {
             let tempdir = tempfile::tempdir().unwrap();
             let image = Reference::try_from(image_url.clone()).expect("create reference failed");
             let mut client = PullClient::new(
                 image,
                 tempdir.path(),
@@ -457,7 +455,7 @@
 
         for image_url in oci_images.iter() {
             let tempdir = tempfile::tempdir().unwrap();
             let image = Reference::try_from(image_url.clone()).expect("create reference failed");
             let mut client = PullClient::new(
                 image,
                 tempdir.path(),
@@ -490,7 +488,7 @@
 
         for image_url in oci_images.iter() {
             let tempdir = tempfile::tempdir().unwrap();
             let image = Reference::try_from(image_url.clone()).expect("create reference failed");
             let mut client = PullClient::new(
                 image,
                 tempdir.path(),

image-rs/src/stream.rs

@@ -2,65 +2,30 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::unpack::unpack;
 use anyhow::{anyhow, bail, Context, Result};
 use sha2::Digest;
 use std::fs;
 use std::io::{Cursor, Read};
 use std::path::{Path, PathBuf};
+use std::sync::mpsc::{channel, Receiver};
 use tokio::io::{AsyncRead, AsyncReadExt};
 
-const CAPACITY: usize = 32768;
-const DIGEST_SHA256: &str = "sha256";
-const DIGEST_SHA512: &str = "sha512";
+use crate::digest::{DigestHasher, LayerDigestHasher, DIGEST_SHA256_PREFIX, DIGEST_SHA512_PREFIX};
+use crate::unpack::unpack;
 
+const CAPACITY: usize = 32768;
 const ERR_BAD_UNCOMPRESSED_DIGEST: &str = "unsupported uncompressed digest format";
 
-pub trait DigestHasher {
-    fn digest_update(&mut self, buf: &[u8]);
-    fn digest_finalize(self) -> String;
-}
-
-#[derive(Clone, Debug)]
-pub enum LayerDigestHasher {
-    Sha256(sha2::Sha256),
-    Sha512(sha2::Sha512),
-}
-
-impl DigestHasher for LayerDigestHasher {
-    fn digest_update(&mut self, buf: &[u8]) {
-        match self {
-            LayerDigestHasher::Sha256(hasher) => {
-                hasher.update(buf);
-            }
-            LayerDigestHasher::Sha512(hasher) => {
-                hasher.update(buf);
-            }
-        }
-    }
-
-    fn digest_finalize(self) -> String {
-        match self {
-            LayerDigestHasher::Sha256(hasher) => {
-                format!("{}:{:x}", DIGEST_SHA256, hasher.finalize())
-            }
-            LayerDigestHasher::Sha512(hasher) => {
-                format!("{}:{:x}", DIGEST_SHA512, hasher.finalize())
-            }
-        }
-    }
-}
-
-// Wrap a flume channel with [`Read`](std::io::Read) support.
+// Wrap a channel with [`Read`](std::io::Read) support.
 // This can bridge the [`AsyncRead`](tokio::io::AsyncRead) from
 // decrypt/decompress and impl Read for unpack.
 struct ChannelRead {
-    rx: flume::Receiver<Vec<u8>>,
+    rx: Receiver<Vec<u8>>,
     current: Cursor<Vec<u8>>,
 }
 
 impl ChannelRead {
-    fn new(rx: flume::Receiver<Vec<u8>>) -> ChannelRead {
+    fn new(rx: Receiver<Vec<u8>>) -> ChannelRead {
         ChannelRead {
             rx,
             current: Cursor::new(vec![]),
@@ -93,35 +58,30 @@ pub async fn stream_processing(
     destination: &Path,
 ) -> Result<String> {
     let dest = destination.to_path_buf();
-    let digest_str = if diff_id.starts_with(DIGEST_SHA256) {
-        let hasher = LayerDigestHasher::Sha256(sha2::Sha256::new());
-
-        channel_processing(layer_reader, hasher, dest)
-            .await
-            .map_err(|e| anyhow!("hasher {} : {:?}", DIGEST_SHA256, e))?
-    } else if diff_id.starts_with(DIGEST_SHA512) {
-        let hasher = LayerDigestHasher::Sha512(sha2::Sha512::new());
-
-        channel_processing(layer_reader, hasher, dest)
-            .await
-            .map_err(|e| anyhow!("hasher {} : {:?}", DIGEST_SHA512, e))?
+    let hasher = if diff_id.starts_with(DIGEST_SHA256_PREFIX) {
+        LayerDigestHasher::Sha256(sha2::Sha256::new())
+    } else if diff_id.starts_with(DIGEST_SHA512_PREFIX) {
+        LayerDigestHasher::Sha512(sha2::Sha512::new())
     } else {
         bail!("{}: {:?}", ERR_BAD_UNCOMPRESSED_DIGEST, diff_id);
     };
 
-    Ok(digest_str)
+    channel_processing(layer_reader, hasher, dest)
+        .await
+        .map_err(|e| anyhow!("hasher {} {:?}", DIGEST_SHA256_PREFIX, e))
 }
 
 async fn channel_processing(
     mut layer_reader: (impl AsyncRead + Unpin),
     mut hasher: LayerDigestHasher,
     destination: PathBuf,
 ) -> Result<String> {
-    let (tx, rx) = flume::unbounded();
+    let (tx, rx) = channel();
     let unpack_thread = std::thread::spawn(move || {
         let mut input = ChannelRead::new(rx);
 
         if let Err(e) = unpack(&mut input, destination.as_path()) {
+            // TODO
             fs::remove_dir_all(destination.as_path())
                 .context("Failed to roll back when unpacking")?;
             return Err(e);
@@ -130,8 +90,8 @@ async fn channel_processing(
         Result::<()>::Ok(())
     });
 
-    let mut buffer = [0; CAPACITY];
     loop {
+        let mut buffer = vec![0u8; CAPACITY];
         let n = layer_reader
             .read(&mut buffer)
             .await
@@ -140,9 +100,9 @@ async fn channel_processing(
             break;
         }
 
-        hasher.digest_update(&buffer[..n]);
-        tx.send_async(buffer[..n].to_vec())
-            .await
+        buffer.resize(n, 0);
+        hasher.digest_update(&buffer);
+        tx.send(buffer)
             .map_err(|e| anyhow!("channel: send failed {:?}", e))?;
     }
 
@@ -182,8 +142,8 @@ mod tests {
         let layer_data = ar.into_inner().unwrap();
 
         let layer_digest = format!(
-            "{}:{:x}",
-            DIGEST_SHA256,
+            "{}{:x}",
+            DIGEST_SHA256_PREFIX,
             sha2::Sha256::digest(layer_data.as_slice())
         );
 
@@ -222,8 +182,8 @@ mod tests {
         let layer_data = ar.into_inner().unwrap();
 
         let layer_digest = format!(
-            "{}:{:x}",
-            DIGEST_SHA256,
+            "{}{:x}",
+            DIGEST_SHA256_PREFIX,
             sha2::Sha256::digest(layer_data.as_slice())
         );
 
@@ -238,8 +198,8 @@ mod tests {
         let tempdir = tempfile::tempdir().unwrap();
         let file_path = tempdir.path().join("layer1");
         let layer_digest = format!(
-            "{}:{:x}",
-            DIGEST_SHA512,
+            "{}{:x}",
+            DIGEST_SHA512_PREFIX,
             sha2::Sha512::digest(layer_data.as_slice())
         );