Skip to content

Commit e830403

Browse files
committed
[AsyncStreaming] Skip empty buffers in AsyncReader.forEachBuffer
When a reader signals end-of-stream with an empty buffer, body was still invoked with that empty buffer. Skip the body call when the buffer is empty.
1 parent 8ee3d2b commit e830403

3 files changed

Lines changed: 41 additions & 19 deletions

File tree

Sources/AsyncStreaming/AsyncReader/AsyncReader+forEach.swift

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111

1212
#if UnstableAsyncStreaming && compiler(>=6.4)
1313

14-
import ContainersPreview
14+
public import ContainersPreview
1515

1616
// swift-format-ignore: AmbiguousTrailingClosureOverload
1717
@available(macOS 10.14.4, iOS 12.2, watchOS 5.2, tvOS 12.2, *)
@@ -47,7 +47,9 @@ extension AsyncReader where Self: ~Copyable, Self: ~Escapable {
4747
var done = false
4848
while !done {
4949
try await self.read { (next, finalElement) throws(Failure) -> Void in
50-
try await body(&next)
50+
if !next.isEmpty {
51+
try await body(&next)
52+
}
5153
if let finalElement {
5254
final = finalElement
5355
done = true
@@ -79,21 +81,15 @@ extension AsyncReader where Self: ~Copyable, Self: ~Escapable {
7981
public consuming func forEachBuffer(
8082
body: (inout Buffer) async -> Void
8183
) async -> FinalElement where ReadFailure == Never {
82-
var finalElement: FinalElement? = nil
83-
while finalElement == nil {
84-
do {
85-
try await self.read { (next, final) -> Void in
86-
await body(&next)
87-
if let final {
88-
finalElement = final
89-
}
90-
}
91-
} catch {
92-
fatalError()
84+
do {
85+
let final: FinalElement? = try await self.forEachBuffer { (buffer) async throws(Never) -> Void in
86+
await body(&buffer)
9387
}
88+
// The force-unwrap is safe since final element must be set at this point
89+
return final!
90+
} catch {
91+
fatalError()
9492
}
95-
// The force-unwrap is safe since final element must be set at this point
96-
return finalElement!
9793
}
9894
}
9995
#endif

Tests/AsyncStreamingTests/AsyncReader/AsyncReader+forEachTests.swift

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,34 @@ struct AsyncReaderforEachBufferTests {
5353
callCount += 1
5454
}
5555

56-
// The reader still emits a terminal call (with an empty buffer + finalElement).
57-
#expect(callCount == 1)
56+
// The reader signals end-of-stream with an empty buffer; body should not be called.
57+
#expect(callCount == 0)
58+
}
59+
60+
@Test
61+
func forEachBufferDoesNotInvokeBodyWithEmptyTerminalBuffer() async throws {
62+
let reader = UniqueArrayAsyncReader(storage: UniqueArray(capacity: 0, copying: []))
63+
var observedCounts: [Int] = []
64+
65+
_ = await reader.forEachBuffer { buffer in
66+
observedCounts.append(buffer.count)
67+
}
68+
69+
#expect(observedCounts.allSatisfy { $0 > 0 })
70+
}
71+
72+
@Test
73+
func forEachBufferThrowingDoesNotInvokeBodyWithEmptyTerminalBuffer() async throws {
74+
enum TestError: Error {}
75+
76+
let reader = UniqueArrayAsyncReader(storage: UniqueArray(capacity: 0, copying: []))
77+
var observedCounts: [Int] = []
78+
79+
_ = try await reader.forEachBuffer { (buffer) throws(TestError) -> Void in
80+
observedCounts.append(buffer.count)
81+
}
82+
83+
#expect(observedCounts.allSatisfy { $0 > 0 })
5884
}
5985

6086
@Test

Tests/AsyncStreamingTests/DuplexChannel/DuplexAsyncChannelTests.swift

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -449,7 +449,7 @@ struct DuplexAsyncChannelTests {
449449
// it through a generic function that only sees the protocol.
450450
func finishViaProtocol<W: CallerAsyncWriter & ~Copyable>(
451451
_ writer: consuming W,
452-
finalElement: consuming W.FinalElement?
452+
finalElement: consuming W.FinalElement
453453
) async throws(W.WriteFailure) where W.WriteElement == Int {
454454
var buf = UniqueArray<Int>(minimumCapacity: 3)
455455
buf.append(7)
@@ -466,7 +466,7 @@ struct DuplexAsyncChannelTests {
466466
_ = readerA
467467
_ = writerB
468468

469-
try await finishViaProtocol(writerA, finalElement: .some(()))
469+
try await finishViaProtocol(writerA, finalElement: ())
470470

471471
var collected: [Int] = []
472472
var sawFinal = false

0 commit comments

Comments
 (0)