Define iterate for RemoteChannel (JuliaLang/julia#48515)

Author: quinnj

src/Distributed.jl

@@ -7,7 +7,7 @@ module Distributed
 
 # imports for extension
 import Base: getindex, wait, put!, take!, fetch, isready, push!, length,
-             hash, ==, kill, close, isopen, showerror
+             hash, ==, kill, close, isopen, showerror, iterate, IteratorSize
 
 # imports for use
 using Base: Process, Semaphore, JLOptions, buffer_writes, @async_unwrap,

src/remotecall.jl

@@ -778,3 +778,23 @@ function getindex(r::RemoteChannel, args...)
     end
     return remotecall_fetch(getindex, r.where, r, args...)
 end
+
+function iterate(c::RemoteChannel, state=nothing)
+    if isopen(c) || isready(c)
+        try
+            return (take!(c), nothing)
+        catch e
+            if isa(e, InvalidStateException) ||
+                (isa(e, RemoteException) &&
+                isa(e.captured.ex, InvalidStateException) &&
+                e.captured.ex.state === :closed)
+                return nothing
+            end
+            rethrow()
+        end
+    else
+        return nothing
+    end
+end
+
+IteratorSize(::Type{<:RemoteChannel}) = SizeUnknown()

test/distributed_exec.jl

@@ -456,6 +456,32 @@ function test_iteration(in_c, out_c)
 end
 
 test_iteration(Channel(10), Channel(10))
+test_iteration(RemoteChannel(() -> Channel(10)), RemoteChannel(() -> Channel(10)))
+
+@everywhere function test_iteration_take(ch)
+    count = 0
+    for x in ch
+        count += 1
+    end
+    return count
+end
+
+@everywhere function test_iteration_put(ch, total)
+    for i in 1:total
+        put!(ch, i)
+    end
+    close(ch)
+end
+
+let ch = RemoteChannel(() -> Channel(1))
+    @async test_iteration_put(ch, 10)
+    @test 10 == @fetchfrom id_other test_iteration_take(ch)
+    # now reverse
+    ch = RemoteChannel(() -> Channel(1))
+    @spawnat id_other test_iteration_put(ch, 10)
+    @test 10 == test_iteration_take(ch)
+end
+
 # make sure exceptions propagate when waiting on Tasks
 @test_throws CompositeException (@sync (@async error("oops")))
 try