Iteration stalling for custom Base.iterate methods using an unbuffered Channel

[depial]

In the code below, my intention was to be able to iterate across a tree, while using a Channel to handle the tree traversal. This implementation works fine if iterated "manually" with a while loop or if print/@show is used in a for loop, but using other functions (e.g. push!) in a for loop stalls after the first item.

After digging around, I've found that this problem can be solved with:

1. a call to yield in one of the Base.iterate methods
2. with a yield placed in the finally clause of the take_unbuffered method from julia/base/channels.jl.

I'm unsure if the second fix would have other, undesirable, ramifications, but there seems to be a yield in the complementary put_unbuffered, whose purpose is to:

immediately give taker a chance to run, but don't block the current task

and this seems to be what is needed in take_unbuffered to make iteration work as expected.

MWE:

abstract type AbstractNode end

function traverse(node::AbstractNode, c::Channel)
    put!(c, node.data)
    foreach(branch -> traverse(branch, c), branches(node))
end

function Base.iterate(node::AbstractNode)
    state = Channel(c -> traverse(node, c))
    isready(state) ? (take!(state), state) : nothing
end
function Base.iterate(node::AbstractNode, state)
    # yield()
    isready(state) ? (take!(state), state) : nothing
end

Base.iterate(::AbstractNode, ::Nothing) = nothing
Base.IteratorSize(::AbstractNode) = Base.SizeUnknown()

struct Leaf <: AbstractNode
    data::Int
end

struct Node <: AbstractNode
    data::Int
    leftbranch::Union{Node, Leaf}
    rightbranch::Union{Node, Leaf}
end

data(node::Union{Node, Leaf}) = node.data
branches(node::Node) = (node.leftbranch, node.rightbranch)
branches(node::Leaf) = ()

The following code will return the full array (e.g. [1, 2, 3, 4, 5, 6, 7]):

tree = Node(1, Node(2, Leaf(3), Leaf(4)), Node(5, Leaf(6), Leaf(7)))

arr = []
next = iterate(tree)
while next !== nothing
    (item, state) = next
    push!(arr, item)
    next = iterate(tree, state)
end
arr

The above code is modified from that in Interfaces, where it is stated that it is a translation of the following for loop:

tree = Node(1, Node(2, Leaf(3), Leaf(4)), Node(5, Leaf(6), Leaf(7)))

arr = []
for item in itr 
    push!(arr, item) 
end
arr

which stalls, and only returns an array with the first item (e.g. [5]). This difference in behavior was the reason for opening this issue.

Some extra examples:

julia> tree = Node(1, Node(2, Leaf(3), Leaf(4)), Node(5, Leaf(6), Leaf(7)));

julia> [n for n in tree]
1-element Vector{Any}:
 1

julia> foreach(n -> println(n), tree)
 1
 2
 3
 4
 5
 6
 7

julia> [n for n in Channel(c -> traverse(tree, c))]
7-element Vector{Int64}:
 1
 2
 3
 4
 5
 6
 7

Apparently, one potential cause for this is the for loop being lazily JIT-compiled, where the calls to println will keep iteration flowing, but calls to push! do not. This idea is what led to me using yield to find a work-around.

Specifically, the fixes I've found:

1. Uncommenting the call to yield() in the MWE's Base.iterate(node::AbstractNode, state) allows for full iteration.
2. The following modification to take_unbuffered, by adding a call to yield() in the finally clause, also allows for full iteration:

function take_unbuffered(c::Channel{T}) where T
    lock(c)
    try
        check_channel_state(c)
        notify(c.cond_put, nothing, false, false)
        return wait(c.cond_take)::T
    finally
        unlock(c)
        yield() # immediately give getter a chance to run, but don't block the current task
    end
end

(where I've added a complementary comment to that found next to the yield() in the put_unbuffered implementation).

Again, I'm unsure if this is an optimal fix or if it could cause unwanted behavior elsewhere, but I believe the issue of the discrepancy between the above while loop and for loop examples would remain.

[jakobnissen]

I believe this is an implementation mistake in your iteration and not an issue with Julia or its Channel implementation. Your definition of iterate is defined as: isready(state) ? (take!(state), state) : nothing. So, if, at any point, the channel is empty, the iteration stops and does not wait for more elements.
And there is no guarantee of when exactly the scheduler will run the task that runs traverse, and fills more elements into the channel, and when the scheduler will run iterate that removes elements from the channel.
Putting in yield in iterate seems to solve the problem, because it allows the scheduler to switch to traverse. But it's not a real fix, because there is still no guarantee the scheduler will actually switch task.

Your real solution here is for iterate to be defined like so:

function Base.iterate(node::AbstractNode)
    state = Channel(c -> traverse(node, c))
    iterate(node, state)
end

function Base.iterate(node::AbstractNode, state::Channel)
    # This iterate call on the channel will block if there are no more elements,
    # and will return `nothing` if there are no more elements AND the channel
    # is closed, which happens when `traverse` has yielded all its elements.
    y = iterate(state)
    isnothing(y) ? y : (first(y), state)
end

The reason this works is that iterate on Channel will take a value if the channel is open or there are elements less, otherwise it will return nothing. Which is exactly the same behaviour you want.

Closing as not planned, but please do write if I'm mistaken.

[depial]

I was opening the issue more because of the behavior difference in the explicit iteration of the while loop and the implicit iteration of the for loop. I was under the impression that these two were supposed to always give the same results.

The other reason that I opened this is because I found that modifying take_unbuffered in a (more or less) symmetric way to put_ubuffered (from jula/base/channels.jl) seemed to fix my iteration problem. The first "fix" simply led me to this second one.

I guess my questions would be:

1. Are there cases when the explicit iteration in the while loop can (and/or should) differ from that in the for loop?
2. Does my modification to take_unbuffered cause undesirable behavior in unbuffered Channels? (e.g. is there still no guarantee it will switch tasks?)

If yes to these, then I feel we can rightly close the issue. If not, it may be reasonable to revisit it.

[jakobnissen]

The for and the while loop forms are guaranteed to be semantically the same. However, there is no guarantee that it produces the exact same machine code (indeed, on my computer, they do produce different machine code), and therefore the two may differ in non-semantic behaviour. Here, by semantics, I mean: "what does Julia guarantee happens when running the code?". Here, the timings and yield-patterns is not covered by the semantics.

Adding a yield() in take_unbuffered still does not guarantee that the producer task will run. It's permitted for the scheduler to simply re-schedule the consumer task. The yield does cause undesirable behaviour in the sense that yield adds overhead which is not normally needed.

[depial]

Thank you for the explanations