apple--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
    private let handlers = Mutex<[Int32: @Sendable (Epoll.Mask) -> Void]>([:])

    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 {
                    let handler = self.handlers.withLock { $0[event.fd] }
                    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.
    func registerFd(
        _ fd: Int32,
        mask: Epoll.Mask = [.input, .output],
        handler: @escaping @Sendable (Epoll.Mask) -> Void
    ) throws {
        self.handlers.withLock { $0[fd] = handler }
        do {
            try self.poller.add(fd, mask: mask)
        } catch {
            self.handlers.withLock { _ = $0.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.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