VsockConnectionStream: Conform to AsyncSequence (#380)

I didn't like how we expose the asyncstream via a public connections
field. We should have the type conform to AsyncSequence and then hide
the underlying stream. This also stops listening on the stdio ports
after we get the initial connection.

Author: dcantah

Sources/Containerization/LinuxProcess.swift

@@ -119,19 +119,21 @@ public final class LinuxProcess: Sendable {
 extension LinuxProcess {
     func setupIO(streams: [VsockConnectionStream?]) async throws -> [FileHandle?] {
         let handles = try await Timeout.run(seconds: 3) {
-            await withTaskGroup(of: (Int, FileHandle?).self) { group in
+            try await withThrowingTaskGroup(of: (Int, FileHandle?).self) { group in
                 var results = [FileHandle?](repeating: nil, count: 3)
 
                 for (index, stream) in streams.enumerated() {
                     guard let stream = stream else { continue }
 
                     group.addTask {
-                        let first = await stream.connections.first(where: { _ in true })
+                        let first = await stream.first(where: { _ in true })
+                        stream.finish()
+                        try self.vm.stopListen(stream.port)
                         return (index, first)
                     }
                 }
 
-                for await (index, fileHandle) in group {
+                for try await (index, fileHandle) in group {
                     results[index] = fileHandle
                 }
                 return results

Sources/Containerization/UnixSocketRelay.swift

@@ -201,7 +201,7 @@ extension SocketRelay {
             $0.t = Task {
                 do {
                     defer { connectionStream.finish() }
-                    for await connection in connectionStream.connections {
+                    for await connection in connectionStream {
                         try await self.handleGuestVsockConn(
                             vsockConn: connection,
                             hostConnectionPath: hostPath,

Sources/Containerization/VsockConnectionStream.swift

@@ -21,9 +21,11 @@ import Virtualization
 #endif
 
 /// A stream of vsock connections.
-public final class VsockConnectionStream: NSObject, Sendable {
+public final class VsockConnectionStream: NSObject, Sendable, AsyncSequence {
+    public typealias Element = FileHandle
+
     /// A stream of connections dialed from the remote.
-    public let connections: AsyncStream<FileHandle>
+    private let connections: AsyncStream<FileHandle>
     /// The port the connections are for.
     public let port: UInt32
 
@@ -39,6 +41,10 @@ public final class VsockConnectionStream: NSObject, Sendable {
     public func finish() {
         self.cont.finish()
     }
+
+    public func makeAsyncIterator() -> AsyncStream<FileHandle>.AsyncIterator {
+        connections.makeAsyncIterator()
+    }
 }
 
 #if os(macOS)
@@ -49,9 +55,18 @@ extension VsockConnectionStream: VZVirtioSocketListenerDelegate {
         from _: VZVirtioSocketDevice
     ) -> Bool {
         let fd = dup(conn.fileDescriptor)
+        guard fd != -1 else {
+            return false
+        }
         conn.close()
 
-        cont.yield(FileHandle(fileDescriptor: fd, closeOnDealloc: false))
+        let fh = FileHandle(fileDescriptor: fd, closeOnDealloc: false)
+        let result = cont.yield(fh)
+        if case .terminated = result {
+            try? fh.close()
+            return false
+        }
+
         return true
     }
 }