fix(plugins): bound discovery when apiserver stalls

Package.swift

@@ -150,6 +150,7 @@ let package = Package(
         .testTarget(
             name: "ContainerCommandsTests",
             dependencies: [
+                "ContainerAPIClient",
                 "ContainerCommands",
                 "ContainerResource",
             ]
@@ -474,6 +475,12 @@ let package = Package(
                 "CAuditToken",
             ]
         ),
+        .testTarget(
+            name: "ContainerXPCTests",
+            dependencies: [
+                "ContainerXPC"
+            ]
+        ),
         .target(
             name: "ContainerOS",
             dependencies: [

Sources/ContainerCommands/Application.swift

@@ -36,6 +36,8 @@ private nonisolated(unsafe) var bootstrapLogger = {
     return log
 }()
 
+typealias APIServerHealthCheck = @Sendable (Duration) async throws -> SystemHealth
+
 public struct Application: AsyncLoggableCommand {
     @OptionGroup
     public var logOptions: Flags.Logging
@@ -127,9 +129,8 @@ public struct Application: AsyncLoggableCommand {
         } catch {
             // --help/-h on the root command (e.g. `container --help`) is intercepted
             // by ArgumentParser and lands here.
-            let containsHelp = fullArgs.contains("-h") || fullArgs.contains("--help")
-            if fullArgs.count <= 2 && containsHelp {
-                let pluginLoader = try? await createPluginLoader()
+            if isRootHelpRequest(fullArgs) {
+                let pluginLoader = await pluginLoaderForHelp()
                 await Self.printModifiedHelpText(pluginLoader: pluginLoader)
                 return
             }
@@ -144,6 +145,18 @@ public struct Application: AsyncLoggableCommand {
     }
 
     public static func createPluginLoader() async throws -> PluginLoader {
+        try await createPluginLoader(healthTimeout: .seconds(10), wallTimeout: .seconds(10), healthCheck: defaultAPIServerHealthCheck(timeout:))
+    }
+
+    static func pluginLoaderForHelp(healthCheck: @escaping APIServerHealthCheck = defaultAPIServerHealthCheck(timeout:)) async -> PluginLoader? {
+        try? await createPluginLoader(healthTimeout: .seconds(1), wallTimeout: .seconds(1), healthCheck: healthCheck)
+    }
+
+    static func isRootHelpRequest(_ arguments: [String]) -> Bool {
+        arguments.count <= 2 && (arguments.contains("-h") || arguments.contains("--help"))
+    }
+
+    private static func createPluginLoader(healthTimeout: Duration, wallTimeout: Duration, healthCheck: @escaping APIServerHealthCheck) async throws -> PluginLoader {
         let installRootPath = CommandLine.executablePath
             .removingLastComponent()
             .removingLastComponent()
@@ -175,9 +188,7 @@ public struct Application: AsyncLoggableCommand {
             AppBundlePluginFactory(logger: bootstrapLogger),
         ]
 
-        guard let systemHealth = try? await ClientHealthCheck.ping(timeout: .seconds(10)) else {
-            throw ContainerizationError(.timeout, message: "unable to retrieve application data root from API server")
-        }
+        let systemHealth = try await apiServerHealth(healthTimeout: healthTimeout, wallTimeout: wallTimeout, healthCheck: healthCheck)
         return try PluginLoader(
             appRoot: systemHealth.appRoot,
             installRoot: systemHealth.installRoot,
@@ -188,6 +199,43 @@ public struct Application: AsyncLoggableCommand {
         )
     }
 
+    private static func defaultAPIServerHealthCheck(timeout: Duration) async throws -> SystemHealth {
+        try await ClientHealthCheck.ping(timeout: timeout)
+    }
+
+    static func apiServerHealth(
+        healthTimeout: Duration,
+        wallTimeout: Duration,
+        healthCheck: @escaping APIServerHealthCheck = defaultAPIServerHealthCheck(timeout:)
+    ) async throws -> SystemHealth {
+        try await withCheckedThrowingContinuation { continuation in
+            let result = HealthCheckResult(continuation)
+
+            let healthTask = Task.detached {
+                do {
+                    let health = try await healthCheck(healthTimeout)
+                    await result.resume(.success(health))
+                } catch {
+                    await result.resume(.failure(error))
+                }
+            }
+
+            let timeoutTask = Task.detached {
+                try? await Task.sleep(for: wallTimeout)
+                await result.resume(
+                    .failure(
+                        ContainerizationError(
+                            .timeout,
+                            message: "unable to retrieve application data root from API server"
+                        )))
+            }
+
+            Task {
+                await result.setTasks(healthTask: healthTask, timeoutTask: timeoutTask)
+            }
+        }
+    }
+
     /// Load the system configuration using `appRoot` / `installRoot` reported by the
     /// daemon. `container system start` MUST have previously been run to start the daemon.
     public static func loadContainerSystemConfig() async throws -> ContainerSystemConfig {
@@ -254,6 +302,41 @@ public struct Application: AsyncLoggableCommand {
     }
 }
 
+private actor HealthCheckResult {
+    private let continuation: CheckedContinuation<SystemHealth, any Error>
+    private var healthTask: Task<Void, Never>?
+    private var timeoutTask: Task<Void, Never>?
+    private var resumed = false
+
+    init(_ continuation: CheckedContinuation<SystemHealth, any Error>) {
+        self.continuation = continuation
+    }
+
+    func setTasks(healthTask: Task<Void, Never>, timeoutTask: Task<Void, Never>) {
+        if resumed {
+            healthTask.cancel()
+            timeoutTask.cancel()
+            return
+        }
+        self.healthTask = healthTask
+        self.timeoutTask = timeoutTask
+    }
+
+    func resume(_ result: Result<SystemHealth, any Error>) {
+        guard !resumed else {
+            return
+        }
+        resumed = true
+        let healthTask = healthTask
+        let timeoutTask = timeoutTask
+        self.healthTask = nil
+        self.timeoutTask = nil
+        healthTask?.cancel()
+        timeoutTask?.cancel()
+        continuation.resume(with: result)
+    }
+}
+
 extension Application {
     // Because we support plugins, we need to modify the help text to display
     // any if we found some.

Sources/ContainerCommands/DefaultCommand.swift

@@ -34,9 +34,8 @@ struct DefaultCommand: AsyncLoggableCommand {
     var remaining: [String] = []
 
     func run() async throws {
-        // See if we have a possible plugin command.
-        let pluginLoader = try? await Application.createPluginLoader()
         guard let command = remaining.first else {
+            let pluginLoader = await Application.pluginLoaderForHelp()
             await Application.printModifiedHelpText(pluginLoader: pluginLoader)
             return
         }
@@ -66,6 +65,9 @@ struct DefaultCommand: AsyncLoggableCommand {
             .map { $0.appendingPathComponent(command).path(percentEncoded: false) }
             .joined(separator: "\n  - ")
 
+        // See if we have a possible plugin command.
+        let pluginLoader = try? await Application.createPluginLoader()
+
         // If plugin loader couldn't be created, the system/APIServer likely isn't running.
         if pluginLoader == nil {
             throw ValidationError(

Sources/ContainerCommands/HelpCommand.swift

@@ -31,7 +31,7 @@ struct HelpCommand: AsyncLoggableCommand {
 
     func run() async throws {
         if subcommandPath.isEmpty {
-            let pluginLoader = try? await Application.createPluginLoader()
+            let pluginLoader = await Application.pluginLoaderForHelp()
             await Application.printModifiedHelpText(pluginLoader: pluginLoader)
             return
         }

Sources/ContainerXPC/XPCClient.swift

@@ -98,42 +98,31 @@ extension XPCClient {
     /// Send the provided message to the service.
     @discardableResult
     public func send(_ message: XPCMessage, responseTimeout: Duration? = nil) async throws -> XPCMessage {
-        try await withThrowingTaskGroup(of: XPCMessage.self, returning: XPCMessage.self) { group in
+        try await withCheckedThrowingContinuation { continuation in
+            let result = XPCResponseResult(continuation)
+
             if let responseTimeout {
-                group.addTask {
-                    try await Task.sleep(for: responseTimeout)
-                    let route = message.string(key: XPCMessage.routeKey) ?? "nil"
-                    throw ContainerizationError(
-                        .internalError,
-                        message: "XPC timeout for request to \(self.service)/\(route)"
-                    )
+                let route = message.string(key: XPCMessage.routeKey) ?? "nil"
+                let timeoutTask = Task.detached {
+                    try? await Task.sleep(for: responseTimeout)
+                    _ = result.resume(
+                        .failure(
+                            ContainerizationError(
+                                .timeout,
+                                message: "XPC timeout for request to \(self.service)/\(route)"
+                            )))
                 }
+                result.setTimeoutTask(timeoutTask)
             }
 
-            group.addTask {
-                try await withCheckedThrowingContinuation { cont in
-                    xpc_connection_send_message_with_reply(self.connection, message.underlying, nil) { reply in
-                        do {
-                            let message = try self.parseReply(reply)
-                            cont.resume(returning: message)
-                        } catch {
-                            cont.resume(throwing: error)
-                        }
-                    }
+            xpc_connection_send_message_with_reply(self.connection, message.underlying, nil) { reply in
+                do {
+                    let message = try self.parseReply(reply)
+                    _ = result.resume(.success(message))
+                } catch {
+                    _ = result.resume(.failure(error))
                 }
             }
-
-            let response = try await group.next()
-            // once one task has finished, cancel the rest.
-            group.cancelAll()
-            // we don't really care about the second error here
-            // as it's most likely a `CancellationError`.
-            try? await group.waitForAll()
-
-            guard let response else {
-                throw ContainerizationError(.invalidState, message: "failed to receive XPC response")
-            }
-            return response
         }
     }
 
@@ -159,4 +148,42 @@ extension XPCClient {
     }
 }
 
+private final class XPCResponseResult: @unchecked Sendable {
+    private let lock = NSLock()
+    private let continuation: CheckedContinuation<XPCMessage, any Error>
+    private var timeoutTask: Task<Void, Never>?
+    private var resumed = false
+
+    init(_ continuation: CheckedContinuation<XPCMessage, any Error>) {
+        self.continuation = continuation
+    }
+
+    func setTimeoutTask(_ task: Task<Void, Never>) {
+        lock.lock()
+        if resumed {
+            lock.unlock()
+            task.cancel()
+            return
+        }
+        timeoutTask = task
+        lock.unlock()
+    }
+
+    func resume(_ result: Result<XPCMessage, any Error>) -> Bool {
+        lock.lock()
+        guard !resumed else {
+            lock.unlock()
+            return false
+        }
+        resumed = true
+        let timeoutTask = timeoutTask
+        self.timeoutTask = nil
+        lock.unlock()
+
+        timeoutTask?.cancel()
+        continuation.resume(with: result)
+        return true
+    }
+}
+
 #endif