Files
nucleic/Sources/NucleicCore/Container/Windows/WslcBrokerClient.swift
T

245 lines
11 KiB
Swift

import Foundation
import NucleicProtocol
// Compiled ONLY on Windows (this directory is excluded elsewhere — Package.swift).
/// The hostd side of the broker pipe (docs/WINDOWS_PORT.md §2.3, §3.3): spawns
/// `nucleic-brokerd.exe` as a supervised child, speaks NDJSON JSON-RPC 2.0 over its stdio
/// via the shared ``JSONRPCConnection``, routes its notifications (process stdio, pull
/// progress, session death) to their consumers, and restarts it with backoff when it dies.
/// wslc state is service-backed, so a restarted broker re-opens the same named session and
/// re-enumerates containers; after a successful reattach `onReattach` fires so the owner can
/// run `ContainerManager.reconcile(activeSessions:)` against the fresh listing.
///
/// The wire contract is pinned from the C# side by `windows/NucleicBroker.Tests` — every
/// method/key used here appears in those tests.
public actor WslcBrokerClient {
public struct BrokerLost: Error, CustomStringConvertible {
public let description = "the wslc broker process is not available"
}
/// Per-process notification sinks, registered by ``WslcProcessHandle``.
struct ProcSinks {
let stdout: @Sendable (Data) -> Void
let stderr: @Sendable (Data) -> Void
let exit: @Sendable (Int32) -> Void
}
/// A process event that arrived before its ``WslcProcessHandle`` existed.
///
/// This is not a theoretical window. The broker enqueues the `proc.exec` response *before*
/// starting the process precisely so the response cannot be overtaken — but a client still
/// cannot register on time, because resolving the request's continuation only SCHEDULES it,
/// while the notification drain is a separate task already reading the next line. A command
/// that finishes immediately therefore lands `proc.stdout` and `proc.exit` before
/// ``register(procId:sinks:)`` runs, and the old code silently dropped both — so `wait()`
/// never returned and the caller hung forever.
///
/// `runCapturing` is the most exposed path, since every short probe it runs (cgroup reads,
/// shim re-seeding) is exactly the fast case. Observed on hardware through the C# spike, which
/// hung twice on it (docs/WINDOWS_PORT.md §13.3).
private enum PendingProcEvent {
case stdout(Data)
case stderr(Data)
case exit(Int32)
}
/// Buffered events per not-yet-registered procId, drained by ``register(procId:sinks:)``.
private var pendingProcEvents: [Int64: [PendingProcEvent]] = [:]
/// The window is a continuation hop, so a handful of events is the realistic maximum; the cap
/// only exists so a procId that never registers (a broker bug) cannot grow without bound.
private static let maxPendingEventsPerProc = 512
private func buffer(_ event: PendingProcEvent, for procId: Int64) {
var events = pendingProcEvents[procId] ?? []
let isExit: Bool = if case .exit = event { true } else { false }
// Always keep the exit — losing it is the hang this mechanism exists to prevent, whereas
// dropping overflow output only truncates a transcript.
guard isExit || events.count < Self.maxPendingEventsPerProc else { return }
events.append(event)
pendingProcEvents[procId] = events
}
private let brokerdPath: String
private var child: ChildProcess?
private var connection: JSONRPCConnection?
private var drainTask: Task<Void, Never>?
private var superviseTask: Task<Void, Never>?
private var procSinks: [Int64: ProcSinks] = [:]
private var restartBackoffNanos: UInt64 = 500_000_000
private var shutdownRequested = false
/// Capabilities from the broker's `hello` (empty until the first successful start).
public private(set) var capabilities: Set<String> = []
public private(set) var wslcVersion: String?
/// Fired after a broker crash+restart once the new broker answered `hello` — the owner
/// re-ensures the session and reconciles containers.
private let onReattach: @Sendable () async -> Void
/// Forwarded `session.down {reason}` notifications (the wslc VM died underneath us).
private let onSessionDown: @Sendable (String) async -> Void
/// Forwarded `image.pullProgress` notifications, folded into the download-progress UI.
private let onPullProgress: @Sendable (_ ref: String, _ status: String, _ current: Int64, _ total: Int64) async -> Void
public init(
brokerdPath: String? = nil,
onReattach: @Sendable @escaping () async -> Void = {},
onSessionDown: @Sendable @escaping (String) async -> Void = { _ in },
onPullProgress: @Sendable @escaping (String, String, Int64, Int64) async -> Void = { _, _, _, _ in }
) {
// Default: the exe ships beside nucleic-hostd in the MSIX payload (§2.4).
self.brokerdPath = brokerdPath
?? ProcessInfo.processInfo.environment["NUCLEIC_BROKERD_PATH"]
?? URL(fileURLWithPath: CommandLine.arguments[0])
.deletingLastPathComponent().appendingPathComponent("nucleic-brokerd.exe").path
self.onReattach = onReattach
self.onSessionDown = onSessionDown
self.onPullProgress = onPullProgress
}
// MARK: - Lifecycle
/// Spawn the broker (if not already running) and complete the `hello` exchange.
public func start() async throws {
guard connection == nil else { return }
shutdownRequested = false
let child = try ChildProcess(spec: ProcessSpec(
executable: brokerdPath, args: [],
cwd: FileManager.default.temporaryDirectory.path))
let connection = JSONRPCConnection(handle: child, includeVersionHeader: true)
self.child = child
self.connection = connection
// The broker never sends client-bound requests; answer any with an error.
await connection.start { request in
try? await connection.replyError(
to: request.id, code: -32601, message: "hostd accepts no requests")
}
drainTask = Task { [weak self] in
for await note in connection.notifications {
await self?.route(note)
}
}
superviseTask = Task { [weak self] in
let code = await child.wait()
await self?.brokerDied(exitCode: code)
}
let hello = try await connection.request("hello")
capabilities = Set((hello["capabilities"]?.arrayValue ?? []).compactMap(\.stringValue))
wslcVersion = hello["wslcVersion"]?.stringValue
restartBackoffNanos = 500_000_000 // healthy start → reset the backoff ladder
}
public func shutdown() async {
shutdownRequested = true
superviseTask?.cancel()
drainTask?.cancel()
await connection?.close()
await child?.terminate()
connection = nil
child = nil
}
private func brokerDied(exitCode: Int32) async {
await connection?.close()
connection = nil
child = nil
drainTask?.cancel()
// Every registered process is dead with the broker: fail their streams so backends
// hit the same recovery path a forced stream close serves (docs/WINDOWS_PORT.md §2.3).
for (_, sinks) in procSinks { sinks.exit(-1) }
procSinks.removeAll()
// Events buffered for handles that never registered belong to the dead broker's procId
// space, which the next broker reuses from 1 — keeping them would misdeliver.
pendingProcEvents.removeAll()
// Exponential-backoff restart loop (§2.3); a successful start() resets the ladder.
while !shutdownRequested {
let backoff = restartBackoffNanos
restartBackoffNanos = min(restartBackoffNanos * 2, 30_000_000_000)
try? await Task.sleep(nanoseconds: backoff)
guard !shutdownRequested else { return }
do {
try await start()
await onReattach()
return
} catch {
continue
}
}
}
// MARK: - Requests
/// One RPC round-trip; `BrokerLost` when the broker is down (a distinguishable error the
/// backends map to the stream-recovery path).
public func call(_ method: String, _ params: JSONValue? = nil) async throws -> JSONValue {
guard let connection else { throw BrokerLost() }
do {
return try await connection.request(method, params: params)
} catch JSONRPCConnection.RPCError.connectionClosed {
throw BrokerLost()
}
}
// MARK: - Process event routing
func register(procId: Int64, sinks: ProcSinks) {
procSinks[procId] = sinks
// Deliver anything that arrived before this handle existed — see `pendingProcEvents`.
guard let buffered = pendingProcEvents.removeValue(forKey: procId) else { return }
for event in buffered {
switch event {
case .stdout(let data): sinks.stdout(data)
case .stderr(let data): sinks.stderr(data)
case .exit(let code):
// The process is already over. Drop the sink as `route` would have.
procSinks[procId] = nil
sinks.exit(code)
return
}
}
}
func unregister(procId: Int64) {
procSinks[procId] = nil
pendingProcEvents[procId] = nil
}
private func route(_ note: JSONRPCConnection.Notification) async {
let params = note.params
switch note.method {
case "proc.stdout", "proc.stderr":
guard let rawId = params["procId"]?.intValue,
let b64 = params["b64"]?.stringValue,
let bytes = Data(base64Encoded: b64)
else { return }
let procId = Int64(rawId)
guard let sinks = procSinks[procId] else {
buffer(note.method == "proc.stderr" ? .stderr(bytes) : .stdout(bytes), for: procId)
return
}
(note.method == "proc.stderr" ? sinks.stderr : sinks.stdout)(bytes)
case "proc.exit":
guard let rawId = params["procId"]?.intValue else { return }
let procId = Int64(rawId)
let code = Int32(params["code"]?.intValue ?? -1)
guard let sinks = procSinks.removeValue(forKey: procId) else {
buffer(.exit(code), for: procId)
return
}
sinks.exit(code)
case "session.down":
await onSessionDown(params["reason"]?.stringValue ?? "unknown")
case "image.pullProgress":
await onPullProgress(
params["ref"]?.stringValue ?? "",
params["status"]?.stringValue ?? "",
Int64(params["current"]?.intValue ?? 0),
Int64(params["total"]?.intValue ?? 0))
default:
break // components.installProgress etc. — routed once onboarding consumes them
}
}
}