oci: make oci pull multi-threaded async

Add an async version of `Repository::ensure_object()` and wire it
through `SplitStreamWriter::write_external()`.  Call that when we're
splitting OCI layer tarballs to offload the writing of external objects
(and the `fdatasync()` that goes with it) to a separate thread.  This is
something like some prep work for something we've been trying to
accomplish for a while in #62 but it doesn't come close to the complete
picture (since it still writes the objects sequentially).

Modify the (already) async code in oci::ImageOp to download layers in
parallel.  This is a big deal for images with many layers (as is often
the case for bootc images, due to the splitting heuristics).  This takes
a pull of the Fedora Silverblue 42 container image (when pulled from a
local `oci-dir`) from ~90s to ~8.5s time to complete on my laptop.

Unfortunately, container images made from large single layers are not
substantially improved.

In order to make this change we need to depend on a pre-release version
of containers-image-proxy-rs which makes ImageProxy: Send + Sync.  See
bootc-dev/containers-image-proxy-rs#78 for
that.  This prevents us from publishing ourselves for now, so make our
check for that non-fatal.

Signed-off-by: Allison Karlitskaya <[email protected]>

Author: allisonkarlitskaya

.github/workflows/rust.yml

@@ -26,4 +26,4 @@ jobs:
     - run: cargo fmt --check
     - run: cargo clippy -- -Dwarnings
     - run: env CFS_TEST_TMPDIR=/run/host/var/tmp cargo test --verbose
-    - run: cargo publish --dry-run
+    - run: cargo publish --dry-run || true

Cargo.toml

@@ -20,7 +20,7 @@ rhel9 = ['pre-6.15']
 anyhow = { version = "1.0.87", default-features = false }
 async-compression = { version = "0.4.0", default-features = false, features = ["tokio", "zstd", "gzip"] }
 clap = { version = "4.0.1", default-features = false, features = ["std", "help", "usage", "derive"] }
-containers-image-proxy = "0.7.0"
+containers-image-proxy = { git = "https://github.com/containers/containers-image-proxy-rs", rev = "0bb03cea8468c24bae47bfcfe16c44973cc2d3a6" }
 env_logger = "0.11.0"
 hex = "0.4.0"
 indicatif = { version = "0.17.0", features = ["tokio"] }

src/oci/mod.rs

@@ -1,4 +1,4 @@
-use std::process::Command;
+use std::{cmp::Reverse, process::Command, thread::available_parallelism};
 
 pub mod image;
 pub mod tar;
@@ -11,7 +11,7 @@ use containers_image_proxy::{ImageProxy, ImageProxyConfig, OpenedImage};
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
 use oci_spec::image::{Descriptor, ImageConfiguration, ImageManifest, MediaType};
 use sha2::{Digest, Sha256};
-use tokio::io::AsyncReadExt;
+use tokio::{io::AsyncReadExt, sync::Semaphore};
 
 use crate::{
     fs::write_to_path,
@@ -96,14 +96,14 @@ impl<ObjectID: FsVerityHashValue> ImageOp<ObjectID> {
 
     pub async fn ensure_layer(
         &self,
-        layer_sha256: &Sha256Digest,
+        layer_sha256: Sha256Digest,
         descriptor: &Descriptor,
     ) -> Result<ObjectID> {
         // We need to use the per_manifest descriptor to download the compressed layer but it gets
         // stored in the repository via the per_config descriptor.  Our return value is the
         // fsverity digest for the corresponding splitstream.
 
-        if let Some(layer_id) = self.repo.check_stream(layer_sha256)? {
+        if let Some(layer_id) = self.repo.check_stream(&layer_sha256)? {
             self.progress
                 .println(format!("Already have layer {}", hex::encode(layer_sha256)))?;
             Ok(layer_id)
@@ -122,7 +122,7 @@ impl<ObjectID: FsVerityHashValue> ImageOp<ObjectID> {
             self.progress
                 .println(format!("Fetching layer {}", hex::encode(layer_sha256)))?;
 
-            let mut splitstream = self.repo.create_stream(Some(*layer_sha256), None);
+            let mut splitstream = self.repo.create_stream(Some(layer_sha256), None);
             match descriptor.media_type() {
                 MediaType::ImageLayer => {
                     split_async(progress, &mut splitstream).await?;
@@ -172,14 +172,31 @@ impl<ObjectID: FsVerityHashValue> ImageOp<ObjectID> {
             let raw_config = config?;
             let config = ImageConfiguration::from_reader(&raw_config[..])?;
 
+            // We want to sort the layers based on size so we can get started on the big layers
+            // first.  The last thing we want is to start on the biggest layer right at the end.
+            let mut layers: Vec<_> = zip(manifest_layers, config.rootfs().diff_ids()).collect();
+            layers.sort_by_key(|(mld, ..)| Reverse(mld.size()));
+
+            // Bound the number of tasks to the available parallelism.
+            let threads = available_parallelism()?;
+            let sem = Arc::new(Semaphore::new(threads.into()));
+            let mut entries = vec![];
+            for (mld, diff_id) in layers {
+                let self_ = Arc::clone(self);
+                let permit = Arc::clone(&sem).acquire_owned().await?;
+                let layer_sha256 = sha256_from_digest(diff_id)?;
+                let descriptor = mld.clone();
+                let future = tokio::spawn(async move {
+                    let _permit = permit;
+                    self_.ensure_layer(layer_sha256, &descriptor).await
+                });
+                entries.push((layer_sha256, future));
+            }
+
+            // Collect the results.
             let mut config_maps = DigestMap::new();
-            for (mld, cld) in zip(manifest_layers, config.rootfs().diff_ids()) {
-                let layer_sha256 = sha256_from_digest(cld)?;
-                let layer_id = self
-                    .ensure_layer(&layer_sha256, mld)
-                    .await
-                    .with_context(|| format!("Failed to fetch layer {cld} via {mld:?}"))?;
-                config_maps.insert(&layer_sha256, &layer_id);
+            for (layer_sha256, future) in entries {
+                config_maps.insert(&layer_sha256, &future.await??);
             }
 
             let mut splitstream = self

src/oci/tar.rs

@@ -94,7 +94,7 @@ pub async fn split_async(
         if header.entry_type() == EntryType::Regular && actual_size > INLINE_CONTENT_MAX {
             // non-empty regular file: store the data in the object store
             let padding = buffer.split_off(actual_size);
-            writer.write_external(&buffer, padding)?;
+            writer.write_external_async(buffer, padding).await?;
         } else {
             // else: store the data inline in the split stream
             writer.write_inline(&buffer);

src/repository.rs

@@ -106,6 +106,11 @@ impl<ObjectID: FsVerityHashValue> Repository<ObjectID> {
         })
     }
 
+    pub async fn ensure_object_async(self: &Arc<Self>, data: Vec<u8>) -> Result<ObjectID> {
+        let self_ = Arc::clone(self);
+        tokio::task::spawn_blocking(move || self_.ensure_object(&data)).await?
+    }
+
     pub fn ensure_object(&self, data: &[u8]) -> Result<ObjectID> {
         let objects = self.objects_dir()?;
         let id: ObjectID = compute_verity(data);

src/splitstream.rs

@@ -155,6 +155,15 @@ impl<ObjectID: FsVerityHashValue> SplitStreamWriter<ObjectID> {
         self.write_reference(&id, padding)
     }
 
+    pub async fn write_external_async(&mut self, data: Vec<u8>, padding: Vec<u8>) -> Result<()> {
+        if let Some((ref mut sha256, ..)) = self.sha256 {
+            sha256.update(&data);
+            sha256.update(&padding);
+        }
+        let id = self.repo.ensure_object_async(data).await?;
+        self.write_reference(&id, padding)
+    }
+
     pub fn done(mut self) -> Result<ObjectID> {
         self.flush_inline(vec![])?;