Fix exec hang and stop timeout

Sources/APIServer/APIServer+Start.swift

@@ -298,7 +298,7 @@ extension APIServer {
             routes[XPCRoute.containerDial] = XPCServer.route(harness.dial)
             routes[XPCRoute.containerStop] = XPCServer.route(harness.stop)
             routes[XPCRoute.containerStartProcess] = XPCServer.route(harness.startProcess)
-            routes[XPCRoute.containerCreateProcess] = XPCServer.route(harness.createProcess)
+            routes[XPCRoute.containerCreateProcess] = harness.createProcess
             routes[XPCRoute.containerResize] = XPCServer.route(harness.resize)
             routes[XPCRoute.containerWait] = XPCServer.route(harness.wait)
             routes[XPCRoute.containerKill] = XPCServer.route(harness.kill)

Sources/Services/ContainerAPIService/Server/Containers/ContainersHarness.swift

@@ -203,7 +203,7 @@ public struct ContainersHarness: Sendable {
     }
 
     @Sendable
-    public func createProcess(_ message: XPCMessage) async throws -> XPCMessage {
+    public func createProcess(_ message: XPCMessage, session: XPCServerSession) async throws -> XPCMessage {
         let id = message.string(key: .id)
         guard let id else {
             throw ContainerizationError(
@@ -228,6 +228,11 @@ public struct ContainersHarness: Sendable {
             stdio: stdio
         )
 
+        await session.onDisconnect { [service] in
+            // Clean up the process if the client connection drops unexpectedly
+            try? await service.kill(id: id, processID: processID, signal: "SIGKILL")
+        }
+
         return message.reply()
     }
 

Sources/Services/RuntimeLinux/Server/RuntimeService.swift

@@ -518,31 +518,52 @@ public actor RuntimeService {
         let signal = try Signal(stopOptions.signal ?? "SIGTERM")
         let timeout: Duration = .seconds(stopOptions.timeoutInSeconds)
 
-        return try await self.lock.withLock { _ in
-            switch await self.state {
-            case .running, .booted:
-                await self.setState(.stopping)
-
-                let ctr = try await self.getContainer()
-                let exitStatus = try await self.gracefulStopContainer(
-                    ctr.container,
-                    signal: signal,
-                    timeout: timeout
-                )
+        return try await withThrowingTaskGroup(of: XPCMessage.self) { group in
+            group.addTask {
+                return try await self.lock.withLock { _ in
+                    switch await self.state {
+                    case .running, .booted:
+                        await self.setState(.stopping)
+
+                        let ctr = try await self.getContainer()
+                        let exitStatus = try await self.gracefulStopContainer(
+                            ctr.container,
+                            signal: signal,
+                            timeout: timeout
+                        )
 
-                do {
-                    if case .stopped = await self.state {
-                        return message.reply()
+                        do {
+                            if case .stopped = await self.state {
+                                return message.reply()
+                            }
+                            try await self.cleanUpContainer(containerInfo: ctr, exitStatus: exitStatus)
+                        } catch {
+                            self.log.error("failed to clean up container", metadata: ["error": "\(error)"])
+                        }
+                        await self.setState(.stopped)
+                    default:
+                        break
                     }
-                    try await self.cleanUpContainer(containerInfo: ctr, exitStatus: exitStatus)
-                } catch {
-                    self.log.error("failed to clean up container", metadata: ["error": "\(error)"])
+                    return message.reply()
                 }
-                await self.setState(.stopped)
-            default:
-                break
             }
-            return message.reply()
+
+            group.addTask {
+                try await Task.sleep(for: timeout)
+                self.log.warning("stop operation timed out; force killing the container")
+                if let ctr = try? await self.getContainer() {
+                    try? await ctr.container.kill(.kill)
+                    try? await self.cleanUpContainer(containerInfo: ctr, exitStatus: ExitStatus(exitCode: 137))
+                    await self.setState(.stopped)
+                }
+                return message.reply()
+            }
+
+            guard let result = try await group.next() else {
+                throw ContainerizationError(.internalError, message: "stop failed")
+            }
+            group.cancelAll()
+            return result
         }
     }
 
@@ -562,7 +583,7 @@ public actor RuntimeService {
         let id = try message.id()
         let signal = try Signal(message.signal())
 
-        try await self.lock.withLock { [self] _ in
+        let skipWait: Bool = try await self.lock.withLock { [self] _ in
             switch await self.state {
             case .running:
                 let ctr = try await getContainer()
@@ -574,11 +595,30 @@ public actor RuntimeService {
                     guard let proc = processInfo.process else {
                         throw ContainerizationError(.invalidState, message: "process \(id) not started")
                     }
-                    try await proc.kill(signal)
-                    return
+                    do {
+                        try await proc.kill(signal)
+                    } catch {
+                        let errStr = String(describing: error)
+                        if errStr.contains("clientIsStopped") {
+                            return true
+                        } else {
+                            throw error
+                        }
+                    }
+                    return false
                 }
 
-                try await ctr.container.kill(signal)
+                do {
+                    try await ctr.container.kill(signal)
+                } catch {
+                    let errStr = String(describing: error)
+                    if errStr.contains("clientIsStopped") {
+                        return true
+                    } else {
+                        throw error
+                    }
+                }
+                return false
             default:
                 throw ContainerizationError(
                     .invalidState,
@@ -587,9 +627,14 @@ public actor RuntimeService {
             }
         }
 
+        if skipWait {
+            self.log.debug("Kill ignored because agent is stopped (process likely already exited)", metadata: ["id": "\(id)"])
+            self.releaseWaiters(for: id, status: ExitStatus(exitCode: 255))
+        }
+
         // SIGKILL is guaranteed by the kernel to terminate the target, so block
         // until we observe the exit.
-        if signal == .kill {
+        if signal == .kill && !skipWait {
             _ = await withCheckedContinuation { cc in
                 self.waitForExit(id: id, cont: cc)
             }