Cgroup2: Add ability to filter stats

vminitd/Sources/Cgroup/Cgroup2Manager.swift

@@ -343,17 +343,12 @@ public struct Cgroup2Manager: Sendable {
         }
     }
 
-    package func stats() throws -> Cgroup2Stats {
-        let pidsStats = try self.readPidsStats()
-        let memoryStats = try self.readMemoryStats()
-        let cpuStats = try self.readCPUStats()
-        let ioStats = try self.readIOStats()
-
-        return Cgroup2Stats(
-            pids: pidsStats,
-            memory: memoryStats,
-            cpu: cpuStats,
-            io: ioStats
+    package func stats(_ categories: Cgroup2StatsCategory = .all) throws -> Cgroup2Stats {
+        Cgroup2Stats(
+            pids: categories.contains(.pids) ? try self.readPidsStats() : nil,
+            memory: categories.contains(.memory) ? try self.readMemoryStats() : nil,
+            cpu: categories.contains(.cpu) ? try self.readCPUStats() : nil,
+            io: categories.contains(.io) ? try self.readIOStats() : nil
         )
     }
 
@@ -525,6 +520,22 @@ public struct Cgroup2Manager: Sendable {
     }
 }
 
+// Selects which cgroup stat groups to read.
+package struct Cgroup2StatsCategory: OptionSet, Sendable {
+    package let rawValue: UInt8
+
+    package init(rawValue: UInt8) {
+        self.rawValue = rawValue
+    }
+
+    package static let pids = Cgroup2StatsCategory(rawValue: 1 << 0)
+    package static let memory = Cgroup2StatsCategory(rawValue: 1 << 1)
+    package static let cpu = Cgroup2StatsCategory(rawValue: 1 << 2)
+    package static let io = Cgroup2StatsCategory(rawValue: 1 << 3)
+
+    package static let all: Cgroup2StatsCategory = [.pids, .memory, .cpu, .io]
+}
+
 package struct Cgroup2Stats: Sendable {
     package var pids: PidsStats?
     package var memory: MemoryStats?

vminitd/Sources/VminitdCore/ManagedContainer.swift

@@ -228,8 +228,8 @@ extension ManagedContainer {
         }
     }
 
-    func stats() throws -> Cgroup2Stats {
-        try self.cgroupManager.stats()
+    func stats(_ categories: Cgroup2StatsCategory = .all) throws -> Cgroup2Stats {
+        try self.cgroupManager.stats(categories)
     }
 
     func getMemoryEvents() throws -> MemoryEvents {

vminitd/Sources/VminitdCore/MemoryMonitor.swift

@@ -107,7 +107,7 @@ package final class MemoryMonitor: Sendable {
                 if events.high > highCountMax {
                     highCountMax = events.high
 
-                    let stats = try cgroupManager.stats()
+                    let stats = try cgroupManager.stats(.memory)
                     let currentUsage = stats.memory?.usage ?? 0
 
                     onThresholdExceeded(currentUsage, events.high)

vminitd/Sources/VminitdCore/Server+GRPC.swift

@@ -1383,13 +1383,14 @@ extension Initd: Com_Apple_Containerization_Sandbox_V3_SandboxContext.SimpleServ
             for containerID in containerIDs {
                 let container = try await state.get(container: containerID)
 
-                // Only fetch cgroup stats if needed
-                let cgStats: Cgroup2Stats?
-                if wantProcess || wantMemory || wantCPU || wantBlockIO {
-                    cgStats = try await container.stats()
-                } else {
-                    cgStats = nil
-                }
+                // Only read the cgroup stat groups that were requested.
+                var cgCategories: Cgroup2StatsCategory = []
+                if wantProcess { cgCategories.insert(.pids) }
+                if wantMemory { cgCategories.insert(.memory) }
+                if wantCPU { cgCategories.insert(.cpu) }
+                if wantBlockIO { cgCategories.insert(.io) }
+
+                let cgStats: Cgroup2Stats? = cgCategories.isEmpty ? nil : try await container.stats(cgCategories)
 
                 // Get network stats only if requested
                 var networkStats: [Com_Apple_Containerization_Sandbox_V3_NetworkStats] = []