Move STDIO input reads off the cooperative executor

Sources/XcodeMCPProxyKit/Stdio/StdioAdapter.swift

@@ -1,40 +1,8 @@
 import Foundation
 import Logging
-import Synchronization
 import XcodeMCPKit
 import XcodeMCPProxyRuntime
 
-private final class StdioReadIteratorReadiness: Sendable {
-    private struct State {
-        var isReady = false
-        var waiters: [CheckedContinuation<Void, Never>] = []
-    }
-
-    private let state = Mutex(State())
-
-    func markReady() {
-        let waiters: [CheckedContinuation<Void, Never>] = state.withLock { state in
-            guard state.isReady == false else { return [] }
-            state.isReady = true
-            let waiters = state.waiters
-            state.waiters.removeAll()
-            return waiters
-        }
-        for waiter in waiters { waiter.resume() }
-    }
-
-    func waitUntilReady() async {
-        await withCheckedContinuation { continuation in
-            let shouldResume = state.withLock { state in
-                guard state.isReady == false else { return true }
-                state.waiters.append(continuation)
-                return false
-            }
-            if shouldResume { continuation.resume() }
-        }
-    }
-}
-
 package struct StdioAdapterShutdownPolicy: Sendable {
     package let requestDrainTimeout: Duration
     package let requestDrainPollInterval: Duration
@@ -65,12 +33,11 @@ actor StdioAdapter {
     }
 
     private let requestTimeout: Duration?
-    private let inputHandle: FileHandle
+    private let inputReader: any StdioInputReading
     private let outputWriter: StdioWriter
     private let logger: Logger
     private let authority: MCPClientSessionAuthority
     private let shutdownPolicy: StdioAdapterShutdownPolicy
-    private let readIteratorReadiness = StdioReadIteratorReadiness()
 
     private var framer = StdioFramer()
     private var requestTasks: [UUID: Task<Void, Never>] = [:]
@@ -130,9 +97,25 @@ actor StdioAdapter {
         output: FileHandle,
         recipe: MCPTransportRecipe,
         shutdownPolicy: StdioAdapterShutdownPolicy
+    ) {
+        self.init(
+            requestTimeout: requestTimeout,
+            inputReader: StdioInputChannel(handle: input),
+            output: output,
+            recipe: recipe,
+            shutdownPolicy: shutdownPolicy
+        )
+    }
+
+    init(
+        requestTimeout: Duration?,
+        inputReader: any StdioInputReading,
+        output: FileHandle,
+        recipe: MCPTransportRecipe,
+        shutdownPolicy: StdioAdapterShutdownPolicy
     ) {
         self.requestTimeout = requestTimeout
-        self.inputHandle = input
+        self.inputReader = inputReader
         let logger = ProxyLogging.make("stdio.adapter")
         self.logger = logger
         self.outputWriter = StdioWriter(handle: output, logger: logger)
@@ -152,18 +135,13 @@ actor StdioAdapter {
         }
         lifecycle = .running
         startAuthorityEventTask()
-        let input = inputHandle
-        let readIteratorReadiness = readIteratorReadiness
+        let inputReader = inputReader
         readTask = Task { [weak self] in
-            var iterator = input.bytes.makeAsyncIterator()
-            // Do not signal readiness before iterator creation: Foundation accesses the
-            // descriptor here and raises NSFileHandleOperationException if close wins.
-            readIteratorReadiness.markReady()
             var terminalError: (any Error)?
             do {
-                while let byte = try await iterator.next() {
+                while let chunk = try await inputReader.read() {
                     guard Task.isCancelled == false else { break }
-                    await self?.handleInput(Data([byte]))
+                    await self?.handleInput(chunk)
                 }
             } catch is CancellationError {
                 // Explicit stop owns completion.
@@ -198,6 +176,7 @@ actor StdioAdapter {
 
     isolated deinit {
         readTask?.cancel()
+        inputReader.stop()
         authorityEventTask?.cancel()
         for task in requestTasks.values { task.cancel() }
         closeTask?.cancel()
@@ -385,10 +364,8 @@ private extension StdioAdapter {
     }
 
     func performClose(drainsOutput: Bool) async {
-        if readTask != nil {
-            await readIteratorReadiness.waitUntilReady()
-            try? inputHandle.close()
-        }
+        inputReader.stop()
+        await inputReader.waitUntilStopped()
         await authority.close()
         await drainRequestTasks()
         let eventTask = authorityEventTask

Sources/XcodeMCPProxyKit/Stdio/StdioInputChannel.swift

@@ -0,0 +1,211 @@
+import Darwin
+import Dispatch
+import Foundation
+import XcodeMCPKit
+
+protocol StdioInputReading: Sendable {
+    func read() async throws -> Data?
+    func stop()
+    func waitUntilStopped() async
+}
+
+private struct StdioInputReadError: Error, CustomStringConvertible, Sendable {
+    let code: Int32
+
+    var description: String {
+        "STDIO input read failed with errno \(code)"
+    }
+}
+
+final class StdioInputChannel: StdioInputReading, @unchecked Sendable {
+    private enum Lifecycle {
+        case open
+        case endOfFile
+        case failed(StdioInputReadError)
+        case stopped
+    }
+
+    private static let maxReadOperationByteCount = 64 * 1024
+
+    private let callbackQueue = DispatchQueue(label: "XcodeMCPProxy.StdioInputChannel.io")
+    private let channel: DispatchIO
+    private let terminal: AsyncTerminalSignal
+
+    // Access is confined to callbackQueue.
+    private var lifecycle: Lifecycle = .open
+    private var bufferedInput = Data()
+    private var pendingRead: CheckedContinuation<Data?, any Error>?
+    private var isReadOperationActive = false
+    private var activeOperationByteCount = 0
+
+    init(handle: FileHandle) {
+        let descriptor = dup(handle.fileDescriptor)
+        precondition(descriptor >= 0, "STDIO input FileHandle must be open")
+
+        let terminal = AsyncTerminalSignal()
+        self.terminal = terminal
+        self.channel = DispatchIO(
+            type: .stream,
+            fileDescriptor: descriptor,
+            queue: callbackQueue
+        ) { _ in
+            _ = Darwin.close(descriptor)
+            terminal.signal()
+        }
+        channel.setLimit(lowWater: 1)
+        channel.setLimit(highWater: Self.maxReadOperationByteCount)
+    }
+
+    func read() async throws -> Data? {
+        try await withTaskCancellationHandler {
+            try Task.checkCancellation()
+            return try await withCheckedThrowingContinuation { continuation in
+                callbackQueue.async { [self] in
+                    registerRead(continuation)
+                }
+            }
+        } onCancel: {
+            stop()
+        }
+    }
+
+    func stop() {
+        callbackQueue.async { [self] in
+            guard case .open = lifecycle else { return }
+            lifecycle = .stopped
+            bufferedInput.removeAll()
+            let continuation = pendingRead
+            pendingRead = nil
+            channel.close(flags: .stop)
+            continuation?.resume(throwing: CancellationError())
+        }
+    }
+
+    func waitUntilStopped() async {
+        await terminal.wait()
+    }
+
+    private func registerRead(
+        _ continuation: CheckedContinuation<Data?, any Error>
+    ) {
+        precondition(pendingRead == nil, "STDIO input supports one consumer")
+
+        if bufferedInput.isEmpty == false {
+            let chunk = bufferedInput
+            bufferedInput = Data()
+            continuation.resume(returning: chunk)
+            return
+        }
+
+        switch lifecycle {
+        case .open:
+            pendingRead = continuation
+            startReadOperationIfNeeded()
+        case .endOfFile:
+            continuation.resume(returning: nil)
+        case .failed(let error):
+            continuation.resume(throwing: error)
+        case .stopped:
+            continuation.resume(throwing: CancellationError())
+        }
+    }
+
+    private func startReadOperationIfNeeded() {
+        guard
+            case .open = lifecycle,
+            pendingRead != nil,
+            bufferedInput.isEmpty,
+            isReadOperationActive == false
+        else {
+            return
+        }
+
+        isReadOperationActive = true
+        activeOperationByteCount = 0
+        channel.read(
+            offset: 0,
+            length: Self.maxReadOperationByteCount,
+            queue: callbackQueue
+        ) { [weak self] done, dispatchData, error in
+            self?.receiveReadResult(
+                done: done,
+                data: dispatchData.map { Data($0) },
+                error: error
+            )
+        }
+    }
+
+    private func receiveReadResult(
+        done: Bool,
+        data: Data?,
+        error: Int32
+    ) {
+        guard case .open = lifecycle else { return }
+
+        let chunk = data ?? Data()
+        if error != 0 {
+            isReadOperationActive = false
+            finishWithError(StdioInputReadError(code: error))
+            return
+        }
+
+        if chunk.isEmpty == false {
+            activeOperationByteCount += chunk.count
+            precondition(
+                activeOperationByteCount <= Self.maxReadOperationByteCount,
+                "DispatchIO exceeded the requested STDIO input byte count"
+            )
+            deliverOrBuffer(chunk)
+        }
+
+        guard done else { return }
+        isReadOperationActive = false
+
+        if chunk.isEmpty {
+            finishAtEndOfFile()
+        } else {
+            startReadOperationIfNeeded()
+        }
+    }
+
+    private func deliverOrBuffer(_ chunk: Data) {
+        if let continuation = pendingRead {
+            pendingRead = nil
+            continuation.resume(returning: chunk)
+        } else {
+            bufferedInput.append(chunk)
+            precondition(
+                bufferedInput.count <= Self.maxReadOperationByteCount,
+                "STDIO input buffer exceeded its bounded read operation"
+            )
+        }
+    }
+
+    private func finishAtEndOfFile() {
+        lifecycle = .endOfFile
+        channel.close()
+        resumeTerminalReadIfPossible()
+    }
+
+    private func finishWithError(_ error: StdioInputReadError) {
+        lifecycle = .failed(error)
+        bufferedInput.removeAll()
+        channel.close(flags: .stop)
+        resumeTerminalReadIfPossible()
+    }
+
+    private func resumeTerminalReadIfPossible() {
+        guard bufferedInput.isEmpty, let continuation = pendingRead else { return }
+        pendingRead = nil
+        switch lifecycle {
+        case .endOfFile:
+            continuation.resume(returning: nil)
+        case .failed(let error):
+            continuation.resume(throwing: error)
+        case .stopped:
+            continuation.resume(throwing: CancellationError())
+        case .open:
+            preconditionFailure("STDIO input terminal read resumed while open")
+        }
+    }
+}