Fixes

orf · Jul 27, 2024 · 9642158 · 9642158
1 parent feea18e
commit 9642158
Show file tree

Hide file tree

Showing 13 changed files with 179 additions and 143 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Taskfile.yaml b/Taskfile.yaml
@@ -13,17 +13,25 @@ tasks:
       - rm -rf oci/{{.image | replace ":" "/" }}/image_new/
       - ./target/release/docker-repack oci/{{.image | replace ":" "/" }}/ --target-size=500MB --exclude='*.pyc' --split-file-threshold=25MB
       - skopeo copy --override-arch=arm64 --override-os=linux oci:oci/{{.image | replace ":" "/" }}/image_new/ docker-daemon:foo:abc
-#      - docker run -it --entrypoint=bash foo:abc
+  #      - docker run -it --entrypoint=bash foo:abc
 
-  run-and-load:
+  run-and-export:
     vars:
-        image: 'nvidia/cuda:12.5.1-cudnn-devel-ubuntu20.04'
-#        image: 'python:3.11'
+      #        image: 'nvidia/cuda:12.5.1-cudnn-devel-ubuntu20.04'
+      image: 'python:3.11'
+      compression_level: 9
+      image_path: 'oci/{{.image | replace ":" "/" }}'
+      image_slug: '{{.image | replace ":" "-" | replace "/" "-" }}'
+      tmp_dir:
+        sh: mktemp -d
     cmds:
       - cargo build --profile=release
       - rm -rf oci/{{.image | replace ":" "/" }}/image_new/
-      - ./target/release/docker-repack oci/{{.image | replace ":" "/" }}/ repack {{.CLI_ARGS}}
-      - skopeo copy --override-arch=arm64 --override-os=linux oci:oci/{{.image | replace ":" "/" }}/image_new/ docker-daemon:foo:abc
+      - ./target/release/docker-repack oci/{{.image | replace ":" "/" }}/ repack --compression={{.compression_level}} {{.CLI_ARGS}}
+      - skopeo copy --override-arch=arm64 --override-os=linux oci:{{.image_path }}/image_new/ docker://orfal/split:repacked-{{.image_slug}}
+      - skopeo copy --override-arch=arm64 --override-os=linux oci:{{.image_path }}/image/ docker://orfal/split:original-{{.image_slug}}
+      - skopeo copy --override-arch=arm64 --override-os=linux oci:{{.image_path }}/image/ dir://{{.tmp_dir}} --dest-compress --dest-compress-format=zstd --dest-compress-level={{.compression_level}}
+      - skopeo copy --override-arch=arm64 --override-os=linux dir://{{.tmp_dir}}  docker://orfal/split:zstd-{{.image_slug}}
 
   test-all:
     cmds:
@@ -42,7 +50,7 @@ tasks:
 
   sync:
     cmds:
-      - for: {var: IMAGES}
+      - for: { var: IMAGES }
         cmd: |
           mkdir -p oci/{{.ITEM | replace ":" "/" }}/image/
           skopeo copy --override-arch=arm64 --override-os=linux docker://docker.io/{{.ITEM}} oci:oci/{{.ITEM | replace ":" "/" }}/image/
diff --git a/docker-repack/Cargo.toml b/docker-repack/Cargo.toml
@@ -26,3 +26,4 @@ strum_macros = "0.26.4"
 globset = { version = "0.4.14", features = ["serde"] }
 file-mode = "0.1.2"
 comfy-table = { version = "7.1.1", default-features = false }
+clap-num = "1.1.1"
diff --git a/docker-repack/src/file_combiner.rs b/docker-repack/src/file_combiner.rs
@@ -1,7 +1,7 @@
 use crate::image_parser::{TarItem, TarItemChunk};
 use std::fmt::Write;
 
-const SCRIPT: &'static str = include_str!("./combine_files.sh");
+const SCRIPT: &str = include_str!("./combine_files.sh");
 
 pub fn generate_combining_script(
     chunked_files: &Vec<(&TarItem, Vec<TarItemChunk>)>,

diff --git a/docker-repack/src/image_parser/compression.rs b/docker-repack/src/image_parser/compression.rs
diff --git a/docker-repack/src/image_parser/image_reader.rs b/docker-repack/src/image_parser/image_reader.rs
@@ -1,11 +1,18 @@
 use crate::image_parser::layer_reader::Layer;
 use oci_spec::image::{ImageConfiguration, ImageIndex, ImageManifest};
+use std::fmt::{Display, Formatter};
 use std::fs::File;
 use std::path::{Path, PathBuf};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
 pub struct SourceLayerID(pub usize);
 
+impl Display for SourceLayerID {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Layer {:<2}", self.0)
+    }
+}
+
 pub struct ImageReader {
     layers: Vec<Layer>,
     pub config: ImageConfiguration,

diff --git a/docker-repack/src/image_parser/image_writer.rs b/docker-repack/src/image_parser/image_writer.rs
@@ -4,32 +4,40 @@ use crate::image_parser::layer_writer::{LayerType, LayerWriter, WrittenLayer};
 use crate::image_parser::{utils, HashAndSize, HashedWriter, ImageReader};
 use anyhow::bail;
 use chrono::Utc;
-use indicatif::{MultiProgress, ProgressDrawTarget};
+use indicatif::MultiProgress;
 use itertools::Itertools;
 use oci_spec::image::{
     Descriptor, HistoryBuilder, ImageIndexBuilder, ImageManifestBuilder, MediaType,
 };
 use rayon::prelude::*;
 use serde_json::json;
 use std::collections::HashMap;
+use std::fmt::{Display, Formatter};
 use std::fs::File;
 use std::io::{BufReader, BufWriter, Read};
 use std::ops::Range;
 use std::path::{Path, PathBuf};
-use byte_unit::{Byte, UnitType};
 use tar::Entry;
-
-const ZSTD_OUTPUT_LEVEL: i32 = 19;
+use zstd::zstd_safe::CompressionLevel;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct NewLayerID(usize);
 
+impl Display for NewLayerID {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "New layer {:<2}", self.0)
+    }
+}
+
+pub type PathKey<'a> = (SourceLayerID, &'a str, Range<u64>);
+pub type PathValue = (NewLayerID, Option<PathBuf>);
+
 pub struct ImageWriter<'a> {
     directory: PathBuf,
     blobs_dir: PathBuf,
     temp_dir: PathBuf,
     layers: Vec<LayerWriter>,
-    paths: HashMap<(SourceLayerID, &'a str, Range<u64>), (NewLayerID, Option<PathBuf>)>,
+    paths: HashMap<PathKey<'a>, PathValue>,
 }
 
 impl<'a> ImageWriter<'a> {
@@ -58,15 +66,15 @@ impl<'a> ImageWriter<'a> {
     ) -> anyhow::Result<(NewLayerID, &LayerWriter)> {
         let layer_id = NewLayerID(self.layers.len());
         let path = self.temp_dir.join(format!("{name}-{}.tar", layer_id.0));
-        let layer = LayerWriter::create_layer(path, type_)?;
+        let layer = LayerWriter::create_layer(layer_id, path, type_)?;
         self.layers.push(layer);
         Ok((layer_id, &self.layers[layer_id.0]))
     }
 
     pub fn add_layer_paths(
         &mut self,
         name: &'static str,
-        paths: impl Iterator<Item=(SourceLayerID, &'a str, Range<u64>, Option<PathBuf>)>,
+        paths: impl Iterator<Item = (SourceLayerID, &'a str, Range<u64>, Option<PathBuf>)>,
         type_: LayerType,
     ) -> anyhow::Result<()> {
         let (layer_id, _) = self.create_new_layer(name, type_)?;
@@ -182,34 +190,31 @@ impl<'a> ImageWriter<'a> {
     }
 
     pub fn finish_writing_layers(&mut self) -> anyhow::Result<Vec<WrittenLayer>> {
-        let finished_layers: Result<Vec<_>, _> = self
-            .layers
-            .drain(0..)
-            .into_iter()
-            .map(|layer| layer.finish())
-            .collect();
+        let finished_layers: Result<Vec<_>, _> =
+            self.layers.drain(0..).map(|layer| layer.finish()).collect();
 
-        Ok(finished_layers?)
+        finished_layers
     }
 
     pub fn compress_layers(
         &self,
         progress: &MultiProgress,
         finished_layers: Vec<WrittenLayer>,
+        compression_level: CompressionLevel,
     ) -> anyhow::Result<Vec<(WrittenLayer, HashAndSize)>> {
         let compressed_layers: Result<Vec<_>, _> = finished_layers
             .into_par_iter()
             .map(|layer| {
-                compress_layer(&progress, self.blobs_dir.clone(), &layer, ZSTD_OUTPUT_LEVEL)
+                compress_layer(progress, self.blobs_dir.clone(), &layer, compression_level)
                     .map(|v| (layer, v))
             })
             .collect();
-        Ok(compressed_layers?)
+        compressed_layers
     }
 
     pub fn write_index(
         self,
-        finished_layers: &Vec<(WrittenLayer, HashAndSize)>,
+        finished_layers: &[(WrittenLayer, HashAndSize)],
         mut image: ImageReader,
     ) -> anyhow::Result<()> {
         let root_fs = image.config.rootfs_mut();
@@ -292,7 +297,12 @@ fn compress_layer(
     encoder.set_pledged_src_size(Some(input_size))?;
 
     let buf_reader = BufReader::new(input_file);
-    let mut progress_reader = progress_reader(progress, input_size, buf_reader);
+    let mut progress_reader = progress_reader(
+        progress,
+        input_size,
+        buf_reader,
+        format!("Compressing {}", layer.id),
+    );
 
     std::io::copy(&mut progress_reader, &mut encoder)?;
 

diff --git a/docker-repack/src/image_parser/layer_reader.rs b/docker-repack/src/image_parser/layer_reader.rs
@@ -10,22 +10,25 @@ pub struct Layer {
     pub id: SourceLayerID,
     pub path: PathBuf,
     pub size: u64,
-    // pub digest: String,
 }
 
 pub fn progress_reader(
     progress: &MultiProgress,
     size: u64,
     file: BufReader<File>,
+    message: String,
 ) -> ProgressBarIter<BufReader<File>> {
     progress
         .add(
             ProgressBar::new(size)
                 .with_style(
-                    ProgressStyle::with_template("{wide_bar} {binary_bytes}/{binary_total_bytes}")
-                        .unwrap(),
+                    ProgressStyle::with_template(
+                        "{msg:>10} {wide_bar} {binary_bytes}/{binary_total_bytes}",
+                    )
+                    .unwrap(),
                 )
-                .with_finish(ProgressFinish::AndClear),
+                .with_finish(ProgressFinish::AndClear)
+                .with_message(message),
         )
         .wrap_read(file)
 }
@@ -39,7 +42,12 @@ impl Layer {
         let file = BufReader::new(file);
         let writer = match progress {
             None => ProgressBar::hidden().wrap_read(file),
-            Some(multi_progress) => progress_reader(multi_progress, self.size, file),
+            Some(multi_progress) => progress_reader(
+                multi_progress,
+                self.size,
+                file,
+                format!("Reading {}", self.id),
+            ),
         };
 
         let decoder = GzDecoder::new(writer);