diff --git a/Sources/Containerization/LinuxProcess.swift b/Sources/Containerization/LinuxProcess.swift index 4dc2a77f..fc2ab1cd 100644 --- a/Sources/Containerization/LinuxProcess.swift +++ b/Sources/Containerization/LinuxProcess.swift @@ -94,6 +94,13 @@ public final class LinuxProcess: Sendable { var pid: Int32 var stdio: StdioHandles var stdinRelay: Task<(), Never>? + var ioTracker: IoTracker? + + struct IoTracker { + let stream: AsyncStream + let cont: AsyncStream.Continuation + let configuredStreams: Int + } } /// The process ID for the container process. This will be -1 @@ -207,16 +214,18 @@ extension LinuxProcess { } } + var configuredStreams = 0 + let (stream, cc) = AsyncStream.makeStream() if let stdout = self.ioSetup.stdout { + configuredStreams += 1 handles[1]?.readabilityHandler = { handle in - // NOTE: We need some way to know when this data is done being piped, - // so DispatchGroup or similar. `availableData` is also pretty poor, - // as it always allocates. We can likely do the read loop ourselves - // with a buffer we allocate once on creation of the process. do { let data = handle.availableData if data.isEmpty { + // This block is called when the producer (the guest) closes + // the fd it is writing into. handles[1]?.readabilityHandler = nil + cc.yield() return } try stdout.writer.write(data) @@ -227,11 +236,13 @@ extension LinuxProcess { } if let stderr = self.ioSetup.stderr { + configuredStreams += 1 handles[2]?.readabilityHandler = { handle in do { let data = handle.availableData if data.isEmpty { handles[2]?.readabilityHandler = nil + cc.yield() return } try stderr.writer.write(data) @@ -240,6 +251,11 @@ extension LinuxProcess { } } } + if configuredStreams > 0 { + self.state.withLock { + $0.ioTracker = .init(stream: stream, cont: cc, configuredStreams: configuredStreams) + } + } return handles } @@ -318,11 +334,13 @@ extension LinuxProcess { @discardableResult public func wait(timeoutInSeconds: Int64? = nil) async throws -> Int32 { do { - return try await self.agent.waitProcess( + let code = try await self.agent.waitProcess( id: self.id, containerID: self.owningContainer, timeoutInSeconds: timeoutInSeconds ) + await self.waitIoComplete() + return code } catch { if error is ContainerizationError { throw error @@ -335,16 +353,37 @@ extension LinuxProcess { } } + /// Wait until the standard output and standard error streams for the process have concluded. + private func waitIoComplete() async { + let ioTracker = self.state.withLock { $0.ioTracker } + guard let ioTracker else { + return + } + do { + try await Timeout.run(seconds: 3) { + var counter = ioTracker.configuredStreams + for await _ in ioTracker.stream { + counter -= 1 + if counter == 0 { + ioTracker.cont.finish() + break + } + } + } + } catch { + self.logger?.error("Timeout waiting for IO to complete for process \(id): \(error)") + } + self.state.withLock { + $0.ioTracker = nil + } + } + /// Cleans up guest state and waits on and closes any host resources (stdio handles). public func delete() async throws { try await self.agent.deleteProcess( id: self.id, containerID: self.owningContainer ) - - // FIXME: Add in IO drain waiting here. We can wait for 2-3 seconds or - // so and then just continue on. - // Now free up stdio handles. try self.state.withLock { $0.stdinRelay?.cancel() diff --git a/Sources/Integration/ProcessTests.swift b/Sources/Integration/ProcessTests.swift index d2f60d93..458d8312 100644 --- a/Sources/Integration/ProcessTests.swift +++ b/Sources/Integration/ProcessTests.swift @@ -17,6 +17,7 @@ import ArgumentParser import Containerization import ContainerizationOCI +import Crypto import Foundation import Logging @@ -154,9 +155,8 @@ extension IntegrationSuite { } } - func testMultipleConcurrentProcessesOutput() async throws { - let id = "test-concurrent-processes-output" - + func testMultipleConcurrentProcessesOutputStress() async throws { + let id = "test-concurrent-processes-output-stress" let bs = try await bootstrap() let container = LinuxContainer( id, @@ -169,36 +169,49 @@ extension IntegrationSuite { try await container.create() try await container.start() - let execConfig = ContainerizationOCI.Process( - args: ["/bin/echo", "hi"], + let baseExecConfig = ContainerizationOCI.Process( + args: ["sh", "-c", "dd if=/dev/random of=/tmp/bytes bs=1M count=20 status=none ; sha256sum /tmp/bytes"], env: ["PATH=\(LinuxContainer.defaultPath)"] ) - + let buffer = BufferWriter() + let exec = try await container.exec( + "expected-value", + configuration: baseExecConfig, + stdout: buffer, + ) + try await exec.start() + let status = try await exec.wait() + if status != 0 { + throw IntegrationError.assert(msg: "process status \(status) != 0") + } + let output = String(data: buffer.data, encoding: .utf8)! + let expected = String(output.split(separator: " ").first!) try await withThrowingTaskGroup(of: Void.self) { group in + let execConfig = ContainerizationOCI.Process( + args: ["cat", "/tmp/bytes"], + env: ["PATH=\(LinuxContainer.defaultPath)"] + ) for i in 0...80 { let idx = i group.addTask { let buffer = BufferWriter() - - var config = execConfig - config.args[1] = "hi\(idx)" - let exec = try await container.exec( "exec-\(idx)", - configuration: config, + configuration: execConfig, stdout: buffer, ) try await exec.start() let status = try await exec.wait() if status != 0 { - throw IntegrationError.assert(msg: "process status \(status) != 0") + throw IntegrationError.assert(msg: "process \(idx) status for \(status) != 0") } - - let output = String(data: buffer.data, encoding: .utf8) - guard output == "hi\(idx)\n" else { + var hasher = SHA256() + hasher.update(data: buffer.data) + let hash = hasher.finalize().digestString.trimmingDigestPrefix + guard hash == expected else { throw IntegrationError.assert( - msg: "process should have returned on stdout 'hi\(idx)' != '\(output!))") + msg: "process \(idx) output \(hash) != expected \(expected)") } try await exec.delete() } @@ -210,12 +223,9 @@ extension IntegrationSuite { // kill the init process. try await container.kill(SIGKILL) - let status = try await container.wait() + try await container.wait() try await container.stop() - print("\(status)") } - } catch { - throw error } } diff --git a/Sources/Integration/Suite.swift b/Sources/Integration/Suite.swift index cc49943c..a21137c1 100644 --- a/Sources/Integration/Suite.swift +++ b/Sources/Integration/Suite.swift @@ -200,7 +200,7 @@ struct IntegrationSuite: AsyncParsableCommand { "process user": testProcessUser, "process home envvar": testProcessHomeEnvvar, "multiple concurrent processes": testMultipleConcurrentProcesses, - "multiple concurrent processes with output": testMultipleConcurrentProcessesOutput, + "multiple concurrent processes with output stress": testMultipleConcurrentProcessesOutputStress, "container hostname": testHostname, "container mount": testMounts, "nested virt": testNestedVirtualizationEnabled, diff --git a/vminitd/Sources/vminitd/ProcessSupervisor.swift b/vminitd/Sources/vminitd/ProcessSupervisor.swift index 640c61ec..35839c00 100644 --- a/vminitd/Sources/vminitd/ProcessSupervisor.swift +++ b/vminitd/Sources/vminitd/ProcessSupervisor.swift @@ -98,7 +98,7 @@ actor ProcessSupervisor { return try process.start() } catch { - self.log?.error("process start failed \(error)") + self.log?.error("process start failed \(error)", metadata: ["process-id": "\(process.id)"]) throw error } } diff --git a/vminitd/Sources/vminitd/StandardIO.swift b/vminitd/Sources/vminitd/StandardIO.swift index bb489043..1d4f045f 100644 --- a/vminitd/Sources/vminitd/StandardIO.swift +++ b/vminitd/Sources/vminitd/StandardIO.swift @@ -130,10 +130,9 @@ final class StandardIO: ManagedProcess.IO & Sendable { try ProcessSupervisor.default.poller.add(readFromFd, mask: EPOLLIN) { mask in if mask.isHangup && !mask.readyToRead { - self.cleanup(readFromFd, buffer: buf, log: self.log) + self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log) return } - // Loop so that in the case that someone wrote > buf.count down the pipe // we properly will drain it fully. while true { @@ -147,7 +146,7 @@ final class StandardIO: ManagedProcess.IO & Sendable { let w = writeTo.write(view) if w.wrote != r.read { self.log?.error("stopping relay: short write for stdio") - self.cleanup(readFromFd, buffer: buf, log: self.log) + self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log) return } } @@ -157,13 +156,13 @@ final class StandardIO: ManagedProcess.IO & Sendable { self.log?.error("failed with errno \(errno) while reading for fd \(readFromFd)") fallthrough case .eof: - self.cleanup(readFromFd, buffer: buf, log: self.log) + self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log) self.log?.debug("closing relay for \(readFromFd)") return case .again: // We read all we could, exit. if mask.isHangup { - self.cleanup(readFromFd, buffer: buf, log: self.log) + self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log) } return default: @@ -173,14 +172,18 @@ final class StandardIO: ManagedProcess.IO & Sendable { } } - func cleanup(_ fd: Int32, buffer: UnsafeMutableBufferPointer, log: Logger?) { + func cleanupRelay(readFd: Int32, writeFd: Int32, buffer: UnsafeMutableBufferPointer, log: Logger?) { do { // We could alternatively just allocate buffers in the constructor, and free them // on close(). buffer.deallocate() - try ProcessSupervisor.default.poller.delete(fd) + try ProcessSupervisor.default.poller.delete(readFd) } catch { - self.log?.error("failed to delete pipe fd from epoll \(fd): \(error)") + self.log?.error("failed to delete pipe fd from epoll \(readFd): \(error)") + } + if Foundation.close(writeFd) != 0 { + let err = POSIXError.fromErrno() + self.log?.error("failed to close write fd for StandardIO relay: \(String(describing:err))") } } diff --git a/vminitd/Sources/vminitd/TerminalIO.swift b/vminitd/Sources/vminitd/TerminalIO.swift index 38c82990..9967e8e4 100644 --- a/vminitd/Sources/vminitd/TerminalIO.swift +++ b/vminitd/Sources/vminitd/TerminalIO.swift @@ -96,10 +96,9 @@ final class TerminalIO: ManagedProcess.IO & Sendable { try ProcessSupervisor.default.poller.add(readFromFd, mask: EPOLLIN) { mask in if mask.isHangup && !mask.readyToRead { - self.cleanup(readFromFd, buffer: buf, log: self.log) + self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log) return } - // Loop so that in the case that someone wrote > buf.count down the pipe // we properly will drain it fully. while true { @@ -113,7 +112,7 @@ final class TerminalIO: ManagedProcess.IO & Sendable { let w = writeTo.write(view) if w.wrote != r.read { self.log?.error("stopping relay: short write for stdio") - self.cleanup(readFromFd, buffer: buf, log: self.log) + self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log) return } } @@ -123,13 +122,13 @@ final class TerminalIO: ManagedProcess.IO & Sendable { self.log?.error("failed with errno \(errno) while reading for fd \(readFromFd)") fallthrough case .eof: - self.cleanup(readFromFd, buffer: buf, log: self.log) + self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log) self.log?.debug("closing relay for \(readFromFd)") return case .again: // We read all we could, exit. if mask.isHangup { - self.cleanup(readFromFd, buffer: buf, log: self.log) + self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log) } return default: @@ -139,14 +138,18 @@ final class TerminalIO: ManagedProcess.IO & Sendable { } } - func cleanup(_ fd: Int32, buffer: UnsafeMutableBufferPointer, log: Logger?) { + func cleanupRelay(readFd: Int32, writeFd: Int32, buffer: UnsafeMutableBufferPointer, log: Logger?) { do { // We could alternatively just allocate buffers in the constructor, and free them // on close(). buffer.deallocate() - try ProcessSupervisor.default.poller.delete(fd) + try ProcessSupervisor.default.poller.delete(readFd) } catch { - self.log?.error("failed to delete pipe fd from epoll \(fd): \(error)") + self.log?.error("failed to delete pipe fd from epoll \(readFd): \(error)") + } + if Foundation.close(writeFd) != 0 { + let err = POSIXError.fromErrno() + self.log?.error("failed to close write fd for TerminalIO relay: \(String(describing:err))") } }