Fix race condition on SandboxService.waiters (#1289)

This PR fixes #1277.

`SandboxService.waiters` had a consistency issue (not exactly race).
`SandboxService.wait` XPC can be executed on arbitrary `id`, and it will
hang forever if no other handler resumes it. Without knowing this
internal, the high level entity can run into this issue, and deadlock.

This PR simplifies the mental model: **`SandboxService.waiters[id]:
ExitWaiter(continuations, exitCode)` can only be in three states: i)
non-existing, ii) existing with nil `exitCode`, and iii) existing with
concrete `exitCode`.**

**If it is non-existing, no handler has been registered to resume it
later. If existing with nil `exitCode`, It is guaranteed the registered
`continuations` will be resumed later with a concrete `exitCode`.
Finally, if already a concrete `exitCode`, a handler has been
registered, and already resumed (with that `exitCode`).**

Thus, `SandboxService.wait` should return immediately if `waiters[id]`
is non-existing or existing with a concrete `exitCode` (as no handler
will resume it later). It should only block when `waiters[id]` is
existing with nil `exitCode` as it is guaranteed to be resumed later. By
doing so, we can guarantee there is no deadlock at all.

For that this PR does followings:
1. Introduce `ExitMonitor` class to updates `continuations` and
`exitCode` all together atomically. Initially, `state` variable saved
the `exitCode`, but it cannot be tied with `continuations` as they are
protected by different primitives (i.e., lock and actor).
2. Gather `waiters` related operations into a single actor method,
guaranteeing those are performed atomically under actor
protection---i.e., we actually don't need Mutex here.
3. Ensure initialized `waiters` are released (i.e., resumed) later
(under any possible circumstances).
4. Move `process.wait` after `process.start` in `io.handleProcess` to
run `SandboxService.wait` only after the `waiters[id]` is initialized.

By doing fourth step, we can guarantee `SandboxService.wait` can meet
only one of two following `ExitMonitor` state: i) existing with nil
`exitCode`, or ii) existing with concrete `exitCode` (in case the
process exited too early). In both cases, `exitCode` is preserved and
returned.

## Type of Change
- [X] Bug fix
- [ ] New feature  
- [ ] Breaking change
- [ ] Documentation update

## Motivation and Context
[Why is this change needed?]

## Testing
- [X] Tested locally
- [ ] Added/updated tests
- [ ] Added/updated docs

Author: JaewonHur

Sources/Services/ContainerAPIService/Client/ProcessIO.swift

@@ -149,6 +149,9 @@ public struct ProcessIO: Sendable {
     public func handleProcess(process: ClientProcess, log: Logger) async throws -> Int32 {
         let signals = AsyncSignalHandler.create(notify: Self.signalSet)
         return try await withThrowingTaskGroup(of: Int32?.self, returning: Int32.self) { group in
+            try await process.start()
+            try closeAfterStart()
+
             let waitAdded = group.addTaskUnlessCancelled {
                 let code = try await process.wait()
                 try await wait()
@@ -160,9 +163,6 @@ public struct ProcessIO: Sendable {
                 return -1
             }
 
-            try await process.start()
-            try closeAfterStart()
-
             if let current = console {
                 let size = try current.size
                 // It's supremely possible the process could've exited already. We shouldn't treat

Sources/Services/ContainerSandboxService/Server/SandboxService.swift

@@ -44,7 +44,7 @@ public actor SandboxService {
     private var container: ContainerInfo?
     private let monitor: ExitMonitor
     private let eventLoopGroup: any EventLoopGroup
-    private let waiters: Mutex<[String: [CheckedContinuation<ExitStatus, Never>]?]> = Mutex([:])
+    private var waiters: [String: ExitWaiter] = [:]
     private let lock: AsyncLock = AsyncLock()
     private let log: Logging.Logger
     private var state: State = .created
@@ -54,6 +54,27 @@ public actor SandboxService {
     private static let sshAuthSocketGuestPath = "/run/host-services/ssh-auth.sock"
     private static let sshAuthSocketEnvVar = "SSH_AUTH_SOCK"
 
+    class ExitWaiter {
+        public var exitCode: Int32? = nil
+        public var continuations: [CheckedContinuation<ExitStatus, Never>] = []
+
+        public func register(_ cc: CheckedContinuation<ExitStatus, Never>) {
+            continuations.append(cc)
+        }
+
+        public func doExit(code: Int32) {
+            for cc in continuations {
+                cc.resume(returning: ExitStatus(exitCode: code))
+            }
+
+            exitCode = code
+        }
+
+        public func exited() -> Bool {
+            exitCode != nil
+        }
+    }
+
     private static func sshAuthSocketHostUrl(config: ContainerConfiguration) -> URL? {
         if config.ssh, let sshSocket = Foundation.ProcessInfo.processInfo.environment[Self.sshAuthSocketEnvVar] {
             return URL(fileURLWithPath: sshSocket)
@@ -225,6 +246,8 @@ public actor SandboxService {
 
             do {
                 try await container.create()
+
+                try await self.initializeWaiters(for: id)
                 try await self.monitor.registerProcess(id: config.id, onExit: self.onContainerExit)
                 if !container.interfaces.isEmpty {
                     try await self.startSocketForwarders(attachment: attachments[0], publishedPorts: config.publishedPorts)
@@ -233,7 +256,7 @@ public actor SandboxService {
             } catch {
                 do {
                     try await self.cleanUpContainer(containerInfo: ctrInfo)
-                    await self.setState(.stopped(nil))
+                    await self.setState(.stopped)
                 } catch {
                     self.log.error("failed to clean up container", metadata: ["error": "\(error)"])
                 }
@@ -320,7 +343,7 @@ public actor SandboxService {
 
         return try await self.lock.withLock { _ in
             switch await self.state {
-            case .created, .stopped(_), .stopping:
+            case .created, .stopped, .stopping:
                 await self.setState(.shuttingDown)
 
             default:
@@ -360,27 +383,27 @@ public actor SandboxService {
 
                 try await self.addNewProcess(id, config, stdio)
 
-                try await self.monitor.registerProcess(
-                    id: id,
-                    onExit: { id, exitStatus in
-                        guard let process = await self.processes[id]?.process else {
-                            throw ContainerizationError(
-                                .invalidState,
-                                message: "ProcessInfo missing for process \(id)"
-                            )
-                        }
-                        self.waiters.withLock {
-                            if let waiters = $0[id], let waiters {
-                                for cc in waiters {
-                                    cc.resume(returning: exitStatus)
-                                }
+                try await self.initializeWaiters(for: id)
+                do {
+                    try await self.monitor.registerProcess(
+                        id: id,
+                        onExit: { id, exitStatus in
+                            await self.releaseWaiters(for: id, status: exitStatus)
+
+                            guard let process = await self.processes[id]?.process else {
+                                throw ContainerizationError(
+                                    .invalidState,
+                                    message: "ProcessInfo missing for process \(id)"
+                                )
                             }
+                            try await process.delete()
+                            try await self.setProcessState(id: id, state: .stopped)
                         }
-                        await self.removeWaiters(for: id)
-                        try await process.delete()
-                        try await self.setProcessState(id: id, state: .stopped(exitStatus.exitCode))
-                    }
-                )
+                    )
+                } catch {
+                    await self.releaseWaiters(for: id, status: ExitStatus(exitCode: -1))
+                    throw error
+                }
 
                 return message.reply()
             default:
@@ -410,7 +433,7 @@ public actor SandboxService {
         var cs: ContainerSnapshot?
 
         switch state {
-        case .created, .stopped(_), .booted, .shuttingDown:
+        case .created, .stopped, .booted, .shuttingDown:
             status = .stopped
         case .stopping:
             status = .stopping
@@ -464,14 +487,14 @@ public actor SandboxService {
                 )
 
                 do {
-                    if case .stopped(_) = await self.state {
+                    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(exitStatus.exitCode))
+                await self.setState(.stopped)
             default:
                 break
             }
@@ -596,39 +619,11 @@ public actor SandboxService {
             throw ContainerizationError(.invalidArgument, message: "missing id in wait xpc message")
         }
 
-        let cachedCode: Int32? = try await self.lock.withLock { _ in
-            let ctrInfo = try await self.getContainer()
-            let ctr = ctrInfo.container
-            if id == ctr.id {
-                switch await self.state {
-                case .stopped(let code):
-                    return code
-                default:
-                    break
-                }
-            } else {
-                guard let processInfo = await self.processes[id] else {
-                    throw ContainerizationError(.notFound, message: "process with id \(id)")
-                }
-                switch processInfo.state {
-                case .stopped(let code):
-                    return code
-                default:
-                    break
-                }
-            }
-            return nil
-        }
-        if let cachedCode {
-            let reply = message.reply()
-            reply.set(key: SandboxKeys.exitCode.rawValue, value: Int64(cachedCode))
-            return reply
-        }
-
         let exitStatus = await withCheckedContinuation { cc in
             // Is this safe since we are in an actor? :(
-            if !self.addWaiter(id: id, cont: cc) {
-                cc.resume(returning: ExitStatus(exitCode: -1))
+            let (added, exitCode) = self.addWaiter(id: id, cont: cc)
+            if !added {
+                cc.resume(returning: ExitStatus(exitCode: exitCode ?? -1))
             }
         }
         let reply = message.reply()
@@ -703,7 +698,7 @@ public actor SandboxService {
             try await self.monitor.track(id: id, waitingOn: waitFunc)
         } catch {
             try? await self.cleanUpContainer(containerInfo: info)
-            self.setState(.stopped(nil))
+            self.setState(.stopped)
             throw error
         }
     }
@@ -828,7 +823,7 @@ public actor SandboxService {
             let ctrInfo = try await getContainer()
 
             switch await self.state {
-            case .stopped(_), .stopping:
+            case .stopped, .stopping:
                 return
             default:
                 break
@@ -839,7 +834,7 @@ public actor SandboxService {
             } catch {
                 self.log.error("failed to clean up container", metadata: ["error": "\(error)"])
             }
-            await setState(.stopped(exitStatus.exitCode))
+            await setState(.stopped)
         }
     }
 
@@ -1077,14 +1072,7 @@ public actor SandboxService {
         await self.stopSocketForwarders()
 
         let status = exitStatus ?? ExitStatus(exitCode: 255)
-        self.waiters.withLock {
-            if let waiters = $0[id], let waiters {
-                for cc in waiters {
-                    cc.resume(returning: status)
-                }
-            }
-        }
-        self.invalidateWaiters(for: id)
+        self.releaseWaiters(for: id, status: status)
     }
 }
 
@@ -1304,34 +1292,31 @@ extension FileHandle: @retroactive ReaderStream, @retroactive Writer {
 // MARK: State handler and bundle creation helpers
 
 extension SandboxService {
-    private func addWaiter(id: String, cont: CheckedContinuation<ExitStatus, Never>) -> Bool {
-        self.waiters.withLock { waiters in
-            var current: [CheckedContinuation<ExitStatus, Never>]
-            switch waiters[id] {
-            case .none:
-                current = []
-            case .some(.some(let value)):
-                current = value
-            case .some(.none):
-                return false
-            }
-
-            current.append(cont)
-            waiters[id] = current
-            return true
+    private func initializeWaiters(for id: String) throws {
+        guard waiters[id] == nil else {
+            throw ContainerizationError(.invalidState, message: "waiter for \(id) already initialized")
         }
+        waiters[id] = ExitWaiter()
     }
 
-    private func removeWaiters(for id: String) {
-        self.waiters.withLock { waiters in
-            waiters[id] = []
+    private func addWaiter(id: String, cont: CheckedContinuation<ExitStatus, Never>) -> (Bool, Int32?) {
+        guard let current = waiters[id] else {
+            // No waiter initialized at all
+            return (false, nil)
         }
-    }
 
-    private func invalidateWaiters(for id: String) {
-        self.waiters.withLock { waiters in
-            waiters[id] = .some(nil)
+        if current.exited() {
+            // Waiter initialzed but already exited
+            return (false, current.exitCode)
         }
+
+        // Waiter initialized and not exited. Guaranteed to exit later.
+        current.register(cont)
+        return (true, nil)
+    }
+
+    private func releaseWaiters(for id: String, status: ExitStatus) {
+        waiters[id]?.doExit(code: status.exitCode)
     }
 
     private func setUnderlyingProcess(_ id: String, _ process: LinuxProcess) throws {
@@ -1387,7 +1372,7 @@ extension SandboxService {
         /// At the beginning of stop() .running will be transitioned to .stopping.
         case stopping
         /// Once a stop is successful, .stopping will transition to .stopped.
-        case stopped(Int32?)
+        case stopped
         /// .shuttingDown will be the last state the sandbox service will ever be in. Shortly
         /// afterwards the process will exit.
         case shuttingDown

Tests/CLITests/Subcommands/Run/TestCLIRunLifecycle.swift

@@ -123,4 +123,19 @@ class TestCLIRunLifecycle: CLITest {
         }
         try? doRemove(name: name)
     }
+
+    @Test func testExecInvalidExcutable() async throws {
+        let name = getTestName()
+        try doLongRun(name: name)
+        defer {
+            try? doStop(name: name)
+        }
+
+        #expect(throws: CLIError.self, "executing invalid executable must throw error, not hang") {
+            try doExec(
+                name: name,
+                cmd: ["foobarbaz"]
+            )
+        }
+    }
 }