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 ~20s time to complete.

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

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

Author: allisonkarlitskaya

Cargo.toml

@@ -20,8 +20,9 @@ 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"
+futures = "0.3.31"
 hex = "0.4.0"
 indicatif = { version = "0.17.0", features = ["tokio"] }
 log = "0.4.8"

src/oci/mod.rs

@@ -8,6 +8,7 @@ use std::{collections::HashMap, io::Read, iter::zip, path::Path, sync::Arc};
 use anyhow::{bail, ensure, Context, Result};
 use async_compression::tokio::bufread::{GzipDecoder, ZstdDecoder};
 use containers_image_proxy::{ImageProxy, ImageProxyConfig, OpenedImage};
+use futures::future::join_all;
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
 use oci_spec::image::{Descriptor, ImageConfiguration, ImageManifest, MediaType};
 use sha2::{Digest, Sha256};
@@ -96,14 +97,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 {layer_sha256:?}"))?;
             Ok(layer_id)
@@ -122,7 +123,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?;
@@ -173,13 +174,22 @@ impl<ObjectID: FsVerityHashValue> ImageOp<ObjectID> {
             let config = ImageConfiguration::from_reader(&raw_config[..])?;
 
             let mut config_maps = DigestMap::new();
+            let mut tasks = vec![];
+            let mut shas = vec![];
             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);
+                let mld = mld.clone();
+                let self_ = Arc::clone(self);
+                let future =
+                    tokio::spawn(async move { self_.ensure_layer(layer_sha256, &mld).await });
+                tasks.push(future);
+                shas.push(layer_sha256);
+            }
+
+            let layer_ids = join_all(tasks).await;
+
+            for (layer_sha256, layer_id) in zip(shas, layer_ids) {
+                config_maps.insert(&layer_sha256, &layer_id??);
             }
 
             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> {
         Self::ensure_object_in_dir(self.objects_dir()?, 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![])?;