Allow callers to run a subprocess and provide low and high water marks when using SequenceOutput to emit standard output and standard error as soon as it arrives.

Sources/Subprocess/AsyncBufferSequence.swift

@@ -15,6 +15,8 @@
 @preconcurrency import SystemPackage
 #endif
 
+internal import Dispatch
+
 #if SubprocessSpan
 @available(SubprocessSpan, *)
 #endif
@@ -30,24 +32,34 @@ public struct AsyncBufferSequence: AsyncSequence, Sendable {
         private var buffer: [UInt8]
         private var currentPosition: Int
         private var finished: Bool
+        private var streamIterator: AsyncThrowingStream<TrackedPlatformDiskIO.StreamStatus, Swift.Error>.AsyncIterator
 
         internal init(diskIO: TrackedPlatformDiskIO) {
             self.diskIO = diskIO
             self.buffer = []
             self.currentPosition = 0
             self.finished = false
+            self.streamIterator = diskIO.readDataStream(upToLength: readBufferSize).makeAsyncIterator()
         }
 
-        public func next() async throws -> SequenceOutput.Buffer? {
-            let data = try await self.diskIO.readChunk(
-                upToLength: readBufferSize
-            )
-            if data == nil {
-                // We finished reading. Close the file descriptor now
+        public mutating func next() async throws -> SequenceOutput.Buffer? {
+            if let status = try await streamIterator.next() {
+                switch status {
+                case .data(let data):
+                    return data
+
+                case .endOfStream(let data):
+                    streamIterator = diskIO.readDataStream(upToLength: readBufferSize).makeAsyncIterator()
+                    return data
+
+                case .endOfFile:
+                    try self.diskIO.safelyClose()
+                    return nil
+                }
+            } else {
                 try self.diskIO.safelyClose()
                 return nil
             }
-            return data
         }
     }
 
@@ -62,6 +74,14 @@ public struct AsyncBufferSequence: AsyncSequence, Sendable {
     }
 }
 
+extension TrackedPlatformDiskIO {
+    internal enum StreamStatus {
+        case data(SequenceOutput.Buffer)
+        case endOfStream(SequenceOutput.Buffer)
+        case endOfFile
+    }
+}
+
 // MARK: - Page Size
 import _SubprocessCShims
 

Sources/Subprocess/Execution.swift

@@ -27,6 +27,8 @@ import Musl
 import WinSDK
 #endif
 
+internal import Dispatch
+
 /// An object that repersents a subprocess that has been
 /// executed. You can use this object to send signals to the
 /// child process as well as stream its output and error.
@@ -107,6 +109,8 @@ extension Execution where Output == SequenceOutput {
         else {
             fatalError("The standard output has already been consumed")
         }
+
+        // TODO: Make buffer size an option
         return AsyncBufferSequence(diskIO: readFd)
     }
 }
@@ -122,14 +126,16 @@ extension Execution where Error == SequenceOutput {
     /// via pipe under the hood and each pipe can only be consumed once.
     public var standardError: AsyncBufferSequence {
         let consumptionState = self.outputConsumptionState.bitwiseXor(
-            OutputConsumptionState.standardOutputConsumed
+            OutputConsumptionState.standardErrorConsumed
         )
 
         guard consumptionState.contains(.standardErrorConsumed),
             let readFd = self.errorPipe.readEnd
         else {
             fatalError("The standard output has already been consumed")
         }
+
+        // TODO: Make buffer size and option
         return AsyncBufferSequence(diskIO: readFd)
     }
 }

Sources/Subprocess/Platforms/Subprocess+Darwin.swift

@@ -57,6 +57,8 @@ public struct PlatformOptions: Sendable {
     /// Creates a session and sets the process group ID
     /// i.e. Detach from the terminal.
     public var createSession: Bool = false
+    public var outputOptions: StreamOptions = .init()
+    public var errorOptions: StreamOptions = .init()
     /// An ordered list of steps in order to tear down the child
     /// process in case the parent task is cancelled before
     /// the child proces terminates.
@@ -126,6 +128,18 @@ extension PlatformOptions {
     #endif
 }
 
+extension PlatformOptions {
+    public struct StreamOptions: Sendable {
+        let lowWater: Int?
+        let highWater: Int?
+
+        init(lowWater: Int? = nil, highWater: Int? = nil) {
+            self.lowWater = lowWater
+            self.highWater = highWater
+        }
+    }
+}
+
 extension PlatformOptions: CustomStringConvertible, CustomDebugStringConvertible {
     internal func description(withIndent indent: Int) -> String {
         let indent = String(repeating: " ", count: indent * 4)
@@ -355,8 +369,8 @@ extension Configuration {
                     output: output,
                     error: error,
                     inputPipe: inputPipe.createInputPipe(),
-                    outputPipe: outputPipe.createOutputPipe(),
-                    errorPipe: errorPipe.createOutputPipe()
+                    outputPipe: outputPipe.createOutputPipe(with: platformOptions.outputOptions),
+                    errorPipe: errorPipe.createOutputPipe(with: platformOptions.errorOptions)
                 )
             }
 

Sources/Subprocess/Platforms/Subprocess+Linux.swift

@@ -127,8 +127,8 @@ extension Configuration {
                     output: output,
                     error: error,
                     inputPipe: inputPipe.createInputPipe(),
-                    outputPipe: outputPipe.createOutputPipe(),
-                    errorPipe: errorPipe.createOutputPipe()
+                    outputPipe: outputPipe.createOutputPipe(with: platformOptions.outputOptions),
+                    errorPipe: errorPipe.createOutputPipe(with: platformOptions.errorOptions)
                 )
             }
 
@@ -176,6 +176,8 @@ public struct PlatformOptions: Sendable {
     // Creates a session and sets the process group ID
     // i.e. Detach from the terminal.
     public var createSession: Bool = false
+    public var outputOptions: StreamOptions = .init()
+    public var errorOptions: StreamOptions = .init()
     /// An ordered list of steps in order to tear down the child
     /// process in case the parent task is cancelled before
     /// the child proces terminates.
@@ -197,6 +199,18 @@ public struct PlatformOptions: Sendable {
     public init() {}
 }
 
+extension PlatformOptions {
+    public struct StreamOptions: Sendable {
+        let lowWater: Int?
+        let highWater: Int?
+
+        init(lowWater: Int? = nil, highWater: Int? = nil) {
+            self.lowWater = lowWater
+            self.highWater = highWater
+        }
+    }
+}
+
 extension PlatformOptions: CustomStringConvertible, CustomDebugStringConvertible {
     internal func description(withIndent indent: Int) -> String {
         let indent = String(repeating: " ", count: indent * 4)

Sources/Subprocess/Platforms/Subprocess+Unix.swift

@@ -412,6 +412,7 @@ extension CreatedPipe {
                     }
                 }
             )
+
             writeEnd = .init(
                 dispatchIO,
                 closeWhenDone: writeFileDescriptor.closeWhenDone
@@ -423,7 +424,7 @@ extension CreatedPipe {
         )
     }
 
-    internal func createOutputPipe() -> OutputPipe {
+    internal func createOutputPipe(with options: PlatformOptions.StreamOptions) -> OutputPipe {
         var readEnd: TrackedPlatformDiskIO? = nil
         if let readFileDescriptor = self.readFileDescriptor {
             let dispatchIO: DispatchIO = DispatchIO(
@@ -437,6 +438,15 @@ extension CreatedPipe {
                     }
                 }
             )
+
+            if let lowWater = options.lowWater {
+                dispatchIO.setLimit(lowWater: lowWater)
+            }
+
+            if let highWater = options.highWater {
+                dispatchIO.setLimit(highWater: highWater)
+            }
+
             readEnd = .init(
                 dispatchIO,
                 closeWhenDone: readFileDescriptor.closeWhenDone
@@ -451,6 +461,48 @@ extension CreatedPipe {
 
 // MARK: - TrackedDispatchIO extensions
 extension TrackedDispatchIO {
+    internal func readDataStream(upToLength maxLength: Int) -> AsyncThrowingStream<StreamStatus, Swift.Error> {
+        return AsyncThrowingStream<StreamStatus, Swift.Error> { continuation in
+            self.dispatchIO.read(
+                offset: 0,
+                length: maxLength,
+                queue: .global()
+            ) { done, data, error in
+                if error != 0 {
+                    continuation.finish(throwing: SubprocessError(
+                        code: .init(.failedToReadFromSubprocess),
+                        underlyingError: .init(rawValue: error)
+                    ))
+                    return
+                }
+
+                // Treat empty data and nil as the same
+                let buffer = data.map { $0.isEmpty ? nil : $0 } ?? nil
+                let status: StreamStatus
+
+                switch (buffer, done) {
+                case (.some(let data), false):
+                    status = .data(SequenceOutput.Buffer(data: data))
+
+                case (.some(let data), true):
+                    status = .endOfStream(SequenceOutput.Buffer(data: data))
+
+                case (nil, false):
+                    return
+
+                case (nil, true):
+                    status = .endOfFile
+                }
+
+                continuation.yield(status)
+
+                if done {
+                    continuation.finish()
+                }
+            }
+        }
+    }
+
 #if SubprocessSpan
     @available(SubprocessSpan, *)
 #endif

Sources/Subprocess/Platforms/Subprocess+Windows.swift

@@ -1043,6 +1043,70 @@ extension CreatedPipe {
 }
 
 extension TrackedFileDescriptor {
+    internal func readDataStream(upToLength maxLength: Int) -> AsyncThrowingStream<StreamStatus, Swift.Error> {
+        return AsyncThrowingStream<StreamStatus, Swift.Error> { continuation in
+            do {
+                var totalBytesRead: Int = 0
+
+                while totalBytesRead < maxLength {
+                    let values = try [UInt8](
+                        unsafeUninitializedCapacity: maxLength
+                    ) { buffer, initializedCount in
+                        guard let baseAddress = buffer.baseAddress else {
+                            initializedCount = 0
+                            return
+                        }
+
+                        var bytesRead: DWORD = 0
+                        let readSucceed = ReadFile(
+                            self.fileDescriptor.platformDescriptor,
+                            UnsafeMutableRawPointer(mutating: baseAddress),
+                            DWORD(maxLength - totalBytesRead),
+                            &bytesRead,
+                            nil
+                        )
+
+                        if !readSucceed {
+                            // Windows throws ERROR_BROKEN_PIPE when the pipe is closed
+                            let error = GetLastError()
+                            if error == ERROR_BROKEN_PIPE {
+                                // We are done reading
+                                initializedCount = 0
+                            } else {
+                                initializedCount = 0
+                                throw SubprocessError(
+                                    code: .init(.failedToReadFromSubprocess),
+                                    underlyingError: .init(rawValue: error)
+                                )
+                            }
+                        } else {
+                            // We successfully read the current round
+                            initializedCount += Int(bytesRead)
+                        }
+                    }
+
+                    if values.count > 0 {
+                        totalBytesRead += values.count
+
+                        if totalBytesRead >= maxLength {
+                            continuation.yield(.endOfStream(SequenceOutput.Buffer(data: values)))
+                            continuation.finish()
+                            return
+                        } else {
+                            continuation.yield(.data(SequenceOutput.Buffer(data: values)))
+                        }
+                    } else {
+                        continuation.yield(.endOfFile)
+                        continuation.finish()
+                        return
+                    }
+                }
+            } catch {
+                continuation.finish(throwing: error)
+            }
+        }
+    }
+
     internal func readChunk(upToLength maxLength: Int) async throws -> SequenceOutput.Buffer? {
         return try await withCheckedThrowingContinuation { continuation in
             self.readUntilEOF(

Tests/SubprocessTests/SubprocessTests+Unix.swift

@@ -665,6 +665,48 @@ extension SubprocessUnixTests {
         #expect(catResult.terminationStatus.isSuccess)
         #expect(catResult.standardError == expected)
     }
+
+    @Test func testSlowDripRedirectedOutputRedirectToSequence() async throws {
+        let threshold: Double = 0.5
+
+        let script = """
+        echo "DONE"
+        sleep \(threshold)
+        """
+
+        var platformOptions = PlatformOptions()
+        platformOptions.outputOptions = .init(lowWater: 0)
+
+        let start = ContinuousClock().now
+
+        let catResult = try await Subprocess.run(
+            .path("/bin/bash"),
+            arguments: ["-c", script],
+            platformOptions: platformOptions,
+            output: .sequence,
+            error: .discarded,
+            body: { (execution, _) in
+                for try await chunk in execution.standardOutput {
+                    let string = chunk.withUnsafeBytes { String(decoding: $0, as: UTF8.self) }
+
+                    if string.hasPrefix("DONE") {
+                        let end = ContinuousClock().now
+
+                        if (end - start) > .seconds(threshold) {
+                            return "Failure"
+
+                        } else {
+                            return "Success"
+                        }
+                    }
+                }
+
+                return "Failure"
+            }
+        )
+        #expect(catResult.terminationStatus.isSuccess)
+        #expect(catResult.value == "Success")
+    }
 }
 
 // MARK: - PlatformOption Tests