Add progress handler for rootfs unpacking

Sources/Containerization/ContainerManager.swift

@@ -376,13 +376,15 @@ public struct ContainerManager: Sendable {
     ///   - readOnly: Whether to mount the root filesystem as read-only.
     ///   - networking: Whether to create a network interface for this container. Defaults to `true`.
     ///     When `false`, no network resources are allocated and `releaseNetwork`/`delete` remain safe to call.
+    ///   - progress: Optional handler for tracking rootfs unpacking progress.
     public mutating func create(
         _ id: String,
         reference: String,
         rootfsSizeInBytes: UInt64 = 8.gib(),
         writableLayerSizeInBytes: UInt64? = nil,
         readOnly: Bool = false,
         networking: Bool = true,
+        progress: ProgressHandler? = nil,
         configuration: (inout LinuxContainer.Configuration) throws -> Void
     ) async throws -> LinuxContainer {
         let image = try await imageStore.get(reference: reference, pull: true)
@@ -393,6 +395,7 @@ public struct ContainerManager: Sendable {
             writableLayerSizeInBytes: writableLayerSizeInBytes,
             readOnly: readOnly,
             networking: networking,
+            progress: progress,
             configuration: configuration
         )
     }
@@ -407,21 +410,24 @@ public struct ContainerManager: Sendable {
     ///   - readOnly: Whether to mount the root filesystem as read-only.
     ///   - networking: Whether to create a network interface for this container. Defaults to `true`.
     ///     When `false`, no network resources are allocated and `releaseNetwork`/`delete` remain safe to call.
+    ///   - progress: Optional handler for tracking rootfs unpacking progress.
     public mutating func create(
         _ id: String,
         image: Image,
         rootfsSizeInBytes: UInt64 = 8.gib(),
         writableLayerSizeInBytes: UInt64? = nil,
         readOnly: Bool = false,
         networking: Bool = true,
+        progress: ProgressHandler? = nil,
         configuration: (inout LinuxContainer.Configuration) throws -> Void
     ) async throws -> LinuxContainer {
         let path = try createContainerRoot(id)
 
         var rootfs = try await unpack(
             image: image,
             destination: path.appendingPathComponent("rootfs.ext4"),
-            size: rootfsSizeInBytes
+            size: rootfsSizeInBytes,
+            progress: progress
         )
         if readOnly {
             rootfs.options.append("ro")
@@ -511,10 +517,10 @@ public struct ContainerManager: Sendable {
         return path
     }
 
-    private func unpack(image: Image, destination: URL, size: UInt64) async throws -> Mount {
+    private func unpack(image: Image, destination: URL, size: UInt64, progress: ProgressHandler? = nil) async throws -> Mount {
         do {
             let unpacker = EXT4Unpacker(blockSizeInBytes: size)
-            return try await unpacker.unpack(image, for: .current, at: destination)
+            return try await unpacker.unpack(image, for: .current, at: destination, progress: progress)
         } catch let err as ContainerizationError {
             if err.code == .exists {
                 return .block(

Sources/Containerization/Image/Unpacker/EXT4Unpacker.swift

@@ -42,19 +42,18 @@ public struct EXT4Unpacker: Unpacker {
         archive: URL,
         compression: ContainerizationArchive.Filter,
         at path: URL
-    ) throws {
+    ) async throws {
         let cleanedPath = try prepareUnpackPath(path: path)
         let filesystem = try EXT4.Formatter(
             FilePath(cleanedPath),
             minDiskSize: blockSizeInBytes
         )
         defer { try? filesystem.close() }
 
-        try filesystem.unpack(
+        try await filesystem.unpack(
             source: archive,
             format: .paxRestricted,
-            compression: compression,
-            progress: nil
+            compression: compression
         )
     }
     #endif
@@ -84,27 +83,59 @@ public struct EXT4Unpacker: Unpacker {
         )
         defer { try? filesystem.close() }
 
+        // Resolve layer paths upfront. When progress reporting is enabled and a layer
+        // uses zstd, decompress once so both the size-scanning pass and the unpack
+        // pass share the same decompressed file.
+        var resolvedLayers: [(file: URL, filter: ContainerizationArchive.Filter)] = []
+        var decompressedFiles: [URL] = []
         for layer in manifest.layers {
             try Task.checkCancellation()
             let content = try await image.getContent(digest: layer.digest)
+            let compression = try compressionFilter(for: layer.mediaType)
+            if progress != nil && compression == .zstd {
+                let decompressed = try ArchiveReader.decompressZstd(content.path)
+                decompressedFiles.append(decompressed)
+                resolvedLayers.append((file: decompressed, filter: .none))
+            } else {
+                resolvedLayers.append((file: content.path, filter: compression))
+            }
+        }
+        defer {
+            for file in decompressedFiles {
+                ArchiveReader.cleanUpDecompressedZstd(file)
+            }
+        }
 
-            let compression: ContainerizationArchive.Filter
-            switch layer.mediaType {
-            case MediaTypes.imageLayer, MediaTypes.dockerImageLayer:
-                compression = .none
-            case MediaTypes.imageLayerGzip, MediaTypes.dockerImageLayerGzip:
-                compression = .gzip
-            case MediaTypes.imageLayerZstd, MediaTypes.dockerImageLayerZstd:
-                compression = .zstd
-            default:
-                throw ContainerizationError(.unsupported, message: "media type \(layer.mediaType) not supported.")
+        if let progress {
+            var totalSize: Int64 = 0
+            var totalItems: Int = 0
+            for layer in resolvedLayers {
+                try Task.checkCancellation()
+                let totals = try EXT4.Formatter.scanArchiveHeaders(
+                    format: .paxRestricted, filter: layer.filter, file: layer.file)
+                totalSize += totals.size
+                totalItems += totals.items
             }
-            try filesystem.unpack(
-                source: content.path,
+            var totalEvents: [ProgressEvent] = []
+            if totalSize > 0 {
+                totalEvents.append(ProgressEvent(event: "add-total-size", value: totalSize))
+            }
+            if totalItems > 0 {
+                totalEvents.append(ProgressEvent(event: "add-total-items", value: totalItems))
+            }
+            if !totalEvents.isEmpty {
+                await progress(totalEvents)
+            }
+        }
+
+        for resolved in resolvedLayers {
+            try Task.checkCancellation()
+            let reader = try ArchiveReader(
                 format: .paxRestricted,
-                compression: compression,
-                progress: progress
+                filter: resolved.filter,
+                file: resolved.file
             )
+            try await filesystem.unpack(reader: reader, progress: progress)
         }
 
         return .block(
@@ -123,4 +154,19 @@ public struct EXT4Unpacker: Unpacker {
         }
         return blockPath
     }
+
+    #if os(macOS)
+    private func compressionFilter(for mediaType: String) throws -> ContainerizationArchive.Filter {
+        switch mediaType {
+        case MediaTypes.imageLayer, MediaTypes.dockerImageLayer:
+            return .none
+        case MediaTypes.imageLayerGzip, MediaTypes.dockerImageLayerGzip:
+            return .gzip
+        case MediaTypes.imageLayerZstd, MediaTypes.dockerImageLayerZstd:
+            return .zstd
+        default:
+            throw ContainerizationError(.unsupported, message: "media type \(mediaType) not supported.")
+        }
+    }
+    #endif
 }

Sources/ContainerizationArchive/ArchiveReader.swift

@@ -125,7 +125,7 @@ public final class ArchiveReader {
     }
 
     /// Decompress a zstd file to a temporary location
-    private static func decompressZstd(_ source: URL) throws -> URL {
+    public static func decompressZstd(_ source: URL) throws -> URL {
         guard let tempDir = createTemporaryDirectory(baseName: "zstd-decompress") else {
             throw ArchiveError.failedToDetectFormat
         }
@@ -148,13 +148,19 @@ public final class ArchiveReader {
         return tempFile
     }
 
+    /// Clean up the temporary directory created by `decompressZstd`.
+    /// The decompressed file is placed inside a unique temporary directory,
+    /// so removing that directory cleans up everything.
+    public static func cleanUpDecompressedZstd(_ file: URL) {
+        try? FileManager.default.removeItem(at: file.deletingLastPathComponent())
+    }
+
     deinit {
         archive_read_free(underlying)
         try? fileHandle?.close()
 
-        // Clean up temp decompressed file
         if let tempFile = tempDecompressedFile {
-            try? FileManager.default.removeItem(at: tempFile.deletingLastPathComponent())
+            Self.cleanUpDecompressedZstd(tempFile)
         }
     }
 }

Sources/ContainerizationEXT4/Formatter+Unpack.swift

@@ -25,7 +25,83 @@ private typealias Hardlinks = [FilePath: FilePath]
 
 extension EXT4.Formatter {
     /// Unpack the provided archive on to the ext4 filesystem.
-    public func unpack(reader: ArchiveReader, progress: ProgressHandler? = nil) throws {
+    public func unpack(reader: ArchiveReader, progress: ProgressHandler? = nil) async throws {
+        try await self.unpackEntries(reader: reader, progress: progress)
+    }
+
+    /// Unpack an archive at the source URL on to the ext4 filesystem.
+    public func unpack(
+        source: URL,
+        format: ContainerizationArchive.Format = .paxRestricted,
+        compression: ContainerizationArchive.Filter = .gzip,
+        progress: ProgressHandler? = nil
+    ) async throws {
+        // For zstd, decompress once and reuse for both passes to avoid double decompression.
+        let fileToRead: URL
+        let readerFilter: ContainerizationArchive.Filter
+        var decompressedFile: URL?
+        if progress != nil && compression == .zstd {
+            let decompressed = try ArchiveReader.decompressZstd(source)
+            fileToRead = decompressed
+            readerFilter = .none
+            decompressedFile = decompressed
+        } else {
+            fileToRead = source
+            readerFilter = compression
+        }
+        defer {
+            if let decompressedFile {
+                ArchiveReader.cleanUpDecompressedZstd(decompressedFile)
+            }
+        }
+
+        if let progress {
+            // First pass: scan headers to get totals (fast, metadata only)
+            let totals = try Self.scanArchiveHeaders(format: format, filter: readerFilter, file: fileToRead)
+            var totalEvents: [ProgressEvent] = []
+            if totals.size > 0 {
+                totalEvents.append(ProgressEvent(event: "add-total-size", value: totals.size))
+            }
+            if totals.items > 0 {
+                totalEvents.append(ProgressEvent(event: "add-total-items", value: totals.items))
+            }
+            if !totalEvents.isEmpty {
+                await progress(totalEvents)
+            }
+        }
+
+        // Unpack pass
+        let reader = try ArchiveReader(
+            format: format,
+            filter: readerFilter,
+            file: fileToRead
+        )
+        try await self.unpackEntries(reader: reader, progress: progress)
+    }
+
+    /// Scan archive headers to count the total number of bytes in regular files
+    /// and the total number of entries.
+    public static func scanArchiveHeaders(
+        format: ContainerizationArchive.Format,
+        filter: ContainerizationArchive.Filter,
+        file: URL
+    ) throws -> (size: Int64, items: Int) {
+        let reader = try ArchiveReader(format: format, filter: filter, file: file)
+        var totalSize: Int64 = 0
+        var totalItems: Int = 0
+        for (entry, _) in reader.makeStreamingIterator() {
+            try Task.checkCancellation()
+            guard entry.path != nil else { continue }
+            totalItems += 1
+            if entry.fileType == .regular, let size = entry.size {
+                totalSize += Int64(size)
+            }
+        }
+        return (size: totalSize, items: totalItems)
+    }
+
+    /// Core unpack logic. When `progress` is nil the handler calls are skipped.
+    private func unpackEntries(reader: ArchiveReader, progress: ProgressHandler?) async throws {
         var hardlinks: Hardlinks = [:]
         // Allocate a single 128KiB reusable buffer for all files to minimize allocations
         // and reduce the number of read calls to libarchive.
@@ -39,35 +115,33 @@ extension EXT4.Formatter {
                 continue
             }
 
-            defer {
-                // Count the number of entries
-                if let progress {
-                    Task {
-                        await progress([
-                            ProgressEvent(event: "add-items", value: 1)
-                        ])
-                    }
-                }
-            }
-
             pathEntry = preProcessPath(s: pathEntry)
             let path = FilePath(pathEntry)
 
             if path.base.hasPrefix(".wh.") {
                 if path.base == ".wh..wh..opq" {  // whiteout directory
                     try self.unlink(path: path.dir, directoryWhiteout: true)
+                    if let progress {
+                        await progress([ProgressEvent(event: "add-items", value: 1)])
+                    }
                     continue
                 }
                 let startIndex = path.base.index(path.base.startIndex, offsetBy: ".wh.".count)
                 let filePath = String(path.base[startIndex...])
                 let dir: FilePath = path.dir
                 try self.unlink(path: dir.join(filePath))
+                if let progress {
+                    await progress([ProgressEvent(event: "add-items", value: 1)])
+                }
                 continue
             }
 
             if let hardlink = entry.hardlink {
                 let hl = preProcessPath(s: hardlink)
                 hardlinks[path] = FilePath(hl)
+                if let progress {
+                    await progress([ProgressEvent(event: "add-items", value: 1)])
+                }
                 continue
             }
             let ts = FileTimestamps(
@@ -84,13 +158,8 @@ extension EXT4.Formatter {
                     uid: entry.owner,
                     gid: entry.group, xattrs: entry.xattrs, fileBuffer: reusableBuffer)
 
-                // Count the size of files
                 if let progress, let size = entry.size {
-                    Task {
-                        await progress([
-                            ProgressEvent(event: "add-size", value: Int64(size))
-                        ])
-                    }
+                    await progress([ProgressEvent(event: "add-size", value: Int64(size))])
                 }
             case .symbolicLink:
                 var symlinkTarget: FilePath?
@@ -102,8 +171,15 @@ extension EXT4.Formatter {
                     uid: entry.owner,
                     gid: entry.group, xattrs: entry.xattrs)
             default:
+                if let progress {
+                    await progress([ProgressEvent(event: "add-items", value: 1)])
+                }
                 continue
             }
+
+            if let progress {
+                await progress([ProgressEvent(event: "add-items", value: 1)])
+            }
         }
         guard hardlinks.acyclic else {
             throw UnpackError.circularLinks
@@ -115,21 +191,6 @@ extension EXT4.Formatter {
         }
     }
 
-    /// Unpack an archive at the source URL on to the ext4 filesystem.
-    public func unpack(
-        source: URL,
-        format: ContainerizationArchive.Format = .paxRestricted,
-        compression: ContainerizationArchive.Filter = .gzip,
-        progress: ProgressHandler? = nil
-    ) throws {
-        let reader = try ArchiveReader(
-            format: format,
-            filter: compression,
-            file: source
-        )
-        try self.unpack(reader: reader, progress: progress)
-    }
-
     private func preProcessPath(s: String) -> String {
         var p = s
         if p.hasPrefix("./") {