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.

Resolves swiftlang#39

Author: rdingman

Sources/Subprocess/AsyncBufferSequence.swift

@@ -15,6 +15,8 @@
 @preconcurrency import SystemPackage
 #endif
 
+internal import Dispatch
+
 #if SubprocessSpan
 @available(SubprocessSpan, *)
 #endif
@@ -27,38 +29,100 @@ public struct AsyncBufferSequence: AsyncSequence, Sendable {
         public typealias Element = SequenceOutput.Buffer
 
         private let diskIO: TrackedPlatformDiskIO
+        private let bufferSize: Int
         private var buffer: [UInt8]
         private var currentPosition: Int
         private var finished: Bool
+        private var streamIterator: AsyncThrowingStream<StreamStatus, Swift.Error>.AsyncIterator
 
-        internal init(diskIO: TrackedPlatformDiskIO) {
+        internal init(diskIO: TrackedPlatformDiskIO, bufferSize: Int) {
             self.diskIO = diskIO
+            self.bufferSize = bufferSize
             self.buffer = []
             self.currentPosition = 0
             self.finished = false
+            self.streamIterator = Self.createDataStream(with: diskIO.dispatchIO, bufferSize: bufferSize).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 = Self.createDataStream(with: diskIO.dispatchIO, bufferSize: bufferSize).makeAsyncIterator()
+                    return data
+
+                case .endOfFile:
+                    try self.diskIO.safelyClose()
+                    return nil
+                }
+            } else {
                 try self.diskIO.safelyClose()
                 return nil
             }
-            return data
+        }
+
+        private enum StreamStatus {
+            case data(SequenceOutput.Buffer)
+            case endOfStream(SequenceOutput.Buffer)
+            case endOfFile
+        }
+
+        private static func createDataStream(with dispatchIO: DispatchIO, bufferSize: Int) -> AsyncThrowingStream<StreamStatus, Swift.Error> {
+            return AsyncThrowingStream<StreamStatus, Swift.Error> { continuation in
+                dispatchIO.read(
+                    offset: 0,
+                    length: bufferSize,
+                    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()
+                    }
+                }
+            }
         }
     }
 
     private let diskIO: TrackedPlatformDiskIO
+    private let bufferSize: Int
 
-    internal init(diskIO: TrackedPlatformDiskIO) {
+    internal init(diskIO: TrackedPlatformDiskIO, bufferSize: Int) {
         self.diskIO = diskIO
+        self.bufferSize = bufferSize
     }
 
     public func makeAsyncIterator() -> Iterator {
-        return Iterator(diskIO: self.diskIO)
+        return Iterator(diskIO: self.diskIO, bufferSize: bufferSize)
     }
 }
 

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,7 +109,16 @@ extension Execution where Output == SequenceOutput {
         else {
             fatalError("The standard output has already been consumed")
         }
-        return AsyncBufferSequence(diskIO: readFd)
+
+        if let lowWater = output.lowWater {
+            readFd.dispatchIO.setLimit(lowWater: lowWater)
+        }
+
+        if let highWater = output.highWater {
+            readFd.dispatchIO.setLimit(highWater: highWater)
+        }
+
+        return AsyncBufferSequence(diskIO: readFd, bufferSize: output.bufferSize)
     }
 }
 
@@ -122,15 +133,24 @@ 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")
         }
-        return AsyncBufferSequence(diskIO: readFd)
+
+        if let lowWater = error.lowWater {
+            readFd.dispatchIO.setLimit(lowWater: lowWater)
+        }
+
+        if let highWater = error.highWater {
+            readFd.dispatchIO.setLimit(highWater: highWater)
+        }
+
+        return AsyncBufferSequence(diskIO: readFd, bufferSize: error.bufferSize)
     }
 }
 

Sources/Subprocess/IO/Output.swift

@@ -210,8 +210,15 @@ public struct BytesOutput: OutputProtocol {
 #endif
 public struct SequenceOutput: OutputProtocol {
     public typealias OutputType = Void
-
-    internal init() {}
+    internal let lowWater: Int?
+    internal let highWater: Int?
+    internal let bufferSize: Int
+
+    internal init(lowWater: Int? = nil, highWater: Int? = nil, bufferSize: Int = readBufferSize) {
+        self.lowWater = lowWater
+        self.highWater = highWater
+        self.bufferSize = bufferSize
+    }
 }
 
 #if SubprocessSpan
@@ -284,6 +291,9 @@ extension OutputProtocol where Self == SequenceOutput {
     /// to the `.standardOutput` (or `.standardError`) property
     /// of `Execution` as `AsyncSequence<Data>`.
     public static var sequence: Self { .init() }
+    public static func sequence(lowWater: Int? = nil, highWater: Int? = nil, bufferSize: Int? = nil) -> Self {
+        .init(lowWater: lowWater, highWater: highWater, bufferSize: bufferSize ?? readBufferSize)
+    }
 }
 
 // MARK: - Span Default Implementations

Tests/SubprocessTests/SubprocessTests+Unix.swift

@@ -665,6 +665,44 @@ 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)
+        """
+
+        let start = ContinuousClock().now
+
+        let catResult = try await Subprocess.run(
+            .path("/bin/bash"),
+            arguments: ["-c", script],
+            output: .sequence(lowWater: 0),
+            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