Add unpack strategy to SnapshotStore (#274)

Define a `UnpackStrategy` function type in the `SnapshotStore` to give
more control over how an image is unpacked.

Previously, we were creating a 512 GB sparse block file for the initial
file system of a container, which is overkill.

With this change, the vminit image is unpacked to a smaller block file,
while container images are unpacked to the 512 GB block

Follows the same pattern as
https://github.com/apple/container/blob/main/Sources/Helpers/RuntimeLinux/RuntimeLinuxHelper.swift#L71

Signed-off-by: Aditya Ramani <a_ramani@apple.com>

Package.resolved

@@ -1,5 +1,5 @@
 {
-  "originHash" : "cc0718ffe17715cc940f40d92b33532b6dd1767065ccc6039f7174402ba5930f",
+  "originHash" : "912a08d7a8cc6a403ab02a8d236337b39276fa0d4f675ceb6d69d5d624af8a49",
   "pins" : [
     {
       "identity" : "async-http-client",
@@ -15,8 +15,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/apple/containerization.git",
       "state" : {
-        "revision" : "1a59de82b052a86edd9e2bd6f023d212a061b0b4",
-        "version" : "0.2.0"
+        "revision" : "bdba5b5740be22a668f7a6d78e2cdd1efc339418",
+        "version" : "0.3.0"
       }
     },
     {

Package.swift

@@ -26,7 +26,7 @@ if let path = ProcessInfo.processInfo.environment["CONTAINERIZATION_PATH"] {
     scDependency = .package(path: path)
     scVersion = "latest"
 } else {
-    scVersion = "0.2.0"
+    scVersion = "0.3.0"
     scDependency = .package(url: "https://github.com/apple/containerization.git", exact: Version(stringLiteral: scVersion))
 }
 

Sources/Helpers/Images/ImagesHelper.swift

@@ -60,6 +60,8 @@ extension ImagesHelper {
             .appendingPathComponent("com.apple.container")
         }()
 
+        private static let unpackStrategy = SnapshotStore.defaultUnpackStrategy
+
         func run() async throws {
             let commandName = ImagesHelper._commandName
             let log = setupLogger()
@@ -89,7 +91,7 @@ extension ImagesHelper {
         private func initializeImagesService(root: URL, log: Logger, routes: inout [String: XPCServer.RouteHandler]) throws {
             let contentStore = RemoteContentStoreClient()
             let imageStore = try ImageStore(path: root, contentStore: contentStore)
-            let snapshotStore = try SnapshotStore(path: root)
+            let snapshotStore = try SnapshotStore(path: root, unpackStrategy: Self.unpackStrategy, log: log)
             let service = try ImagesService(contentStore: contentStore, imageStore: imageStore, snapshotStore: snapshotStore, log: log)
             let harness = ImagesServiceHarness(service: service, log: log)
 

Sources/Services/ContainerImagesService/Server/SnapshotStore.swift

@@ -21,6 +21,7 @@ import ContainerizationExtras
 import ContainerizationOCI
 import ContainerizationOS
 import Foundation
+import Logging
 import TerminalProgress
 
 public actor SnapshotStore {
@@ -28,14 +29,33 @@ public actor SnapshotStore {
     private static let snapshotInfoFileName = "snapshot-info"
     private static let ingestDirName = "ingest"
 
+    /// Return the Unpacker to use for a given image.
+    /// If the given platform for the image cannot be unpacked return `nil`.
+    public typealias UnpackStrategy = @Sendable (Containerization.Image, Platform) async throws -> Unpacker?
+
+    public static let defaultUnpackStrategy: UnpackStrategy = { image, platform in
+        guard platform.os == "linux" else {
+            return nil
+        }
+        var minBlockSize = 512.gib()
+        if image.reference == ClientDefaults.get(key: .defaultInitImage) {
+            minBlockSize = 512.mib()
+        }
+        return EXT4Unpacker(blockSizeInBytes: minBlockSize)
+    }
+
     let path: URL
     let fm = FileManager.default
     let ingestDir: URL
+    let unpackStrategy: UnpackStrategy
+    let log: Logger?
 
-    public init(path: URL) throws {
+    public init(path: URL, unpackStrategy: @escaping UnpackStrategy, log: Logger?) throws {
         let root = path.appendingPathComponent("snapshots")
         self.path = root
         self.ingestDir = self.path.appendingPathComponent(Self.ingestDirName)
+        self.unpackStrategy = unpackStrategy
+        self.log = log
         try self.fm.createDirectory(at: root, withIntermediateDirectories: true)
         try self.fm.createDirectory(at: self.ingestDir, withIntermediateDirectories: true)
     }
@@ -62,7 +82,8 @@ public actor SnapshotStore {
             guard let platform = desc.platform else {
                 throw ContainerizationError(.internalError, message: "Missing platform for descriptor \(desc.digest)")
             }
-            guard Self.shouldUnpackPlatform(platform) else {
+            guard let unpacker = try await self.unpackStrategy(image, platform) else {
+                self.log?.warning("Skipping unpack for \(image.reference) for platform \(platform.description). No unpacker configured.")
                 continue
             }
             let currentSubTask = await taskManager.startTask()
@@ -79,7 +100,8 @@ public actor SnapshotStore {
             let tempSnapshotPath = tempDir.appendingPathComponent(Self.snapshotFileName, isDirectory: false)
             let infoPath = tempDir.appendingPathComponent(Self.snapshotInfoFileName, isDirectory: false)
             do {
-                let mount = try await image.unpack(for: platform, at: tempSnapshotPath, progress: ContainerizationProgressAdapter.handler(from: taskUpdateProgress))
+                let progress = ContainerizationProgressAdapter.handler(from: taskUpdateProgress)
+                let mount = try await unpacker.unpack(image, for: platform, at: tempSnapshotPath, progress: progress)
                 let fs = Filesystem.block(
                     format: mount.type,
                     source: self.snapshotPath(desc).absolutePath(),
@@ -186,13 +208,6 @@ public actor SnapshotStore {
         try self.fm.createDirectory(at: uniqueDirectoryURL, withIntermediateDirectories: true, attributes: nil)
         return uniqueDirectoryURL
     }
-
-    private static func shouldUnpackPlatform(_ platform: Platform) -> Bool {
-        guard platform.os == "linux" else {
-            return false
-        }
-        return true
-    }
 }
 
 extension FileManager {