oci: add MT xz support

Signed-off-by: Benny Zlotnik <bzlotnik@redhat.com>
Assisted-by: claude-opus-4.6

Author: bennyz

src/fls/decompress.rs

@@ -7,7 +7,7 @@ use tokio::io::AsyncReadExt;
 use tokio::process::{Child, Command};
 use tokio::sync::mpsc;
 
-fn mb_to_bytes(mb: u64) -> u64 {
+pub(crate) fn mb_to_bytes(mb: u64) -> u64 {
     mb.saturating_mul(1024 * 1024)
 }
 
@@ -25,6 +25,29 @@ pub(crate) fn create_xz_decoder<R: Read>(
     Ok(liblzma::read::XzDecoder::new_stream(reader, stream))
 }
 
+pub(crate) fn create_mt_xz_decoder<R: Read + Send + 'static>(
+    reader: R,
+    xz_memlimit_mb: u64,
+) -> Result<Box<dyn Read + Send>, String> {
+    let num_threads = std::thread::available_parallelism()
+        .map(|n| n.get() as u32)
+        .unwrap_or(2);
+    let memlimit = mb_to_bytes(xz_memlimit_mb);
+    eprintln!(
+        "XZ decompression: {} threads, memory limit {}MB",
+        num_threads, xz_memlimit_mb
+    );
+    let stream = liblzma::stream::MtStreamBuilder::new()
+        .threads(num_threads)
+        .memlimit_threading(memlimit)
+        .memlimit_stop(memlimit)
+        .decoder()
+        .map_err(|e| format!("Failed to create MT XZ decoder: {}", e))?;
+    Ok(Box::new(liblzma::read::XzDecoder::new_stream(
+        reader, stream,
+    )))
+}
+
 /// Determines the appropriate decompression command based on URL extension
 fn get_decompressor_command(url: &str) -> &'static str {
     let extension = url.rsplit('.').next().unwrap_or("").to_lowercase();
@@ -130,23 +153,7 @@ pub(crate) fn start_inprocess_decompressor(
                 ChannelReader::new_byte_bounded(buffer_rx).with_progress(consumed_progress_tx);
 
             let mut decoder: Box<dyn Read + Send> = match compression {
-                Compression::Xz => {
-                    let num_threads = std::thread::available_parallelism()
-                        .map(|n| n.get() as u32)
-                        .unwrap_or(2);
-                    let memlimit = mb_to_bytes(xz_memlimit_mb);
-                    eprintln!(
-                        "XZ decompression: {} threads, memory limit {}MB",
-                        num_threads, xz_memlimit_mb
-                    );
-                    let stream = liblzma::stream::MtStreamBuilder::new()
-                        .threads(num_threads)
-                        .memlimit_threading(memlimit)
-                        .memlimit_stop(memlimit)
-                        .decoder()
-                        .map_err(|e| format!("Failed to create MT XZ decoder: {}", e))?;
-                    Box::new(liblzma::read::XzDecoder::new_stream(channel_reader, stream))
-                }
+                Compression::Xz => create_mt_xz_decoder(channel_reader, xz_memlimit_mb)?,
                 Compression::Gzip => Box::new(flate2::read::GzDecoder::new(channel_reader)),
                 Compression::None => Box::new(channel_reader),
                 Compression::Zstd => {

src/fls/oci/from_oci.rs

@@ -56,7 +56,7 @@ struct DownloadContext {
 struct RawDiskDownloadParams {
     http_tx: ByteBoundedSender<bytes::Bytes>,
     writer_handle: tokio::task::JoinHandle<Result<u64, std::io::Error>>,
-    external_decompressor: Option<tokio::process::Child>,
+    is_compressed: bool,
     decompressed_progress_rx: mpsc::UnboundedReceiver<u64>,
     raw_written_progress_rx: mpsc::UnboundedReceiver<u64>,
 }
@@ -69,12 +69,6 @@ struct ProcessingHandles {
     tar_extractor_handle: tokio::task::JoinHandle<Result<(), String>>,
 }
 
-/// Components returned by external decompressor pipeline setup
-struct ExternalDecompressorPipeline {
-    writer_handle: tokio::task::JoinHandle<Result<u64, std::io::Error>>,
-    decompressor: tokio::process::Child,
-}
-
 /// Components returned by pipeline setup
 struct TarPipelineComponents {
     http_tx: ByteBoundedSender<bytes::Bytes>,
@@ -1423,159 +1417,36 @@ async fn coordinate_download_and_processing(
     Ok(())
 }
 
-/// Setup external decompressor pipeline for XZ compression
-async fn setup_external_decompressor_pipeline(
-    http_rx: ByteBoundedReceiver<bytes::Bytes>,
-    block_writer: AsyncBlockWriter,
-    decompressed_progress_tx: mpsc::UnboundedSender<u64>,
-    debug: bool,
-) -> Result<ExternalDecompressorPipeline, Box<dyn std::error::Error>> {
-    // XZ: Use external xzcat process
-    let (mut decompressor, decompressor_name) = start_decompressor_process("disk.img.xz").await?;
-
-    let decompressor_stdin = decompressor.stdin.take().unwrap();
-    let decompressor_stdout = decompressor.stdout.take().unwrap();
-    let decompressor_stderr = decompressor.stderr.take().unwrap();
-
-    let (error_tx, error_rx) = mpsc::unbounded_channel::<String>();
-
-    // Spawn stderr reader
-    tokio::spawn(spawn_stderr_reader(
-        decompressor_stderr,
-        error_tx.clone(),
-        decompressor_name,
-    ));
-
-    // Spawn error processor
-    tokio::spawn(process_error_messages(error_rx));
-
-    // Spawn blocking task: read from channel and write to xzcat stdin
-    // First, create a sync file handle from the async stdin
-    #[cfg(unix)]
-    let stdin_fd = {
-        use std::os::unix::io::{AsRawFd, FromRawFd};
-        let raw_fd = decompressor_stdin.as_raw_fd();
-        // Duplicate the fd so we can use it in blocking context
-        let dup_fd = unsafe { libc::dup(raw_fd) };
-        if dup_fd == -1 {
-            return Err(std::io::Error::last_os_error().into());
-        }
-        // SAFETY: dup_fd is a valid file descriptor (we checked above)
-        unsafe { std::fs::File::from_raw_fd(dup_fd) }
-    };
-    #[cfg(not(unix))]
-    let stdin_fd: std::fs::File = {
-        return Err("XZ streaming decompression is not supported on non-unix platforms".into());
-    };
-
-    // Drop the original async stdin (the dup'd fd still points to the pipe)
-    drop(decompressor_stdin);
-
-    let stdin_writer_handle = {
-        tokio::task::spawn_blocking(move || {
-            use std::io::Write as _;
-            let reader = ChannelReader::new_byte_bounded(http_rx);
-            let mut reader = reader;
-            let mut stdin = stdin_fd;
-            let mut buffer = vec![0u8; 1024 * 1024]; // 1MB chunks
-
-            loop {
-                match reader.read(&mut buffer) {
-                    Ok(0) => break, // EOF
-                    Ok(n) => {
-                        if let Err(e) = stdin.write_all(&buffer[..n]) {
-                            return Err(format!("Error writing to xzcat: {}", e));
-                        }
-                    }
-                    Err(e) => return Err(format!("Error reading stream: {}", e)),
-                }
-            }
-
-            drop(stdin);
-            Ok::<(), String>(())
-        })
-    };
-
-    // Spawn task: xzcat stdout -> block writer with sparse detection
-    let writer = block_writer;
-    let progress_tx = decompressed_progress_tx;
-    let writer_handle = tokio::spawn(async move {
-        let mut stdout = decompressor_stdout;
-        let mut buffer = vec![0u8; 8 * 1024 * 1024]; // 8MB buffer
-
-        // Auto-detect sparse image format from initial data
-        let mut detector = FormatDetector::new();
-        let mut parser: Option<SparseParser> = None;
-        let mut format_determined = false;
-
-        loop {
-            let n = match tokio::io::AsyncReadExt::read(&mut stdout, &mut buffer).await {
-                Ok(0) => {
-                    finalize_format_at_eof(&mut detector, format_determined, &writer, debug)
-                        .await?;
-                    break;
-                }
-                Ok(n) => n,
-                Err(e) => return Err(e),
-            };
-
-            let _ = progress_tx.send(n as u64);
-
-            process_buffer_with_format_detection(
-                &buffer,
-                n,
-                &mut detector,
-                &mut parser,
-                &mut format_determined,
-                &writer,
-                debug,
-            )
-            .await?;
-        }
-
-        // Wait for stdin writer to finish and propagate any errors
-        stdin_writer_handle
-            .await
-            .map_err(|e| std::io::Error::other(format!("Stdin writer task failed: {}", e)))?
-            .map_err(|e| std::io::Error::other(format!("Stdin writer error: {}", e)))?;
-
-        writer.close().await
-    });
-
-    Ok(ExternalDecompressorPipeline {
-        writer_handle,
-        decompressor,
-    })
-}
-
-/// Setup in-process decompression pipeline for Gzip or None compression
 async fn setup_inprocess_decompression_pipeline(
     http_rx: ByteBoundedReceiver<bytes::Bytes>,
     block_writer: AsyncBlockWriter,
     decompressed_progress_tx: mpsc::UnboundedSender<u64>,
     compression_type: Compression,
     debug: bool,
+    xz_memlimit_mb: u64,
 ) -> Result<tokio::task::JoinHandle<Result<u64, std::io::Error>>, Box<dyn std::error::Error>> {
-    // Gzip or None: decompress in-process and write directly to block writer
     let writer = block_writer;
     let progress_tx = decompressed_progress_tx;
 
-    // Create an async channel for decompressed data
     let (data_tx, mut data_rx) = mpsc::channel::<Vec<u8>>(16);
 
-    // Spawn blocking task: read, decompress, send to async channel
     let reader_handle = tokio::task::spawn_blocking(move || {
         let reader = ChannelReader::new_byte_bounded(http_rx);
 
-        // Apply in-process gzip decompression if needed
         let processed_reader: Box<dyn std::io::Read + Send> = match compression_type {
             Compression::Gzip => {
                 if debug {
                     eprintln!("[DEBUG] Applying in-process gzip decompression");
                 }
                 Box::new(flate2::read::GzDecoder::new(reader))
             }
-            _ => Box::new(reader),
+            Compression::Xz => {
+                crate::fls::decompress::create_mt_xz_decoder(reader, xz_memlimit_mb)?
+            }
+            Compression::None => Box::new(reader),
+            Compression::Zstd => {
+                return Err("Zstd in-process decompression is not supported".to_string());
+            }
         };
 
         let mut reader = processed_reader;
@@ -1657,7 +1528,7 @@ async fn coordinate_raw_disk_download(
     let mut progress =
         ProgressTracker::new(options.common.newline_progress, options.common.show_memory);
     progress.set_content_length(Some(layer_size));
-    progress.set_is_compressed(params.external_decompressor.is_some());
+    progress.set_is_compressed(params.is_compressed);
     progress.bytes_received = initial_buffer.len() as u64;
     let update_interval = Duration::from_secs_f64(options.common.progress_interval_secs);
 
@@ -1772,24 +1643,6 @@ async fn coordinate_raw_disk_download(
     progress.decompress_duration = Some(elapsed);
     progress.write_duration = Some(elapsed);
 
-    // Wait for external decompressor process and check exit status
-    if let Some(mut decompressor) = params.external_decompressor {
-        match decompressor.wait().await {
-            Ok(status) => {
-                if !status.success() {
-                    return Err(format!(
-                        "Decompressor process failed with exit code: {:?}",
-                        status.code()
-                    )
-                    .into());
-                }
-            }
-            Err(e) => {
-                return Err(format!("Failed to wait for decompressor process: {}", e).into());
-            }
-        }
-    }
-
     // Final progress update
     let _ = progress.update_progress(Some(layer_size), update_interval, true);
 
@@ -2520,48 +2373,31 @@ async fn flash_raw_disk_image_directly(
         byte_bounded_channel::<bytes::Bytes>(max_buffer_bytes, buffer_capacity);
     let (decompressed_progress_tx, decompressed_progress_rx) = mpsc::unbounded_channel::<u64>();
 
-    // For gzip and none, we can decompress in-process and write directly to block writer
-    // For XZ, we need the external xzcat process
-    let needs_external_decompressor = compression_type == Compression::Xz;
+    if compression_type == Compression::Zstd {
+        return Err("Zstd in-process decompression is not supported".into());
+    }
 
     if options.common.debug {
         eprintln!(
-            "[DEBUG] Compression type: {:?}, using external decompressor: {}",
-            compression_type, needs_external_decompressor
+            "[DEBUG] Compression type: {:?}, using in-process decompressor",
+            compression_type
         );
     }
 
-    // Spawn the processing pipeline based on compression type
-    let (writer_handle, external_decompressor) = if needs_external_decompressor {
-        // XZ: Use external xzcat process
-        let pipeline = setup_external_decompressor_pipeline(
-            http_rx,
-            block_writer,
-            decompressed_progress_tx,
-            options.common.debug,
-        )
-        .await?;
-
-        (pipeline.writer_handle, Some(pipeline.decompressor))
-    } else {
-        // Gzip or None: decompress in-process and write directly to block writer
-        let handle = setup_inprocess_decompression_pipeline(
-            http_rx,
-            block_writer,
-            decompressed_progress_tx,
-            compression_type,
-            options.common.debug,
-        )
-        .await?;
-
-        (handle, None)
-    };
+    let writer_handle = setup_inprocess_decompression_pipeline(
+        http_rx,
+        block_writer,
+        decompressed_progress_tx,
+        compression_type,
+        options.common.debug,
+        options.common.xz_memlimit_mb,
+    )
+    .await?;
 
-    // Coordinate download and processing
     let params = RawDiskDownloadParams {
         http_tx,
         writer_handle,
-        external_decompressor,
+        is_compressed: compression_type != Compression::None,
         decompressed_progress_rx,
         raw_written_progress_rx,
     };