From 4bfe236d53beb206fb2b514b107ac8f1dc16af9f Mon Sep 17 00:00:00 2001 From: Soeren Schoenbrod Date: Mon, 13 Jul 2026 14:22:16 +0000 Subject: [PATCH] fix: don't drop items when close races with the final put! The consumer-side drain checks in take!, take!(ch, output), wait and iterate read the head/tail indices, observed an empty buffer, and *then* separately read the closed flag. Because close() is ordered after every put! but the empty-snapshot and closed reads are distinct operations, a producer's final put! could become visible to the consumer only after its close(): the consumer would observe the channel as closed-and-empty and throw / end iteration, silently dropping the last item(s). Fix the drain checks to re-read the indices after observing closed: since close() is ordered after every put!, a re-read then observes all published items, so throwing only happens when the channel is genuinely drained. Also simplify iterate to always attempt the (now race-free) take! instead of guarding it with the racy `isopen(ch) || isready(ch)` snapshot; take! already drains before throwing and propagates bound-task exceptions. Add a threaded regression test that produces N items then closes and asserts the consumer iterates all N, over many trials. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/PipeChannels.jl | 59 +++++++++++++++++++++++---------------------- test/runtests.jl | 31 ++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 29 deletions(-) diff --git a/src/PipeChannels.jl b/src/PipeChannels.jl index d1112d9..78e4ea4 100644 --- a/src/PipeChannels.jl +++ b/src/PipeChannels.jl @@ -218,8 +218,10 @@ function Base.wait(ch::PipeChannel) if ch.head[] != ch.tail[] return nothing end - # Check if closed and empty - if ch.closed[] + # Check if closed and empty. Re-read the indices after observing `closed` + # (see take!): close() is ordered after every put!, so a re-read here sees + # any item published just before close(), avoiding a spurious throw. + if ch.closed[] && ch.head[] == ch.tail[] check_closed_and_throw(ch) end # Spin-wait @@ -306,7 +308,11 @@ function Base.take!(ch::PipeChannel{T}) where {T} # Check if buffer is empty if tail == head - if ch.closed[] + # Re-read head after observing `closed`. `close()` is ordered after + # every `put!`, so if the channel still looks empty here it is truly + # drained. Without the re-read a `put!` made visible just before + # `close()` could be skipped, dropping the last item(s). + if ch.closed[] && tail == ch.head[] check_closed_and_throw(ch) end # Spin-wait @@ -336,31 +342,24 @@ Note: The `@inline` annotation is critical for avoiding heap allocation of the returned `(value, nothing)` tuple when `T` is not an isbits type. """ @inline function Base.iterate(ch::PipeChannel{T}, state=nothing) where {T} - # Fast path: if channel is open or has data, try to take - if isopen(ch) || isready(ch) - try - value = take!(ch) - return (value, nothing) - catch e - if e isa InvalidStateException - # Check if there's a stored exception we should throw instead - # This handles the case where take! threw InvalidStateException but - # there's a more specific exception from a bound task - stored = ch.excp - if stored !== nothing && !isa(stored, InvalidStateException) - throw(stored) - end - return nothing - end - rethrow() - end - else - # Channel is closed and empty - check for stored exception from bind - # This matches Julia's Channel behavior: propagate bound task exceptions - e = ch.excp - if e !== nothing && !isa(e, InvalidStateException) - throw(e) - end + # Always attempt `take!` rather than guarding with `isopen(ch) || isready(ch)`. + # + # `take!` is race-free: it re-reads the head/tail indices in its own loop, so it + # returns a still-buffered value even after the channel has been closed, and only + # throws once the channel is both closed and drained. It also throws a bound + # task's stored exception (via `check_closed_and_throw`) when one is present. + # + # The previous `isopen(ch) || isready(ch)` guard took a racy snapshot: if a + # producer's final `put!` became visible to the consumer only after its `close()`, + # the consumer could observe the channel as closed-and-empty and return `nothing`, + # silently dropping the last item(s). + try + value = take!(ch) + return (value, nothing) + catch e + # Plain closed-and-drained: end iteration. A bound task's real exception is + # rethrown here (take! threw it directly), matching Julia's Channel behavior. + e isa InvalidStateException || rethrow() return nothing end end @@ -534,7 +533,9 @@ function Base.take!(ch::PipeChannel{T}, output::AbstractVector{T}) where {T} # Check if buffer is empty if tail == head - if ch.closed[] + # Re-read head after observing `closed` (see single-item take!): this + # avoids dropping items published just before close() became visible. + if ch.closed[] && tail == ch.head[] check_closed_and_throw(ch) end # Spin-wait for data diff --git a/test/runtests.jl b/test/runtests.jl index b5a4959..f014844 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -448,4 +448,35 @@ using PipeChannels end end + @testset "iterate does not drop items when close races with final put!" begin + # Regression test: previously `iterate` guarded `take!` with + # `isopen(ch) || isready(ch)`, which could observe the channel as + # closed-and-empty and drop the last item(s) when a producer's final + # `put!` became visible only after its `close()`. Stress the window by + # producing on one thread and consuming (iterating) on another over many + # trials. + if Threads.nthreads() >= 2 + n = 3 + failures = 0 + for _ in 1:5000 + ch = PipeChannel{Int}(16) + producer = Threads.@spawn begin + for i in 1:n + put!(ch, i) + end + close(ch) + end + got = Int[] + for v in ch + push!(got, v) + end + wait(producer) + got == collect(1:n) || (failures += 1) + end + @test failures == 0 + else + @info "Skipping iterate/close race test (need at least 2 threads)" + end + end + end