vminitd: Add ability to set log level (#531)

This adds a small flag named --log-level to be able to set the log level
of vminitd and the pause command. I imagine this won't be the last time
we'll have to add flags to vminitd so I went ahead and converted vminitd
to be a Argument Parser based binary as well.

vminitd/Package.swift

@@ -48,6 +48,7 @@ let package = Package(
         .executableTarget(
             name: "vminitd",
             dependencies: [
+                .product(name: "ArgumentParser", package: "swift-argument-parser"),
                 .product(name: "Logging", package: "swift-log"),
                 .product(name: "Containerization", package: "containerization"),
                 .product(name: "ContainerizationNetlink", package: "containerization"),

vminitd/Sources/vminitd/Application.swift

@@ -14,46 +14,44 @@
 // limitations under the License.
 //===----------------------------------------------------------------------===//
 
+import ArgumentParser
 import ContainerizationOS
 import Foundation
 import Logging
 
 @main
-struct Application {
+struct Application: AsyncParsableCommand {
+    static let configuration = CommandConfiguration(
+        commandName: "vminitd",
+        abstract: "Virtual machine init daemon",
+        version: "0.1.0",
+        subcommands: [
+            InitCommand.self,
+            PauseCommand.self,
+        ],
+        defaultSubcommand: InitCommand.self
+    )
+
     static func main() async throws {
-        LoggingSystem.bootstrap(StreamLogHandler.standardError)
+        // Swift has issues spawning threads if /proc isn't mounted,
+        // so we do this synchronously before any async code runs.
+        try mountProc()
 
-        // 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)
+        var command = try parseAsRoot()
+        if let asyncCommand = command as? AsyncParsableCommand {
+            nonisolated(unsafe) var unsafeCommand = asyncCommand
+            try await unsafeCommand.run()
+        } else {
+            try command.run()
         }
     }
 
-    // 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 {
+    private static func mountProc() throws {
         // Is it already mounted (would only be true in debug builds where we re-exec ourselves)?
         if isProcMounted() {
             return
         }
 
-        log.info("mounting /proc")
-
         let mnt = ContainerizationOS.Mount(
             type: "proc",
             source: "proc",
@@ -63,7 +61,7 @@ struct Application {
         try mnt.mount(createWithPerms: 0o755)
     }
 
-    static func isProcMounted() -> Bool {
+    private static func isProcMounted() -> Bool {
         guard let data = try? String(contentsOfFile: "/proc/mounts", encoding: .utf8) else {
             return false
         }
@@ -81,3 +79,36 @@ struct Application {
         return false
     }
 }
+
+struct LogLevelOption: ParsableArguments {
+    @Option(name: .long, help: "Set the log level (trace, debug, info, notice, warning, error, critical)")
+    var logLevel: String = "info"
+
+    func resolvedLogLevel() -> Logger.Level {
+        switch logLevel.lowercased() {
+        case "trace":
+            return .trace
+        case "debug":
+            return .debug
+        case "info":
+            return .info
+        case "notice":
+            return .notice
+        case "warning":
+            return .warning
+        case "error":
+            return .error
+        case "critical":
+            return .critical
+        default:
+            return .info
+        }
+    }
+}
+
+func makeLogger(label: String, level: Logger.Level) -> Logger {
+    LoggingSystem.bootstrap(StreamLogHandler.standardError)
+    var log = Logger(label: label)
+    log.logLevel = level
+    return log
+}

vminitd/Sources/vminitd/InitCommand.swift

@@ -14,6 +14,7 @@
 // limitations under the License.
 //===----------------------------------------------------------------------===//
 
+import ArgumentParser
 import Cgroup
 import Containerization
 import ContainerizationError
@@ -28,12 +29,19 @@ import Musl
 import LCShim
 #endif
 
-struct InitCommand {
+struct InitCommand: AsyncParsableCommand {
+    static let configuration = CommandConfiguration(
+        commandName: "init",
+        abstract: "Run the init daemon"
+    )
+
     private static let foregroundEnvVar = "FOREGROUND"
     private static let vsockPort = 1024
 
-    static func run(log: Logger) async throws {
-        var log = log
+    @OptionGroup var options: LogLevelOption
+
+    mutating func run() async throws {
+        let log = makeLogger(label: "vminitd", level: options.resolvedLogLevel())
         try Self.adjustLimits(log)
 
         // when running under debug mode, launch vminitd as a sub process of pid1
@@ -43,11 +51,11 @@ struct InitCommand {
         log.info("DEBUG mode active, checking FOREGROUND env var")
         let environment = ProcessInfo.processInfo.environment
         let foreground = environment[Self.foregroundEnvVar]
-        log.info("checking for shim var \(foregroundEnvVar)=\(String(describing: foreground))")
+        log.info("checking for shim var \(Self.foregroundEnvVar)=\(String(describing: foreground))")
 
         if foreground == nil {
-            try runInForeground(log)
-            exit(0)
+            try Self.runInForeground(log, logLevel: options.logLevel)
+            _exit(0)
         }
 
         log.info("FOREGROUND is set, running as subprocess, setting subreaper")
@@ -57,8 +65,6 @@ struct InitCommand {
         CZ_set_sub_reaper()
         #endif
 
-        log.logLevel = .debug
-
         signal(SIGPIPE, SIG_IGN)
 
         log.info("vminitd booting")
@@ -137,7 +143,7 @@ struct InitCommand {
 
         do {
             log.info("serving vminitd API")
-            try await server.serve(port: vsockPort)
+            try await server.serve(port: Self.vsockPort)
             log.info("vminitd API returned, syncing filesystems")
 
             #if os(Linux)
@@ -150,14 +156,14 @@ struct InitCommand {
             Musl.sync()
             #endif
 
-            exit(1)
+            _exit(1)
         }
     }
 
-    private static func runInForeground(_ log: Logger) throws {
+    private static func runInForeground(_ log: Logger, logLevel: String) throws {
         log.info("running vminitd under pid1")
 
-        var command = Command("/sbin/vminitd")
+        var command = Command("/sbin/vminitd", arguments: ["init", "--log-level", logLevel])
         command.attrs = .init(setsid: true)
         command.stdin = .standardInput
         command.stdout = .standardOutput

vminitd/Sources/vminitd/PauseCommand.swift

@@ -14,12 +14,22 @@
 // limitations under the License.
 //===----------------------------------------------------------------------===//
 
+import ArgumentParser
 import Dispatch
 import Logging
 import Musl
 
-struct PauseCommand {
-    static func run(log: Logger) throws {
+struct PauseCommand: ParsableCommand {
+    static let configuration = CommandConfiguration(
+        commandName: "pause",
+        abstract: "Run the pause container"
+    )
+
+    @OptionGroup var options: LogLevelOption
+
+    mutating func run() throws {
+        let log = makeLogger(label: "pause", level: options.resolvedLogLevel())
+
         if getpid() != 1 {
             log.warning("pause should be the first process")
         }