EXT4: Optimize unpack (#378)

Optimize unpack a little by trying to reduce allocations in the hot
path. Today for every file we read the entire file into memory and then
pass the data blob to the ext4 writer to eventually be written to the
sparse file. Before being written to the sparse file the data is copied
*again* to a temp buffer before finally hitting write(2) in FileHandle.

This change moves things around such that we can pass an optional buffer
to the ext4 create() (so we can reuse a buffer for file writes), as well
as stops reading entire files into memory by passing the archive entry
itself (wrapped in a ReaderStream object albeit) down to the writer.

Testing with unpacking every platform for
`docker.io/jenkins/jenkins:lts` on an M1 Max:

Old Avg (5 runs): 7.43s
New Avg (5 runs): 5.31s

Author: dcantah

Sources/ContainerizationArchive/Reader.swift

@@ -17,6 +17,32 @@
 import CArchive
 import Foundation
 
+/// A protocol for reading data in chunks, compatible with both `InputStream` and zero-allocation archive readers.
+public protocol ReadableStream {
+    /// Reads up to `maxLength` bytes into the provided buffer.
+    /// Returns the number of bytes actually read, 0 for EOF, or -1 for error.
+    func read(_ buffer: UnsafeMutablePointer<UInt8>, maxLength: Int) -> Int
+}
+
+extension InputStream: ReadableStream {}
+
+/// Small wrapper type to read data from an archive entry.
+public struct ArchiveEntryReader: ReadableStream {
+    private weak var reader: ArchiveReader?
+
+    init(reader: ArchiveReader) {
+        self.reader = reader
+    }
+
+    /// Reads up to `maxLength` bytes into the provided buffer.
+    /// Returns the number of bytes actually read, 0 for EOF, or -1 for error.
+    public func read(_ buffer: UnsafeMutablePointer<UInt8>, maxLength: Int) -> Int {
+        guard let archive = reader?.underlying else { return -1 }
+        let bytesRead = archive_read_data(archive, buffer, maxLength)
+        return bytesRead < 0 ? -1 : bytesRead
+    }
+}
+
 /// A class responsible for reading entries from an archive file.
 public final class ArchiveReader {
     /// A pointer to the underlying `archive` C structure.
@@ -99,6 +125,29 @@ extension ArchiveReader: Sequence {
         }
     }
 
+    /// Returns an iterator that yields archive entries.
+    public func makeStreamingIterator() -> StreamingIterator {
+        StreamingIterator(reader: self)
+    }
+
+    public struct StreamingIterator: Sequence, IteratorProtocol {
+        var reader: ArchiveReader
+
+        public func makeIterator() -> StreamingIterator {
+            self
+        }
+
+        public mutating func next() -> (WriteEntry, ArchiveEntryReader)? {
+            let entry = WriteEntry()
+            let result = archive_read_next_header2(reader.underlying, entry.underlying)
+            if result == ARCHIVE_EOF {
+                return nil
+            }
+            let streamReader = ArchiveEntryReader(reader: reader)
+            return (entry, streamReader)
+        }
+    }
+
     internal func readDataForEntry(_ entry: WriteEntry) -> Data {
         let bufferSize = Int(Swift.min(entry.size ?? 4096, 4096))
         var entry = Data()

Sources/ContainerizationEXT4/EXT4+Formatter.swift

@@ -16,6 +16,7 @@
 
 //  swiftlint: disable discouraged_direct_init shorthand_operator syntactic_sugar
 
+import ContainerizationArchive
 import ContainerizationOS
 import Foundation
 import SystemPackage
@@ -24,7 +25,7 @@ extension EXT4 {
     /// The `EXT4.Formatter` class provides methods to format a block device with the ext4 filesystem.
     /// It allows customization of block size and maximum disk size.
     public class Formatter {
-        private let blockSize: UInt32
+        let blockSize: UInt32
         private var size: UInt64
         private let groupDescriptorSize: UInt32 = 32
 
@@ -264,7 +265,7 @@ extension EXT4 {
         //    - path: The FilePath representing the path where the file, directory, or symlink should be created.
         //    - link: An optional FilePath representing the target path for a symlink. If `nil`, a regular file or directory will be created. Preceding '/' should be omitted
         //    - mode: The permissions to set for the created file, directory, or symlink.
-        //    - buf: An `InputStream` object providing the contents for the created file. Ignored when creating directories or symlinks.
+        //    - buf: A `ReadableStream` object providing the contents for the created file. Ignored when creating directories or symlinks.
         //
         //  - Note:
         //    - This function recursively creates parent directories if they don't already exist. The `uid` and `gid` of the created parent directories are set to the values of their parent's `uid` and `gid`.
@@ -295,11 +296,12 @@ extension EXT4 {
             link: FilePath? = nil,  // to create symbolic links
             mode: UInt16,
             ts: FileTimestamps = FileTimestamps(),
-            buf: InputStream? = nil,
+            buf: (any ReadableStream)? = nil,
             uid: UInt32? = nil,
             gid: UInt32? = nil,
             xattrs: [String: Data]? = nil,
-            recursion: Bool = false
+            recursion: Bool = false,
+            fileBuffer: UnsafeMutableBufferPointer<UInt8>? = nil
         ) throws {
             if let nodePtr = self.tree.lookup(path: path) {
                 let node = nodePtr.pointee
@@ -539,14 +541,29 @@ extension EXT4 {
             if mode.isReg() {
                 startBlock = self.currentBlock
                 if let buf {  // in case of empty files, this will be nil
-                    let tempBuf = Ptr<UInt8>.allocate(capacity: Int(self.blockSize))
-                    defer { tempBuf.deallocate() }
-                    while case let block = buf.read(tempBuf.underlying, maxLength: Int(self.blockSize)), block > 0 {
+                    let tempBuf: UnsafeMutablePointer<UInt8>
+                    let bufferSize: Int
+                    let shouldDeallocate: Bool
+                    if let fileBuffer {
+                        tempBuf = fileBuffer.baseAddress!
+                        bufferSize = fileBuffer.count
+                        shouldDeallocate = false
+                    } else {
+                        tempBuf = UnsafeMutablePointer<UInt8>.allocate(capacity: Int(self.blockSize))
+                        bufferSize = Int(self.blockSize)
+                        shouldDeallocate = true
+                    }
+                    defer {
+                        if shouldDeallocate {
+                            tempBuf.deallocate()
+                        }
+                    }
+                    while case let block = buf.read(tempBuf, maxLength: bufferSize), block > 0 {
                         size += UInt64(block)
                         if size > EXT4.MaxFileSize {
                             throw Error.fileTooBig(size)
                         }
-                        let data = UnsafeRawBufferPointer(start: tempBuf.underlying, count: block)
+                        let data = UnsafeRawBufferPointer(start: tempBuf, count: block)
                         try withUnsafeLittleEndianBuffer(of: data) { b in
                             try self.handle.write(contentsOf: b)
                         }
@@ -565,29 +582,6 @@ extension EXT4 {
             throw Error.unsupportedFiletype
         }
 
-        public func setOwner(path: FilePath, uid: UInt16? = nil, gid: UInt16? = nil, recursive: Bool = false) throws {
-            // ensure that target exists
-            guard let pathPtr = self.tree.lookup(path: path) else {
-                throw Error.notFound(path)
-            }
-            let pathNode = pathPtr.pointee
-            let pathInodePtr = self.inodes[Int(pathNode.inode) - 1]
-            var pathInode = pathInodePtr.pointee
-            if let uid {
-                pathInode.uid = uid
-            }
-            if let gid {
-                pathInode.gid = gid
-            }
-            pathInodePtr.initialize(to: pathInode)
-            if recursive {
-                for childPtr in pathNode.children {
-                    let child = childPtr.pointee
-                    try self.setOwner(path: path.join(child.name), uid: uid, gid: gid, recursive: recursive)
-                }
-            }
-        }
-
         //  Completes the formatting of an ext4 filesystem after writing the necessary structures.
         //
         //  This function is responsible for finalizing the formatting process of an ext4 filesystem

Sources/ContainerizationEXT4/Formatter+Unpack.swift

@@ -27,7 +27,13 @@ extension EXT4.Formatter {
     /// Unpack the provided archive on to the ext4 filesystem.
     public func unpack(reader: ArchiveReader, progress: ProgressHandler? = nil) throws {
         var hardlinks: Hardlinks = [:]
-        for (entry, data) in reader {
+        // Allocate a single 128KiB reusable buffer for all files to minimize allocations
+        // and reduce the number of read calls to libarchive.
+        let bufferSize = 128 * 1024
+        let reusableBuffer = UnsafeMutableBufferPointer<UInt8>.allocate(capacity: bufferSize)
+        defer { reusableBuffer.deallocate() }
+
+        for (entry, streamReader) in reader.makeStreamingIterator() {
             try Task.checkCancellation()
             guard var pathEntry = entry.path else {
                 continue
@@ -73,20 +79,16 @@ extension EXT4.Formatter {
                     gid: entry.group,
                     xattrs: entry.xattrs)
             case .regular:
-                let inputStream = InputStream(data: data)
-                inputStream.open()
                 try self.create(
-                    path: path, mode: EXT4.Inode.Mode(.S_IFREG, entry.permissions), ts: ts, buf: inputStream,
+                    path: path, mode: EXT4.Inode.Mode(.S_IFREG, entry.permissions), ts: ts, buf: streamReader,
                     uid: entry.owner,
-                    gid: entry.group, xattrs: entry.xattrs)
-                inputStream.close()
+                    gid: entry.group, xattrs: entry.xattrs, fileBuffer: reusableBuffer)
 
                 // Count the size of files
-                if let progress {
+                if let progress, let size = entry.size {
                     Task {
-                        let size = Int64(data.count)
                         await progress([
-                            ProgressEvent(event: "add-size", value: size)
+                            ProgressEvent(event: "add-size", value: Int64(size))
                         ])
                     }
                 }

Sources/cctl/ImageCommand.swift

@@ -123,6 +123,7 @@ extension Application {
                     print("Reference resolved to \(reference.description)")
                 }
 
+                var startTime = ContinuousClock.now
                 let image = try await Images.withAuthentication(ref: normalizedReference) { auth in
                     try await imageStore.pull(reference: normalizedReference, platform: platform, insecure: http, auth: auth)
                 }
@@ -132,7 +133,9 @@ extension Application {
                     Application.exit(withError: POSIXError(.EACCES))
                 }
 
-                print("image pulled")
+                var duration = ContinuousClock.now - startTime
+                print("Image pull took: \(duration)\n")
+
                 guard let unpackPath else {
                     return
                 }
@@ -144,6 +147,7 @@ extension Application {
 
                 let unpacker = EXT4Unpacker.init(blockSizeInBytes: 2.gib())
 
+                startTime = ContinuousClock.now
                 if let platform {
                     let name = platform.description.replacingOccurrences(of: "/", with: "-")
                     let _ = try await unpacker.unpack(image, for: platform, at: unpackUrl.appending(component: name))
@@ -160,6 +164,8 @@ extension Application {
                         print("created snapshot for platform \(descPlatform.description)")
                     }
                 }
+                duration = ContinuousClock.now - startTime
+                print("\nUnpacking took: \(duration)")
             }
         }