vminitd: Add memory threshold monitoring (#427)

Add a small bit of logic to monitor if vminitd goes over a (somewhat
arbitrary) memory threshold. On average, when running one container it
seems to hover around 25MiB, so this is mostly to catch cases where it's
exceeding a limit we don't deem normal.

Author: dcantah

Sources/Containerization/Vminitd.swift

@@ -54,24 +54,14 @@ extension Vminitd: VirtualMachineAgent {
 
         try await setenv(key: "PATH", value: LinuxProcessConfiguration.defaultPath)
 
+        // Vminitd mounts /proc, /sys, /sys/fs/cgroup and /run automatically.
         let mounts: [ContainerizationOCI.Mount] = [
-            .init(type: "sysfs", source: "sysfs", destination: "/sys"),
             .init(type: "tmpfs", source: "tmpfs", destination: "/tmp"),
             .init(type: "devpts", source: "devpts", destination: "/dev/pts", options: ["gid=5", "mode=620", "ptmxmode=666"]),
-            .init(type: "cgroup2", source: "none", destination: "/sys/fs/cgroup"),
         ]
         for mount in mounts {
             try await self.mount(mount)
         }
-
-        // Setup root cg subtree_control.
-        let data = "+memory +pids +io +cpu +cpuset +hugetlb".data(using: .utf8)!
-        try await writeFile(
-            path: "/sys/fs/cgroup/cgroup.subtree_control",
-            data: data,
-            flags: .init(),
-            mode: 0
-        )
     }
 
     public func writeFile(path: String, data: Data, flags: WriteFileFlags, mode: UInt32) async throws {

vminitd/Sources/Cgroup/Cgroup2Manager.swift

@@ -234,6 +234,38 @@ package struct Cgroup2Manager: Sendable {
         }
     }
 
+    package func setMemoryHigh(bytes: UInt64) throws {
+        self.logger?.debug(
+            "setting memory.high",
+            metadata: [
+                "path": "\(self.path.path)",
+                "bytes": "\(bytes)",
+            ])
+
+        try Self.writeValue(
+            path: self.path,
+            value: String(bytes),
+            fileName: "memory.high"
+        )
+    }
+
+    package func getMemoryEvents() throws -> MemoryEvents {
+        let content = try readFileContent(fileName: "memory.events")
+        let values = parseKeyValuePairs(content)
+
+        return MemoryEvents(
+            low: values["low"] ?? 0,
+            high: values["high"] ?? 0,
+            max: values["max"] ?? 0,
+            oom: values["oom"] ?? 0,
+            oomKill: values["oom_kill"] ?? 0
+        )
+    }
+
+    package func getMemoryEventsPath() -> String {
+        self.path.appending(path: "memory.events").path
+    }
+
     package func kill() throws {
         try Self.writeValue(
             path: self.path,
@@ -608,6 +640,28 @@ package struct IOEntry: Sendable {
     }
 }
 
+package struct MemoryEvents: Sendable {
+    package var low: UInt64
+    package var high: UInt64
+    package var max: UInt64
+    package var oom: UInt64
+    package var oomKill: UInt64
+
+    package init(
+        low: UInt64 = 0,
+        high: UInt64 = 0,
+        max: UInt64 = 0,
+        oom: UInt64 = 0,
+        oomKill: UInt64 = 0
+    ) {
+        self.low = low
+        self.high = high
+        self.max = max
+        self.oom = oom
+        self.oomKill = oomKill
+    }
+}
+
 extension Cgroup2Manager {
     package enum Error: Swift.Error, CustomStringConvertible {
         case notCgroup

vminitd/Sources/vminitd/Application.swift

@@ -14,6 +14,7 @@
 // limitations under the License.
 //===----------------------------------------------------------------------===//
 
+import Cgroup
 import Containerization
 import ContainerizationError
 import ContainerizationOS
@@ -91,23 +92,86 @@ struct Application {
         CZ_set_sub_reaper()
         #endif
 
+        log.logLevel = .debug
+
         signal(SIGPIPE, SIG_IGN)
 
-        // Because the sysctl rpc wouldn't make sense if this didn't always exist, we
-        // ALWAYS mount /proc.
-        guard Musl.mount("proc", "/proc", "proc", 0, "") == 0 else {
-            log.error("failed to mount /proc")
-            exit(1)
-        }
-        guard Musl.mount("tmpfs", "/run", "tmpfs", 0, "") == 0 else {
-            log.error("failed to mount /run")
-            exit(1)
+        log.info("vminitd booting")
+
+        // Set of mounts necessary to be mounted prior to taking any RPCs.
+        // 1. /proc as the sysctl rpc wouldn't make sense if it wasn't there.
+        // 2. /run as that is where we store container state.
+        // 3. /sys as we need it for /sys/fs/cgroup
+        // 4. /sys/fs/cgroup to add the agent to a cgroup, as well as containers later.
+        let mounts = [
+            ContainerizationOS.Mount(
+                type: "proc",
+                source: "proc",
+                target: "/proc",
+                options: []
+            ),
+            ContainerizationOS.Mount(
+                type: "tmpfs",
+                source: "tmpfs",
+                target: "/run",
+                options: []
+            ),
+            ContainerizationOS.Mount(
+                type: "sysfs",
+                source: "sysfs",
+                target: "/sys",
+                options: []
+            ),
+            ContainerizationOS.Mount(
+                type: "cgroup2",
+                source: "none",
+                target: "/sys/fs/cgroup",
+                options: []
+            ),
+        ]
+
+        for mnt in mounts {
+            log.info("mounting \(mnt.target)")
+
+            try mnt.mount(createWithPerms: 0o755)
         }
         try Binfmt.mount()
 
-        log.logLevel = .debug
+        let cgManager = Cgroup2Manager(
+            group: URL(filePath: "/vminitd"),
+            logger: log
+        )
+        try cgManager.create()
+        try cgManager.toggleAllAvailableControllers(enable: true)
+
+        // Set memory.high threshold to 75 MiB
+        let threshold: UInt64 = 75 * 1024 * 1024
+        try cgManager.setMemoryHigh(bytes: threshold)
+        try cgManager.addProcess(pid: getpid())
+
+        let memoryMonitor = try MemoryMonitor(
+            cgroupManager: cgManager,
+            threshold: threshold,
+            logger: log
+        ) { [log] (currentUsage, highMark) in
+            log.warning(
+                "vminitd memory threshold exceeded",
+                metadata: [
+                    "threshold_bytes": "\(threshold)",
+                    "current_bytes": "\(currentUsage)",
+                    "high_events_total": "\(highMark)",
+                ])
+        }
+
+        let t = Thread { [log] in
+            do {
+                try memoryMonitor.run()
+            } catch {
+                log.error("memory monitor failed: \(error)")
+            }
+        }
+        t.start()
 
-        log.info("vminitd booting")
         let eg = MultiThreadedEventLoopGroup(numberOfThreads: System.coreCount)
         let server = Initd(log: log, group: eg)
 

vminitd/Sources/vminitd/MemoryMonitor.swift

@@ -0,0 +1,158 @@
+//===----------------------------------------------------------------------===//
+// Copyright © 2025 Apple Inc. and the Containerization 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.
+//===----------------------------------------------------------------------===//
+
+#if os(Linux)
+
+import Cgroup
+import Foundation
+import Logging
+
+#if canImport(Musl)
+import Musl
+#elseif canImport(Glibc)
+import Glibc
+#endif
+
+package final class MemoryMonitor: Sendable {
+    private static let inotifyEventSize = 0x10
+
+    private let cgroupManager: Cgroup2Manager
+    private let threshold: UInt64
+    private let logger: Logger
+    private let inotifyFd: Int32
+    private let watchDescriptor: Int32
+    private let onThresholdExceeded: @Sendable (UInt64, UInt64) -> Void
+
+    package init(
+        cgroupManager: Cgroup2Manager,
+        threshold: UInt64,
+        logger: Logger,
+        onThresholdExceeded: @escaping @Sendable (UInt64, UInt64) -> Void
+    ) throws {
+        self.cgroupManager = cgroupManager
+        self.threshold = threshold
+        self.logger = logger
+        self.onThresholdExceeded = onThresholdExceeded
+
+        let fd = inotify_init()
+        guard fd != -1 else {
+            throw Error.inotifyInit(errno: errno)
+        }
+        self.inotifyFd = fd
+
+        let eventsPath = cgroupManager.getMemoryEventsPath()
+        let wd = inotify_add_watch(
+            inotifyFd,
+            eventsPath,
+            UInt32(IN_MODIFY)
+        )
+        guard wd != -1 else {
+            close(fd)
+            throw Error.inotifyAddWatch(errno: errno, path: eventsPath)
+        }
+        self.watchDescriptor = wd
+    }
+
+    /// Run the monitoring loop. Call this from a dedicated thread.
+    /// This function blocks until an error occurs.
+    package func run() throws {
+        let eventsPath = cgroupManager.getMemoryEventsPath()
+
+        logger.info(
+            "Started memory monitoring",
+            metadata: [
+                "threshold_bytes": "\(threshold)",
+                "events_path": "\(eventsPath)",
+            ])
+
+        // Read initial state
+        var highCountMax: UInt64 = 0
+        do {
+            let events = try cgroupManager.getMemoryEvents()
+            highCountMax = events.high
+        } catch {
+            throw Error.readMemoryEvents(error: error)
+        }
+
+        let bufSize = Self.inotifyEventSize * 10
+        var buffer = [UInt8](repeating: 0, count: bufSize)
+        while true {
+            let bytesRead = buffer.withUnsafeMutableBytes { ptr in
+                read(inotifyFd, ptr.baseAddress!, bufSize)
+            }
+
+            if bytesRead < 0 {
+                if errno == EINTR {
+                    continue
+                }
+                throw Error.readFailed(errno: errno)
+            }
+
+            do {
+                let events = try cgroupManager.getMemoryEvents()
+
+                if events.high > highCountMax {
+                    highCountMax = events.high
+
+                    let stats = try cgroupManager.stats()
+                    let currentUsage = stats.memory?.usage ?? 0
+
+                    onThresholdExceeded(currentUsage, events.high)
+                }
+
+                if events.oom > 0 || events.oomKill > 0 {
+                    logger.error(
+                        "OOM events detected",
+                        metadata: [
+                            "oom_events": "\(events.oom)",
+                            "oom_kill_events": "\(events.oomKill)",
+                        ])
+                }
+            } catch {
+                throw Error.readMemoryEvents(error: error)
+            }
+        }
+    }
+
+    deinit {
+        inotify_rm_watch(inotifyFd, watchDescriptor)
+        close(inotifyFd)
+    }
+}
+
+extension MemoryMonitor {
+    package enum Error: Swift.Error, CustomStringConvertible {
+        case inotifyInit(errno: Int32)
+        case inotifyAddWatch(errno: Int32, path: String)
+        case readFailed(errno: Int32)
+        case readMemoryEvents(error: Swift.Error)
+
+        package var description: String {
+            switch self {
+            case .inotifyInit(let errno):
+                return "failed to initialize inotify: errno \(errno)"
+            case .inotifyAddWatch(let errno, let path):
+                return "failed to add inotify watch on \(path): errno \(errno)"
+            case .readFailed(let errno):
+                return "failed to read inotify events: errno \(errno)"
+            case .readMemoryEvents(let error):
+                return "failed to read memory events: \(error)"
+            }
+        }
+    }
+}
+
+#endif