Vminitd: Add pause command

vminitd/Sources/vminitd/Application.swift

@@ -14,183 +14,70 @@
 // limitations under the License.
 //===----------------------------------------------------------------------===//
 
-import Cgroup
-import Containerization
-import ContainerizationError
 import ContainerizationOS
 import Foundation
 import Logging
-import NIOCore
-import NIOPosix
-
-#if os(Linux)
-import Musl
-import LCShim
-#endif
 
 @main
 struct Application {
-    private static let foregroundEnvVar = "FOREGROUND"
-    private static let vsockPort = 1024
-    private static let standardErrorLock = NSLock()
-
-    private static func runInForeground(_ log: Logger) throws {
-        log.info("running vminitd under pid1")
-
-        var command = Command("/sbin/vminitd")
-        command.attrs = .init(setsid: true)
-        command.stdin = .standardInput
-        command.stdout = .standardOutput
-        command.stderr = .standardError
-        command.environment = ["\(foregroundEnvVar)=1"]
-
-        try command.start()
-        _ = try command.wait()
-    }
-
-    private static func adjustLimits() throws {
-        var limits = rlimit()
-        guard getrlimit(RLIMIT_NOFILE, &limits) == 0 else {
-            throw POSIXError(.init(rawValue: errno)!)
-        }
-        limits.rlim_cur = 65536
-        limits.rlim_max = 65536
-        guard setrlimit(RLIMIT_NOFILE, &limits) == 0 else {
-            throw POSIXError(.init(rawValue: errno)!)
-        }
-    }
-
-    @Sendable
-    private static func standardError(label: String) -> StreamLogHandler {
-        standardErrorLock.withLock {
-            StreamLogHandler.standardError(label: label)
+    static func main() async throws {
+        LoggingSystem.bootstrap(StreamLogHandler.standardError)
+
+        // Parse command line arguments
+        let args = CommandLine.arguments
+        let command = args.count > 1 ? args[1] : "init"
+
+        switch command {
+        case "pause":
+            let log = Logger(label: "pause")
+
+            log.info("Running pause command")
+            try PauseCommand.run(log: log)
+        case "init":
+            fallthrough
+        default:
+            let log = Logger(label: "vminitd")
+
+            log.info("Running init command")
+            try Self.mountProc(log: log)
+            try await InitCommand.run(log: log)
         }
     }
 
-    static func main() async throws {
-        LoggingSystem.bootstrap(standardError)
-        var log = Logger(label: "vminitd")
-
-        try adjustLimits()
-
-        // when running under debug mode, launch vminitd as a sub process of pid1
-        // so that we get a chance to collect better logs and errors before pid1 exists
-        // and the kernel panics.
-        #if DEBUG
-        let environment = ProcessInfo.processInfo.environment
-        let foreground = environment[Self.foregroundEnvVar]
-        log.info("checking for shim var \(foregroundEnvVar)=\(String(describing: foreground))")
-
-        if foreground == nil {
-            try runInForeground(log)
-            exit(0)
+    // Swift seems like it has some fun issues trying to spawn threads if /proc isn't around, so we
+    // do this before calling our first async function.
+    static func mountProc(log: Logger) throws {
+        // Is it already mounted (would only be true in debug builds where we re-exec ourselves)?
+        if isProcMounted() {
+            return
         }
 
-        // since we are not running as pid1 in this mode we must set ourselves
-        // as a subpreaper so that all child processes are reaped by us and not
-        // passed onto our parent.
-        CZ_set_sub_reaper()
-        #endif
+        log.info("mounting /proc")
 
-        log.logLevel = .debug
-
-        signal(SIGPIPE, SIG_IGN)
-
-        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()
-
-        let cgManager = Cgroup2Manager(
-            group: URL(filePath: "/vminitd"),
-            logger: log
+        let mnt = ContainerizationOS.Mount(
+            type: "proc",
+            source: "proc",
+            target: "/proc",
+            options: []
         )
-        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())
+        try mnt.mount(createWithPerms: 0o755)
+    }
 
-        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)",
-                ])
+    static func isProcMounted() -> Bool {
+        guard let data = try? String(contentsOfFile: "/proc/mounts", encoding: .utf8) else {
+            return false
         }
 
-        let t = Thread { [log] in
-            do {
-                try memoryMonitor.run()
-            } catch {
-                log.error("memory monitor failed: \(error)")
+        for line in data.split(separator: "\n") {
+            let fields = line.split(separator: " ")
+            if fields.count >= 2 {
+                let mountPoint = String(fields[1])
+                if mountPoint == "/proc" {
+                    return true
+                }
             }
         }
-        t.start()
-
-        let eg = MultiThreadedEventLoopGroup(numberOfThreads: System.coreCount)
-        let server = Initd(log: log, group: eg)
-
-        do {
-            log.info("serving vminitd API")
-            try await server.serve(port: vsockPort)
-            log.info("vminitd API returned, syncing filesystems")
 
-            #if os(Linux)
-            Musl.sync()
-            #endif
-        } catch {
-            log.error("vminitd boot error \(error)")
-
-            #if os(Linux)
-            Musl.sync()
-            #endif
-
-            exit(1)
-        }
+        return false
     }
 }