Feat: enable configurable default disk capacity for native and builder containers

Sources/ContainerClient/Flags.swift

@@ -74,6 +74,12 @@ public struct Flags {
             help: "Amount of memory (1MiByte granularity), with optional K, M, G, T, or P suffix"
         )
         public var memory: String?
+
+        @Option(
+            name: .shortAndLong,
+            help: "Disk capacity / storage size for the container"
+        )
+        public var storage: String?
     }
 
     public struct Registry: ParsableArguments {

Sources/ContainerClient/Parser.swift

@@ -84,14 +84,18 @@ public struct Parser {
         try .init(from: platform)
     }
 
-    public static func resources(cpus: Int64?, memory: String?) throws -> ContainerConfiguration.Resources {
+    public static func resources(cpus: Int64?, memory: String?, storage: String?) throws -> ContainerConfiguration.Resources {
         var resource = ContainerConfiguration.Resources()
         if let cpus {
             resource.cpus = Int(cpus)
         }
         if let memory {
             resource.memoryInBytes = try Parser.memoryString(memory).mib()
         }
+        if let storage {
+            let storageInMiB = try Parser.memoryString(storage)
+            resource.storage = UInt64(storageInMiB.mib())
+        }
         return resource
     }
 

Sources/ContainerClient/Utility.swift

@@ -168,9 +168,35 @@ public struct Utility {
         var config = ContainerConfiguration(id: id, image: description, process: pc)
         config.platform = requestedPlatform
 
+        let effectiveStorage: String?
+        if let storage = resource.storage {
+            do {
+                _ = try Parser.memoryString(storage)
+            } catch {
+                throw ContainerizationError(
+                    .invalidArgument,
+                    message: "invalid storage value '\(storage)' for --storage"
+                )
+            }
+            effectiveStorage = storage
+        } else if let defaultStorage: String = DefaultsStore.getOptional(key: .defaultContainerStorage) {
+            do {
+                _ = try Parser.memoryString(defaultStorage)
+            } catch {
+                throw ContainerizationError(
+                    .invalidArgument,
+                    message: "invalid default container storage value '\(defaultStorage)'; update it with `container property set defaultContainerStorage`"
+                )
+            }
+            effectiveStorage = defaultStorage
+        } else {
+            effectiveStorage = nil
+        }
+
         config.resources = try Parser.resources(
             cpus: resource.cpus,
-            memory: resource.memory
+            memory: resource.memory,
+            storage: effectiveStorage
         )
 
         let tmpfs = try Parser.tmpfsMounts(management.tmpFs)

Sources/ContainerCommands/BuildCommand.swift

@@ -81,6 +81,12 @@ extension Application {
         )
         var memory: String = "2048MB"
 
+        @Option(
+            name: .long,
+            help: "Disk capacity for the builder container"
+        )
+        var storage: String?
+
         @Flag(name: .long, help: "Do not use cache")
         var noCache: Bool = false
 
@@ -140,12 +146,12 @@ extension Application {
 
                 progress.set(description: "Dialing builder")
 
-                let builder: Builder? = try await withThrowingTaskGroup(of: Builder.self) { [vsockPort, cpus, memory] group in
+                let builder: Builder? = try await withThrowingTaskGroup(of: Builder.self) { [vsockPort, cpus, memory, storage] group in
                     defer {
                         group.cancelAll()
                     }
 
-                    group.addTask { [vsockPort, cpus, memory] in
+                    group.addTask { [vsockPort, cpus, memory, storage] in
                         while true {
                             do {
                                 let container = try await ClientContainer.get(id: "buildkit")
@@ -166,6 +172,7 @@ extension Application {
                                 try await BuilderStart.start(
                                     cpus: cpus,
                                     memory: memory,
+                                    storage: storage,
                                     progressUpdate: progress.handler
                                 )
 

Sources/ContainerCommands/Builder/BuilderStart.swift

@@ -44,6 +44,12 @@ extension Application {
         )
         var memory: String = "2048MB"
 
+        @Option(
+            name: .long,
+            help: "Disk capacity for the builder container"
+        )
+        var storage: String?
+
         @OptionGroup
         var global: Flags.Global
 
@@ -60,11 +66,11 @@ extension Application {
                 progress.finish()
             }
             progress.start()
-            try await Self.start(cpus: self.cpus, memory: self.memory, progressUpdate: progress.handler)
+            try await Self.start(cpus: self.cpus, memory: self.memory, storage: self.storage, progressUpdate: progress.handler)
             progress.finish()
         }
 
-        static func start(cpus: Int64?, memory: String?, progressUpdate: @escaping ProgressUpdateHandler) async throws {
+        static func start(cpus: Int64?, memory: String?, storage: String?, progressUpdate: @escaping ProgressUpdateHandler) async throws {
             await progressUpdate([
                 .setDescription("Fetching BuildKit image"),
                 .setItemsName("blobs"),
@@ -88,6 +94,32 @@ extension Application {
 
             let builderPlatform = ContainerizationOCI.Platform(arch: "arm64", os: "linux", variant: "v8")
 
+            // Decide which storage string to use for the builder
+            let effectiveStorage: String?
+            if let storage {
+                do {
+                    _ = try Parser.memoryString(storage)
+                } catch {
+                    throw ContainerizationError(
+                        .invalidArgument,
+                        message: "invalid storage value '\(storage)' for --storage"
+                    )
+                }
+                effectiveStorage = storage
+            } else if let defaultStorage: String = DefaultsStore.getOptional(key: .defaultBuilderStorage) {
+                do {
+                    _ = try Parser.memoryString(defaultStorage)
+                } catch {
+                    throw ContainerizationError(
+                        .invalidArgument,
+                        message: "invalid default builder storage value '\(defaultStorage)'; update it with `container property set defaultBuilderStorage`"
+                    )
+                }
+                effectiveStorage = defaultStorage
+            } else {
+                effectiveStorage = nil
+            }
+
             let existingContainer = try? await ClientContainer.get(id: "buildkit")
             if let existingContainer {
                 let existingImage = existingContainer.configuration.image.reference
@@ -103,19 +135,28 @@ extension Application {
                     }
                     return false
                 }()
+
                 let memChanged = try {
                     if let memory {
-                        let memoryInBytes = try Parser.resources(cpus: nil, memory: memory).memoryInBytes
-                        if existingResources.memoryInBytes != memoryInBytes {
-                            return true
-                        }
+                        let memoryInMiB = try Parser.memoryString(memory)
+                        let memoryInBytes = UInt64(memoryInMiB.mib())
+                        return existingResources.memoryInBytes != memoryInBytes
                     }
                     return false
                 }()
 
+                let storageChanged = try {
+                    if let effectiveStorage {
+                        let storageInMiB = try Parser.memoryString(effectiveStorage)
+                        let storageInBytes = UInt64(storageInMiB.mib())
+                        return existingResources.storage != storageInBytes
+                    }
+                    return existingResources.storage != 0
+                }()
+
                 switch existingContainer.status {
                 case .running:
-                    guard imageChanged || cpuChanged || memChanged else {
+                    guard imageChanged || cpuChanged || memChanged || storageChanged else {
                         // If image, mem and cpu are the same, continue using the existing builder
                         return
                     }
@@ -125,7 +166,7 @@ extension Application {
                 case .stopped:
                     // If the builder is stopped and matches our requirements, start it
                     // Otherwise, delete it and create a new one
-                    guard imageChanged || cpuChanged || memChanged else {
+                    guard imageChanged || cpuChanged || memChanged || storageChanged else {
                         try await existingContainer.startBuildKit(progressUpdate, nil)
                         return
                     }
@@ -184,10 +225,12 @@ extension Application {
 
             let resources = try Parser.resources(
                 cpus: cpus,
-                memory: memory
+                memory: memory,
+                storage: effectiveStorage
             )
 
             var config = ContainerConfiguration(id: id, image: imageDesc, process: processConfig)
+            config.platform = builderPlatform
             config.resources = resources
             config.mounts = [
                 .init(

Sources/ContainerCommands/System/Property/PropertySet.swift

@@ -74,6 +74,9 @@ extension Application {
                     throw ContainerizationError(.invalidArgument, message: "invalid CIDRv4 address: \(value)")
                 }
                 DefaultsStore.set(value: value, key: key)
+            case .defaultBuilderStorage, .defaultContainerStorage:
+                _ = try Parser.memoryString(value)
+                DefaultsStore.set(value: value, key: key)
             }
         }
     }

Sources/ContainerPersistence/DefaultsStore.swift

@@ -26,6 +26,8 @@ public enum DefaultsStore {
         case buildRosetta = "build.rosetta"
         case defaultDNSDomain = "dns.domain"
         case defaultBuilderImage = "image.builder"
+        case defaultBuilderStorage = "builder.storage"
+        case defaultContainerStorage = "container.storage"
         case defaultInitImage = "image.init"
         case defaultKernelBinaryPath = "kernel.binaryPath"
         case defaultKernelURL = "kernel.url"
@@ -69,6 +71,8 @@ public enum DefaultsStore {
         let allKeys: [(Self.Keys, (Self.Keys) -> Any?)] = [
             (.buildRosetta, { Self.getBool(key: $0) }),
             (.defaultBuilderImage, { Self.get(key: $0) }),
+            (.defaultBuilderStorage, { Self.getOptional(key: $0) }),
+            (.defaultContainerStorage, { Self.getOptional(key: $0) }),
             (.defaultInitImage, { Self.get(key: $0) }),
             (.defaultKernelBinaryPath, { Self.get(key: $0) }),
             (.defaultKernelURL, { Self.get(key: $0) }),
@@ -124,6 +128,10 @@ extension DefaultsStore.Keys {
             return "If defined, the local DNS domain to use for containers with unqualified names."
         case .defaultBuilderImage:
             return "The image reference for the utility container that `container build` uses."
+        case .defaultBuilderStorage:
+            return "Default disk capacity for the builder container."
+        case .defaultContainerStorage:
+            return "Default disk capacity for native containers."
         case .defaultInitImage:
             return "The image reference for the default initial filesystem image."
         case .defaultKernelBinaryPath:
@@ -145,6 +153,10 @@ extension DefaultsStore.Keys {
             return String.self
         case .defaultBuilderImage:
             return String.self
+        case .defaultBuilderStorage:
+            return String.self
+        case .defaultContainerStorage:
+            return String.self
         case .defaultInitImage:
             return String.self
         case .defaultKernelBinaryPath:
@@ -168,6 +180,10 @@ extension DefaultsStore.Keys {
         case .defaultBuilderImage:
             let tag = String(cString: get_container_builder_shim_version())
             return "ghcr.io/apple/container-builder-shim/builder:\(tag)"
+        case .defaultBuilderStorage:
+            return ""
+        case .defaultContainerStorage:
+            return ""
         case .defaultInitImage:
             let tag = String(cString: get_swift_containerization_version())
             guard tag != "latest" else {

Sources/Services/ContainerSandboxService/SandboxService.swift

@@ -805,6 +805,7 @@ public actor SandboxService {
     ) throws {
         czConfig.cpus = config.resources.cpus
         czConfig.memoryInBytes = config.resources.memoryInBytes
+        czConfig.storageInBytes = config.resources.storage
         czConfig.sysctl = config.sysctls.reduce(into: [String: String]()) {
             $0[$1.key] = $1.value
         }

Tests/CLITests/Subcommands/Build/CLIBuilderLifecycleTest.swift

@@ -14,6 +14,7 @@
 // limitations under the License.
 //===----------------------------------------------------------------------===//
 
+import ContainerClient
 import Foundation
 import Testing
 
@@ -33,5 +34,47 @@ extension TestCLIBuildBase {
                 #expect(status == "stopped", "BuildKit container is not stopped")
             }
         }
+
+        @Test func testBuilderStorageFlag() async throws {
+            do {
+                let requestedStorage = "4096MB"
+                let expectedMiB = Int(try Parser.memoryString(requestedStorage))
+                let expectedBytes = UInt64(expectedMiB.mib())
+
+                try? builderStop()
+                try? builderDelete(force: true)
+
+                let (_, _, err, status) = try run(arguments: [
+                    "builder",
+                    "start",
+                    "--storage", requestedStorage,
+                ])
+                try #require(status == 0, "builder start failed: \(err)")
+
+                try waitForBuilderRunning()
+
+                defer {
+                    try? builderStop()
+                    try? builderDelete(force: true)
+                }
+
+                let buildkitName = "buildkit"
+                let buildkit = try await ClientContainer.get(id: buildkitName)
+                let resources = buildkit.configuration.resources
+
+                guard let storageBytes = resources.storage else {
+                    Issue.record("expected builder resources.storage to be set for --storage \(requestedStorage)")
+                    return
+                }
+
+                #expect(
+                    storageBytes == expectedBytes,
+                    "expected builder storage \(expectedBytes) bytes for \(requestedStorage), got \(storageBytes) bytes"
+                )
+            } catch {
+                Issue.record("failed to verify builder storage: \(error)")
+                return
+            }
+        }
     }
 }

Tests/CLITests/Subcommands/Run/TestCLIRunCommand.swift

@@ -578,6 +578,41 @@ class TestCLIRunCommand: CLITest {
         }
     }
 
+    @Test
+    func testRunCommandStorage() async throws {
+        do {
+            let name = getTestName()
+            let requestedStorage = "2048MB"
+            let expectedMiB = Int(try Parser.memoryString(requestedStorage))
+            let expectedBytes = UInt64(expectedMiB.mib())
+
+            try doLongRun(name: name, args: ["--storage", requestedStorage])
+            defer {
+                try? doStop(name: name)
+            }
+
+            // Inspect configuration via the client instead of df
+            let container = try await ClientContainer.get(id: name)
+            let resources = container.configuration.resources
+
+            guard let storageBytes = resources.storage else {
+                Issue.record("expected container resources.storage to be set for --storage \(requestedStorage)")
+                return
+            }
+
+            #expect(
+                storageBytes == expectedBytes,
+                "expected container storage \(expectedBytes) bytes for \(requestedStorage), got \(storageBytes) bytes"
+            )
+
+            try doStop(name: name)
+        } catch {
+            Issue.record("failed to run container \(error)")
+            return
+        }
+
+    }
+
     func getDefaultDomain() throws -> String? {
         let (_, output, err, status) = try run(arguments: ["system", "property", "get", "dns.domain"])
         try #require(status == 0, "default DNS domain retrieval returned status \(status): \(err)")