containerization/Sources/Containerization/EmergencyVMHandle.swift

//===----------------------------------------------------------------------===//
// Copyright © 2025-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(macOS)
import Foundation
import Synchronization
@preconcurrency import Virtualization

/// A once-only escape hatch for stopping a VM when its normal lifecycle path is unavailable.
///
/// The handle deliberately retains only Virtualization's VM object, its required dispatch queue,
/// and caller-supplied generation metadata. Requesting a stop does not enter a container state
/// gate, an instance lifecycle lock, or a guest-agent/RPC path.
public final class EmergencyVMHandle: @unchecked Sendable {
    private struct State {
        enum Phase {
            case ready
            case requested
            case stopped
        }

        var phase: Phase = .ready
        var waiters: [UUID: CheckedContinuation<Bool, Never>] = [:]
    }

    /// Stable identity for this exact retained VM, independent of caller generation numbering.
    public let id = UUID()
    public let generation: UInt64

    private nonisolated(unsafe) let virtualMachine: VZVirtualMachine
    private let queue: DispatchQueue
    private let state = Mutex(State())

    init(virtualMachine: VZVirtualMachine, queue: DispatchQueue, generation: UInt64) {
        self.virtualMachine = virtualMachine
        self.queue = queue
        self.generation = generation
    }

    /// Whether Virtualization has reported stop completion for this exact VM.
    public var hasStopped: Bool {
        state.withLock {
            if case .stopped = $0.phase { return true }
            return false
        }
    }

    /// Enqueue a hard VM stop exactly once. Returns `true` only for the caller which requested it.
    /// Completion is intentionally not awaited: this API is the independent request lane used when
    /// normal lifecycle tasks may themselves be unable to make progress.
    @discardableResult
    public func requestStop() -> Bool {
        let shouldRequest = state.withLock { state in
            guard case .ready = state.phase else { return false }
            state.phase = .requested
            return true
        }
        guard shouldRequest else { return false }

        queue.async { [self] in
            if virtualMachine.state == .stopped {
                finishStopped()
                return
            }
            if virtualMachine.state == .stopping {
                scheduleStopStateObservation()
                return
            }
            virtualMachine.stop { [weak self] _ in
                self?.scheduleStopStateObservation(immediate: true)
            }
        }
        return true
    }

    /// Wait until Virtualization reports this exact VM stopped, but never beyond `deadline`.
    /// `false` is an explicit force-retirement signal for the caller; the stop request itself stays
    /// independent and lock-free even if the normal instance lifecycle lane is wedged.
    public func waitUntilStopped(deadline: ContinuousClock.Instant) async -> Bool {
        let waiterID = UUID()
        return await withCheckedContinuation { continuation in
            let alreadyStopped = state.withLock { state in
                guard case .stopped = state.phase else {
                    state.waiters[waiterID] = continuation
                    return false
                }
                return true
            }
            if alreadyStopped {
                continuation.resume(returning: true)
                return
            }

            Task { [weak self] in
                do {
                    try await ContinuousClock().sleep(until: deadline)
                } catch {
                    // Cancellation is also a bounded failure for the waiting caller. The retained
                    // stop request continues independently on Virtualization's queue.
                }
                self?.finishWaiter(waiterID, stopped: false)
            }
        }
    }

    private func scheduleStopStateObservation(immediate: Bool = false) {
        queue.asyncAfter(deadline: .now() + (immediate ? 0 : 0.05)) { [weak self] in
            guard let self else { return }
            guard case .requested = state.withLock({ $0.phase }) else { return }
            if virtualMachine.state == .stopped {
                finishStopped()
            } else {
                scheduleStopStateObservation()
            }
        }
    }

    private func finishStopped() {
        let waiters = state.withLock { state -> [CheckedContinuation<Bool, Never>] in
            guard case .stopped = state.phase else {
                state.phase = .stopped
                let waiters = Array(state.waiters.values)
                state.waiters.removeAll()
                return waiters
            }
            return []
        }
        for waiter in waiters { waiter.resume(returning: true) }
    }

    private func finishWaiter(_ id: UUID, stopped: Bool) {
        let waiter = state.withLock { $0.waiters.removeValue(forKey: id) }
        waiter?.resume(returning: stopped)
    }
}
#endif