From 80f1273f3f6c9a1eeb00d2692220f5fb8f948a3e Mon Sep 17 00:00:00 2001
From: Danny Canter <danny_canter@apple.com>
Date: Fri, 29 May 2026 18:20:03 -0700
Subject: [PATCH] Cgroup2: Add ability to filter stats

We funnily enough filter on the host, send the filters through
to the guest, but we still collect all stats and then just drop
whatever we didn't ask for which is very wasteful. This adds a
change to just not collect what we don't ask for.
---
 vminitd/Sources/Cgroup/Cgroup2Manager.swift   | 33 ++++++++++++-------
 .../VminitdCore/ManagedContainer.swift        |  4 +--
 .../Sources/VminitdCore/MemoryMonitor.swift   |  2 +-
 vminitd/Sources/VminitdCore/Server+GRPC.swift | 15 +++++----
 4 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/vminitd/Sources/Cgroup/Cgroup2Manager.swift b/vminitd/Sources/Cgroup/Cgroup2Manager.swift
index bd5f6b46..353b45f8 100644
--- a/vminitd/Sources/Cgroup/Cgroup2Manager.swift
+++ b/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?
diff --git a/vminitd/Sources/VminitdCore/ManagedContainer.swift b/vminitd/Sources/VminitdCore/ManagedContainer.swift
index 9efa5c40..545046a8 100644
--- a/vminitd/Sources/VminitdCore/ManagedContainer.swift
+++ b/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 {
diff --git a/vminitd/Sources/VminitdCore/MemoryMonitor.swift b/vminitd/Sources/VminitdCore/MemoryMonitor.swift
index 7a6f9fc0..d6dcd1c5 100644
--- a/vminitd/Sources/VminitdCore/MemoryMonitor.swift
+++ b/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)
diff --git a/vminitd/Sources/VminitdCore/Server+GRPC.swift b/vminitd/Sources/VminitdCore/Server+GRPC.swift
index 15e8aa56..d65f4519 100644
--- a/vminitd/Sources/VminitdCore/Server+GRPC.swift
+++ b/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] = []