shows snapshotSize if available with fallback to compressedSize; fix tests

Bortnyak · Bortnyak · commit dfa0ceb47965 · 2026-01-18T23:17:47.000+01:00
diff --git a/Package.swift b/Package.swift
@@ -251,6 +251,15 @@ let package = Package(
                 "ContainerXPC",
             ],
             path: "Sources/Helpers/NetworkVmnet"
+        )
+        .testTarget(
+            name: "ContainerImagesServiceTests",
+            dependencies: [
+                .product(name: "Containerization", package: "containerization"),
+                .product(name: "ContainerizationOCI", package: "containerization"),
+                "ContainerImagesService",
+                "ContainerPersistence",
+            ]
         ),
         .target(
             name: "ContainerNetworkService",
diff --git a/Sources/ContainerCommands/Image/ImageList.swift b/Sources/ContainerCommands/Image/ImageList.swift
@@ -80,7 +80,9 @@ extension Application {
                     }
 
                     let created = config.created ?? ""
-                    let size = descriptor.size + manifest.config.size + manifest.layers.reduce(0, { (l, r) in l + r.size })
+                    let compressedSize = descriptor.size + manifest.config.size + manifest.layers.reduce(0, { (l, r) in l + r.size })
+                    let snapshotSize = (try? await image.getSnapshotSize(platform: platform)) ?? 0
+                    let size = snapshotSize > 0 ? Int64(snapshotSize) : compressedSize
                     let formattedSize = formatter.string(fromByteCount: size)
 
                     let processedReferenceString = try ClientImage.denormalizeReference(image.reference)
diff --git a/Sources/Helpers/Images/ImagesHelper.swift b/Sources/Helpers/Images/ImagesHelper.swift
@@ -102,6 +102,7 @@ extension ImagesHelper {
             routes[ImagesServiceXPCRoute.imageDiskUsage.rawValue] = harness.calculateDiskUsage
             routes[ImagesServiceXPCRoute.snapshotDelete.rawValue] = harness.deleteSnapshot
             routes[ImagesServiceXPCRoute.snapshotGet.rawValue] = harness.getSnapshot
+            routes[ImagesServiceXPCRoute.snapshotSize.rawValue] = harness.getSnapshotSize
         }
 
         private func initializeContentService(root: URL, log: Logger, routes: inout [String: XPCServer.RouteHandler]) throws {
diff --git a/Sources/Services/ContainerAPIService/Client/ClientImage.swift b/Sources/Services/ContainerAPIService/Client/ClientImage.swift
@@ -420,6 +420,17 @@ extension ClientImage {
         return fs
     }
 
+    public func getSnapshotSize(platform: Platform) async throws -> UInt64 {
+        let client = Self.newXPCClient()
+        let request = Self.newRequest(.snapshotSize)
+
+        try request.set(description: description)
+        try request.set(platform: platform)
+
+        let response = try await client.send(request)
+        return response.uint64(key: .imageSize)
+    }
+
     @discardableResult
     public func getCreateSnapshot(platform: Platform, progressUpdate: ProgressUpdateHandler? = nil) async throws -> Filesystem {
         do {
diff --git a/Sources/Services/ContainerImagesService/Client/ImageServiceXPCRoutes.swift b/Sources/Services/ContainerImagesService/Client/ImageServiceXPCRoutes.swift
@@ -40,6 +40,7 @@ public enum ImagesServiceXPCRoute: String {
     case imageUnpack
     case snapshotDelete
     case snapshotGet
+    case snapshotSize
 }
 
 extension XPCMessage {
diff --git a/Sources/Services/ContainerImagesService/Server/ImageService.swift b/Sources/Services/ContainerImagesService/Server/ImageService.swift
@@ -206,6 +206,15 @@ extension ImagesService {
         let img = try await self._get(description)
         return try await self.snapshotStore.get(for: img, platform: platform)
     }
+
+    public func getSnapshotSize(description: ImageDescription, platform: Platform) async throws -> UInt64 {
+        self.log.info("ImagesService: \(#function) - description: \(description), platform: \(String(describing: platform))")
+        let img = try await self._get(description)
+        let descriptor = try await img.descriptor(for: platform)
+        // getSnapshotSize returns 0 if snapshot doesn't exist, and can throw on filesystem errors
+        // We catch errors and return 0 to match the behavior when snapshot doesn't exist
+        return (try? await self.snapshotStore.getSnapshotSize(descriptor: descriptor)) ?? 0
+    }
 }
 
 // MARK: Static Methods
diff --git a/Sources/Services/ContainerImagesService/Server/ImagesServiceHarness.swift b/Sources/Services/ContainerImagesService/Server/ImagesServiceHarness.swift
@@ -282,4 +282,28 @@ extension ImagesServiceHarness {
         reply.set(key: .filesystem, value: fsData)
         return reply
     }
+
+    @Sendable
+    public func getSnapshotSize(_ message: XPCMessage) async throws -> XPCMessage {
+        let descriptionData = message.dataNoCopy(key: .imageDescription)
+        guard let descriptionData else {
+            throw ContainerizationError(
+                .invalidArgument,
+                message: "missing image description"
+            )
+        }
+        let description = try JSONDecoder().decode(ImageDescription.self, from: descriptionData)
+        let platformData = message.dataNoCopy(key: .ociPlatform)
+        guard let platformData else {
+            throw ContainerizationError(
+                .invalidArgument,
+                message: "missing OCI platform"
+            )
+        }
+        let platform = try JSONDecoder().decode(ContainerizationOCI.Platform.self, from: platformData)
+        let size = try await self.service.getSnapshotSize(description: description, platform: platform)
+        let reply = message.reply()
+        reply.set(key: .imageSize, value: size)
+        return reply
+    }
 }
diff --git a/Tests/ContainerImagesServiceTests/SnapshotStoreTests.swift b/Tests/ContainerImagesServiceTests/SnapshotStoreTests.swift
@@ -0,0 +1,184 @@
+//===----------------------------------------------------------------------===//
+// Copyright © 2025 Apple Inc. and the container project authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//===----------------------------------------------------------------------===//
+
+import ContainerImagesService
+import ContainerizationOCI
+import Foundation
+import Testing
+
+@testable import ContainerImagesService
+
+struct SnapshotStoreTests {
+
+    @Test("getSnapshotSize should return 0 when snapshot directory does not exist")
+    func testGetSnapshotSizeWhenSnapshotDoesNotExist() async throws {
+        // prepare
+        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+        try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
+        defer {
+            try? FileManager.default.removeItem(at: tempDir)
+        }
+
+        let snapshotStore = try SnapshotStore(
+            path: tempDir,
+            unpackStrategy: SnapshotStore.defaultUnpackStrategy,
+            log: nil
+        )
+
+        // Create a descriptor for a non-existent snapshot
+        let fakeDigest = "sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"
+        let platform = Platform(arch: "amd64", os: "linux")
+
+        // Create descriptor directly with the properties
+        let descriptor = Descriptor(
+            mediaType: "application/vnd.oci.image.manifest.v1+json",
+            digest: fakeDigest,
+            size: 1024,
+            platform: platform
+        )
+
+        let size = try await snapshotStore.getSnapshotSize(descriptor: descriptor)
+        
+        #expect(size == 0, "Expected size to be 0 when snapshot doesn't exist, got \(size)")
+    }
+
+    @Test("getSnapshotSize should return correct size when snapshot exists")
+    func testGetSnapshotSizeWhenSnapshotExists() async throws {
+        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+        try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
+        defer {
+            try? FileManager.default.removeItem(at: tempDir)
+        }
+
+        let snapshotStore = try SnapshotStore(
+            path: tempDir,
+            unpackStrategy: SnapshotStore.defaultUnpackStrategy,
+            log: nil
+        )
+
+        // Create a test descriptor
+        let fakeDigest = "sha256:abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890"
+        let descriptor = Descriptor(
+            mediaType: "application/vnd.oci.image.manifest.v1+json",
+            digest: fakeDigest,
+            size: 1024,
+            platform: Platform(arch: "amd64", os: "linux")
+        )
+
+        // Create snapshot directory
+        let snapshotDir =
+            tempDir
+            .appendingPathComponent("snapshots")
+            .appendingPathComponent(descriptor.digest.trimmingDigestPrefix)
+        try FileManager.default.createDirectory(at: snapshotDir, withIntermediateDirectories: true)
+
+        // Create test files: 1KB snapshot + 512 bytes info = 1536 bytes total
+        let testFile = snapshotDir.appendingPathComponent("snapshot")
+        try Data(repeating: 0, count: 1024).write(to: testFile)
+
+        let infoFile = snapshotDir.appendingPathComponent("snapshot-info")
+        try Data(repeating: 0, count: 512).write(to: infoFile)
+
+        let size = try await snapshotStore.getSnapshotSize(descriptor: descriptor)
+
+        #expect(size >= 1536, "Expected at least 1536 bytes (1024 + 512), got \(size) allocated")
+    }
+
+    @Test("getSnapshotSize should handle multiple files in snapshot directory")
+    func testGetSnapshotSizeWithMultipleFiles() async throws {
+        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+        try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
+        defer {
+            try? FileManager.default.removeItem(at: tempDir)
+        }
+
+        let snapshotStore = try SnapshotStore(
+            path: tempDir,
+            unpackStrategy: SnapshotStore.defaultUnpackStrategy,
+            log: nil
+        )
+
+        let fakeDigest = "sha256:fedcba0987654321fedcba0987654321fedcba0987654321fedcba0987654321"
+        let platform = Platform(arch: "arm64", os: "linux")
+        let descriptor = Descriptor(
+            mediaType: "application/vnd.oci.image.manifest.v1+json",
+            digest: fakeDigest,
+            size: 2048,
+            platform: platform
+        )
+
+        let snapshotDir =
+            tempDir
+            .appendingPathComponent("snapshots", isDirectory: true)
+            .appendingPathComponent(descriptor.digest.trimmingDigestPrefix, isDirectory: true)
+        try FileManager.default.createDirectory(at: snapshotDir, withIntermediateDirectories: true)
+
+        // Create multiple files
+        let file1 = snapshotDir.appendingPathComponent("snapshot", isDirectory: false)
+        try Data(repeating: 1, count: 2048).write(to: file1)  // 2KB
+
+        let file2 = snapshotDir.appendingPathComponent("snapshot-info", isDirectory: false)
+        try Data(repeating: 2, count: 1024).write(to: file2)  // 1KB
+
+        let subdir = snapshotDir.appendingPathComponent("subdir", isDirectory: true)
+        try FileManager.default.createDirectory(at: subdir, withIntermediateDirectories: true)
+        let file3 = subdir.appendingPathComponent("extra", isDirectory: false)
+        try Data(repeating: 3, count: 512).write(to: file3)  // 512 bytes
+
+        let expectedMinSize: UInt64 = 2048 + 1024 + 512  // 3.5KB
+
+        let size = try await snapshotStore.getSnapshotSize(descriptor: descriptor)
+
+        #expect(size >= expectedMinSize, "Expected size to include all files (at least \(expectedMinSize) bytes), got \(size)")
+    }
+
+    @Test("getSnapshotSize should handle empty snapshot directory")
+    func testGetSnapshotSizeWithEmptyDirectory() async throws {
+        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+        try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
+        defer {
+            try? FileManager.default.removeItem(at: tempDir)
+        }
+
+        let snapshotStore = try SnapshotStore(
+            path: tempDir,
+            unpackStrategy: SnapshotStore.defaultUnpackStrategy,
+            log: nil
+        )
+
+        let fakeDigest = "sha256:0000000000000000000000000000000000000000000000000000000000000000"
+        let platform = Platform(arch: "amd64", os: "linux")
+        let descriptor = Descriptor(
+            mediaType: "application/vnd.oci.image.manifest.v1+json",
+            digest: fakeDigest,
+            size: 0,
+            platform: platform
+        )
+
+        // Create empty snapshot directory
+        let snapshotDir =
+            tempDir
+            .appendingPathComponent("snapshots", isDirectory: true)
+            .appendingPathComponent(descriptor.digest.trimmingDigestPrefix, isDirectory: true)
+        try FileManager.default.createDirectory(at: snapshotDir, withIntermediateDirectories: true)
+
+
+        let size = try await snapshotStore.getSnapshotSize(descriptor: descriptor)
+
+        // Empty directory should return 0 or a small overhead value
+        #expect(size >= 0, "Expected size to be non-negative, got \(size)")
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -80,7 +80,9 @@ extension Application {`
`80`	`80`	`}`
`81`	`81`
`82`	`82`	`let created = config.created ?? ""`
`83`		`- let size = descriptor.size + manifest.config.size + manifest.layers.reduce(0, { (l, r) in l + r.size })`
	`83`	`+ let compressedSize = descriptor.size + manifest.config.size + manifest.layers.reduce(0, { (l, r) in l + r.size })`
	`84`	`+ let snapshotSize = (try? await image.getSnapshotSize(platform: platform)) ?? 0`
	`85`	`+ let size = snapshotSize > 0 ? Int64(snapshotSize) : compressedSize`
`84`	`86`	`let formattedSize = formatter.string(fromByteCount: size)`
`85`	`87`
`86`	`88`	`let processedReferenceString = try ClientImage.denormalizeReference(image.reference)`
Original file line number	Diff line number	Diff line change
`@@ -102,6 +102,7 @@ extension ImagesHelper {`
`102`	`102`	`routes[ImagesServiceXPCRoute.imageDiskUsage.rawValue] = harness.calculateDiskUsage`
`103`	`103`	`routes[ImagesServiceXPCRoute.snapshotDelete.rawValue] = harness.deleteSnapshot`
`104`	`104`	`routes[ImagesServiceXPCRoute.snapshotGet.rawValue] = harness.getSnapshot`
	`105`	`+ routes[ImagesServiceXPCRoute.snapshotSize.rawValue] = harness.getSnapshotSize`
`105`	`106`	`}`
`106`	`107`
`107`	`108`	`private func initializeContentService(root: URL, log: Logger, routes: inout [String: XPCServer.RouteHandler]) throws {`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@ public enum ImagesServiceXPCRoute: String {`
`40`	`40`	`case imageUnpack`
`41`	`41`	`case snapshotDelete`
`42`	`42`	`case snapshotGet`
	`43`	`+ case snapshotSize`
`43`	`44`	`}`
`44`	`45`
`45`	`46`	`extension XPCMessage {`