containerization/vminitd/Sources/VminitdCore/ProcessSupervisor.swift

//===----------------------------------------------------------------------===//
// Copyright © 2026 Apple Inc. and the Containerization project authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//===----------------------------------------------------------------------===//

#if os(Linux)

import ContainerizationOS
import Foundation
import Logging
import Synchronization

final class ProcessSupervisor: Sendable {
    private let poller: Epoll

    /// [Nucleic vendored patch] Handler table keyed by fd, each entry stamped with the
    /// registration generation echoed back through epoll (`Epoll.Event.generation`).
    /// Dispatch compares the event's generation against the live entry's, so an event
    /// queued for a CLOSED registration of a recycled fd number can never fire the new
    /// registration's handler (a stale EPOLLHUP used to be able to tear down a brand-new
    /// healthy connection that inherited the number mid-batch).
    private struct HandlerTable {
        var nextGeneration: UInt32 = 1
        var entries: [Int32: (generation: UInt32, handler: @Sendable (Epoll.Mask) -> Void)] = [:]
    }
    private let handlers = Mutex<HandlerTable>(HandlerTable())

    private let queue: DispatchQueue
    // `DispatchSourceSignal` is thread-safe.
    private nonisolated(unsafe) let source: DispatchSourceSignal

    private struct State {
        var processes: [any ContainerProcess] = []
        var log: Logger?
    }

    private let state: Mutex<State>
    private let reaperCommandRunner = ReaperCommandRunner()

    func setLog(_ log: Logger?) {
        self.state.withLock { $0.log = log }
    }

    static let `default` = ProcessSupervisor()

    private init() {
        let queue = DispatchQueue(label: "process-supervisor")
        self.source = DispatchSource.makeSignalSource(signal: SIGCHLD, queue: queue)
        self.queue = queue
        self.poller = try! Epoll()
        self.state = Mutex(State())
        let t = Thread {
            while true {
                guard let events = self.poller.wait() else {
                    return
                }
                if events.isEmpty {
                    return
                }
                for event in events {
                    // [Nucleic vendored patch] Dispatch only when the event belongs to the
                    // CURRENT registration of this fd number — an earlier handler in this
                    // batch may have closed the fd and something else re-registered the
                    // recycled number already (see HandlerTable).
                    let entry = self.handlers.withLock { $0.entries[event.fd] }
                    guard let entry, entry.generation == event.generation else { continue }
                    entry.handler(event.mask)
                }
            }
        }
        t.start()
    }

    /// Register a file descriptor for epoll monitoring with a handler.
    ///
    /// The handler is stored before the fd is added to epoll, ensuring no
    /// events are missed.
    ///
    /// [Nucleic vendored patch] Refuses (EEXIST) an fd that is already registered rather
    /// than clobbering its handler: the old overwrite-then-fail-EEXIST path destroyed the
    /// existing registration's handler AND removed the map entry, leaving the fd armed in
    /// epoll with no handler — a silently dead relay (the TerminalIO shared-fd case).
    func registerFd(
        _ fd: Int32,
        mask: Epoll.Mask = [.input, .output],
        handler: @escaping @Sendable (Epoll.Mask) -> Void
    ) throws {
        let generation: UInt32 = try self.handlers.withLock { table in
            guard table.entries[fd] == nil else { throw POSIXError(.EEXIST) }
            let generation = table.nextGeneration
            // 0 is reserved for the poller's internal eventFD registration.
            table.nextGeneration = table.nextGeneration &+ 1
            if table.nextGeneration == 0 { table.nextGeneration = 1 }
            table.entries[fd] = (generation, handler)
            return generation
        }
        do {
            try self.poller.add(fd, mask: mask, generation: generation)
        } catch {
            self.handlers.withLock { _ = $0.entries.removeValue(forKey: fd) }
            throw error
        }
    }

    /// Remove a file descriptor from epoll monitoring and discard its handler.
    func unregisterFd(_ fd: Int32) throws {
        self.handlers.withLock { _ = $0.entries.removeValue(forKey: fd) }
        try self.poller.delete(fd)
    }

    func ready() {
        self.source.setEventHandler {
            self.handleSignal()
        }
        self.source.resume()
    }

    private func handleSignal() {
        dispatchPrecondition(condition: .onQueue(queue))

        let exited = Reaper.reap()

        for (pid, status) in exited {
            reaperCommandRunner.notifyExit(pid: pid, status: status)
        }

        self.state.withLock { state in
            state.log?.debug("received SIGCHLD, reaping processes")
            state.log?.debug("finished wait4 of \(exited.count) processes")
            state.log?.debug("checking for exit of managed process", metadata: ["exits": "\(exited)", "processes": "\(state.processes.count)"])

            let exitedProcesses = state.processes.filter { proc in
                exited.contains { pid, _ in
                    proc.pid == pid
                }
            }

            for proc in exitedProcesses {
                guard let pid = proc.pid else {
                    continue
                }

                if let status = exited[pid] {
                    state.log?.debug(
                        "managed process exited",
                        metadata: [
                            "pid": "\(pid)",
                            "status": "\(status)",
                            "count": "\(state.processes.count - 1)",
                        ])
                    proc.setExit(status)
                    state.processes.removeAll(where: { $0.pid == pid })
                }
            }
        }
    }

    func start(process: any ContainerProcess) async throws -> Int32 {
        self.state.withLock { state in
            state.log?.debug("in supervisor lock to start process")
            state.processes.append(process)
        }
        do {
            return try await process.start()
        } catch {
            self.state.withLock { state in
                state.processes.removeAll(where: { $0.id == process.id })
            }
            throw error
        }
    }

    /// Get a Runc instance configured with the reaper command runner
    func getRuncWithReaper(_ base: Runc = Runc()) -> Runc {
        var runc = base
        runc.commandRunner = reaperCommandRunner
        return runc
    }

    deinit {
        source.cancel()
        poller.shutdown()
    }
}

#endif