run cleanupRelay only once in TerminalIO and StandardIO (#153)

elijah-wright · web-flow · commit 538ea78c335c · 2025-06-26T12:03:17.000-04:00
diff --git a/Sources/Integration/ProcessTests.swift b/Sources/Integration/ProcessTests.swift
@@ -205,7 +205,7 @@ extension IntegrationSuite {
 
                         let status = try await exec.wait()
                         if status != 0 {
-                            throw IntegrationError.assert(msg: "process \(idx) status for  \(status) != 0")
+                            throw IntegrationError.assert(msg: "process \(idx) status \(status) != 0")
                         }
                         var hasher = SHA256()
                         hasher.update(data: buffer.data)
diff --git a/vminitd/Sources/vminitd/ManagedProcess.swift b/vminitd/Sources/vminitd/ManagedProcess.swift
@@ -131,7 +131,11 @@ final class ManagedProcess: Sendable {
 extension ManagedProcess {
     func start() throws -> Int32 {
         try self.lock.withLock {
-            log.debug("starting managed process")
+            log.debug(
+                "starting managed process",
+                metadata: [
+                    "id": "\(id)"
+                ])
 
             // Start the underlying process.
             try process.start()
@@ -154,7 +158,8 @@ extension ManagedProcess {
             log.debug(
                 "started managed process",
                 metadata: [
-                    "pid": "\(i)"
+                    "pid": "\(i)",
+                    "id": "\(id)",
                 ])
 
             return i
diff --git a/vminitd/Sources/vminitd/ProcessSupervisor.swift b/vminitd/Sources/vminitd/ProcessSupervisor.swift
@@ -64,10 +64,15 @@ actor ProcessSupervisor {
         let exited = Reaper.reap()
         self.log?.debug("finished wait4 of \(exited.count) processes")
 
-        for proc in processes {
-            let pid = proc.pid
-            self.log?.debug("checking for exit of managed process", metadata: ["pid": "\(pid)", "exits": "\(exited)"])
+        self.log?.debug("checking for exit of managed process", metadata: ["exits": "\(exited)", "processes": "\(processes.count)"])
+        let exitedProcesses = self.processes.filter { proc in
+            exited.contains { pid, _ in
+                proc.pid == pid
+            }
+        }
 
+        for proc in exitedProcesses {
+            let pid = proc.pid
             if pid <= 0 {
                 continue
             }
@@ -80,7 +85,6 @@ actor ProcessSupervisor {
                         "status": "\(status)",
                         "count": "\(processes.count - 1)",
                     ])
-
                 proc.setExit(status)
                 self.processes.removeAll(where: { $0.pid == pid })
             }
@@ -95,7 +99,6 @@ actor ProcessSupervisor {
 
         do {
             self.processes.append(process)
-
             return try process.start()
         } catch {
             self.log?.error("process start failed \(error)", metadata: ["process-id": "\(process.id)"])
diff --git a/vminitd/Sources/vminitd/StandardIO.swift b/vminitd/Sources/vminitd/StandardIO.swift
@@ -78,7 +78,7 @@ final class StandardIO: ManagedProcess.IO & Sendable {
                 port: stdinPort,
                 cid: VsockType.hostCID
             )
-            let stdinSocket = try Socket(type: type)
+            let stdinSocket = try Socket(type: type, closeOnDeinit: false)
             try stdinSocket.connect()
             self.stdinSocket = stdinSocket
 
@@ -93,7 +93,8 @@ final class StandardIO: ManagedProcess.IO & Sendable {
                 port: stdoutPort,
                 cid: VsockType.hostCID
             )
-            let stdoutSocket = try Socket(type: type)
+            // These fd's get closed when cleanupRelay is called
+            let stdoutSocket = try Socket(type: type, closeOnDeinit: false)
             try stdoutSocket.connect()
             self.stdoutSocket = stdoutSocket
 
@@ -108,7 +109,7 @@ final class StandardIO: ManagedProcess.IO & Sendable {
                 port: stderrPort,
                 cid: VsockType.hostCID
             )
-            let stderrSocket = try Socket(type: type)
+            let stderrSocket = try Socket(type: type, closeOnDeinit: false)
             try stderrSocket.connect()
             self.stderrSocket = stderrSocket
 
@@ -125,12 +126,19 @@ final class StandardIO: ManagedProcess.IO & Sendable {
     func relay(readFromFd: Int32, writeToFd: Int32) throws {
         let readFrom = OSFile(fd: readFromFd)
         let writeTo = OSFile(fd: writeToFd)
-        // `buf` isn't used concurrently.
+        // `buf` and `didCleanup` aren't used concurrently.
         nonisolated(unsafe) let buf = UnsafeMutableBufferPointer<UInt8>.allocate(capacity: Int(getpagesize()))
+        nonisolated(unsafe) var didCleanup = false
+
+        let cleanupRelay: @Sendable () -> Void = {
+            if didCleanup { return }
+            didCleanup = true
+            self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+        }
 
         try ProcessSupervisor.default.poller.add(readFromFd, mask: EPOLLIN) { mask in
             if mask.isHangup && !mask.readyToRead {
-                self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+                cleanupRelay()
                 return
             }
             // Loop so that in the case that someone wrote > buf.count down the pipe
@@ -146,7 +154,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.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+                        cleanupRelay()
                         return
                     }
                 }
@@ -156,13 +164,13 @@ final class StandardIO: ManagedProcess.IO & Sendable {
                     self.log?.error("failed with errno \(errno) while reading for fd \(readFromFd)")
                     fallthrough
                 case .eof:
-                    self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+                    cleanupRelay()
                     self.log?.debug("closing relay for \(readFromFd)")
                     return
                 case .again:
                     // We read all we could, exit.
                     if mask.isHangup {
-                        self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+                        cleanupRelay()
                     }
                     return
                 default:
diff --git a/vminitd/Sources/vminitd/TerminalIO.swift b/vminitd/Sources/vminitd/TerminalIO.swift
@@ -65,7 +65,7 @@ final class TerminalIO: ManagedProcess.IO & Sendable {
                 port: stdinPort,
                 cid: VsockType.hostCID
             )
-            let stdinSocket = try Socket(type: type)
+            let stdinSocket = try Socket(type: type, closeOnDeinit: false)
             try stdinSocket.connect()
             self.stdinSocket = stdinSocket
 
@@ -80,7 +80,7 @@ final class TerminalIO: ManagedProcess.IO & Sendable {
                 port: stdoutPort,
                 cid: VsockType.hostCID
             )
-            let stdoutSocket = try Socket(type: type)
+            let stdoutSocket = try Socket(type: type, closeOnDeinit: false)
             try stdoutSocket.connect()
             self.stdoutSocket = stdoutSocket
 
@@ -94,12 +94,19 @@ final class TerminalIO: ManagedProcess.IO & Sendable {
     func relay(readFromFd: Int32, writeToFd: Int32) throws {
         let readFrom = OSFile(fd: readFromFd)
         let writeTo = OSFile(fd: writeToFd)
-        // `buf` isn't used concurrently.
+        // `buf` and `didCleanup` aren't used concurrently.
         nonisolated(unsafe) let buf = UnsafeMutableBufferPointer<UInt8>.allocate(capacity: Int(getpagesize()))
+        nonisolated(unsafe) var didCleanup = false
+
+        let cleanupRelay: @Sendable () -> Void = {
+            if didCleanup { return }
+            didCleanup = true
+            self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+        }
 
         try ProcessSupervisor.default.poller.add(readFromFd, mask: EPOLLIN) { mask in
             if mask.isHangup && !mask.readyToRead {
-                self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+                cleanupRelay()
                 return
             }
             // Loop so that in the case that someone wrote > buf.count down the pipe
@@ -115,7 +122,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.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+                        cleanupRelay()
                         return
                     }
                 }
@@ -125,13 +132,13 @@ final class TerminalIO: ManagedProcess.IO & Sendable {
                     self.log?.error("failed with errno \(errno) while reading for fd \(readFromFd)")
                     fallthrough
                 case .eof:
-                    self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+                    cleanupRelay()
                     self.log?.debug("closing relay for \(readFromFd)")
                     return
                 case .again:
                     // We read all we could, exit.
                     if mask.isHangup {
-                        self.cleanupRelay(readFd: readFromFd, writeFd: writeToFd, buffer: buf, log: self.log)
+                        cleanupRelay()
                     }
                     return
                 default:

Original file line number	Diff line number	Diff line change
`@@ -205,7 +205,7 @@ extension IntegrationSuite {`
`205`	`205`
`206`	`206`	`let status = try await exec.wait()`
`207`	`207`	`if status != 0 {`
`208`		`- throw IntegrationError.assert(msg: "process \(idx) status for \(status) != 0")`
	`208`	`+ throw IntegrationError.assert(msg: "process \(idx) status \(status) != 0")`
`209`	`209`	`}`
`210`	`210`	`var hasher = SHA256()`
`211`	`211`	`hasher.update(data: buffer.data)`