Add a way to detect if a process' io has completed (#110)

This change adds a new private method `waitIoComplete` on the
`LinuxProcess` type.

This method is called internally when the user calls `wait` for a
process - and it tries to give the IO streams some time to clear their
buffers.

Internally, this method sets up an `AsyncStream` down which an item is
sent when the vsock connection for either stdout/stderr is terminated.
We get this termination signal when the readability handler for the
associated fd fires with a no available data.

Inside the guest - once we are done relaying the IO from the process
into the socket connection, we close the socket fd which triggers the
above.

All this logic is wrapped around a timeout of 3 seconds, just to ensure
the method does not block forever.

---------

Signed-off-by: Aditya Ramani <a_ramani@apple.com>
This commit is contained in:
Aditya Ramani
2025-06-13 11:23:41 -04:00
committed by GitHub
parent a56fdb7046
commit a4a0cdfae1
6 changed files with 102 additions and 47 deletions
+48 -9
View File
@@ -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<Void>
let cont: AsyncStream<Void>.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<Void>.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()
+30 -20
View File
@@ -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
}
}
+1 -1
View File
@@ -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,
@@ -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
}
}
+11 -8
View File
@@ -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<UInt8>, log: Logger?) {
func cleanupRelay(readFd: Int32, writeFd: Int32, buffer: UnsafeMutableBufferPointer<UInt8>, 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))")
}
}
+11 -8
View File
@@ -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<UInt8>, log: Logger?) {
func cleanupRelay(readFd: Int32, writeFd: Int32, buffer: UnsafeMutableBufferPointer<UInt8>, 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))")
}
}