add Threads.foreach for convenient multithreaded Channel consumption (

JuliaLang#34543)

Co-authored-by: Takafumi Arakaki <[email protected]>
Co-authored-by: Alex Arslan <[email protected]>
Co-authored-by: Valentin Churavy <[email protected]>

NEWS.md

@@ -14,7 +14,6 @@ Language changes
 Compiler/Runtime improvements
 -----------------------------
 
-
 * All platforms can now use `@executable_path` within `jl_load_dynamic_library()`.
  This allows executable-relative paths to be embedded within executables on all
  platforms, not just MacOS, which the syntax is borrowed from. ([#35627])
@@ -33,14 +32,17 @@ Build system changes
 
 New library functions
 ---------------------
+
 * New function `Base.kron!` and corresponding overloads for various matrix types for performing Kronecker product in-place. ([#31069]).
+* New function `Base.Threads.foreach(f, channel::Channel)` for multithreaded `Channel` consumption. ([#34543]).
 
 New library features
 --------------------
 
 
 Standard library changes
 ------------------------
+
 * The `nextprod` function now accepts tuples and other array types for its first argument ([#35791]).
 * The function `isapprox(x,y)` now accepts the `norm` keyword argument also for numeric (i.e., non-array) arguments `x` and `y` ([#35883]).
 * `view`, `@view`, and `@views` now work on `AbstractString`s, returning a `SubString` when appropriate ([#35879]).

base/Base.jl

@@ -223,6 +223,7 @@ include("threads.jl")
 include("lock.jl")
 include("channels.jl")
 include("task.jl")
+include("threads_overloads.jl")
 include("weakkeydict.jl")
 
 # Logging

base/threadingconstructs.jl

@@ -180,3 +180,11 @@ macro spawn(expr)
  end
  end
 end
+
+# This is a stub that can be overloaded for downstream structures like `Channel`
+function foreach end
+
+# Scheduling traits that can be employed for downstream overloads
+abstract type AbstractSchedule end
+struct StaticSchedule <: AbstractSchedule end
+struct FairSchedule <: AbstractSchedule end

base/threads_overloads.jl

@@ -0,0 +1,51 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
+"""
+ Threads.foreach(f, channel::Channel;
+ schedule::Threads.AbstractSchedule=Threads.FairSchedule(),
+ ntasks=Threads.nthreads())
+
+Similar to `foreach(f, channel)`, but iteration over `channel` and calls to
+`f` are split across `ntasks` tasks spawned by `Threads.@spawn`. This function
+will wait for all internally spawned tasks to complete before returning.
+
+If `schedule isa FairSchedule`, `Threads.foreach` will attempt to spawn tasks in a
+manner that enables Julia's scheduler to more freely load-balance work items across
+threads. This approach generally has higher per-item overhead, but may perform
+better than `StaticSchedule` in concurrence with other multithreaded workloads.
+
+If `schedule isa StaticSchedule`, `Threads.foreach` will spawn tasks in a manner
+that incurs lower per-item overhead than `FairSchedule`, but is less amenable
+to load-balancing. This approach thus may be more suitable for fine-grained,
+uniform workloads, but may perform worse than `FairSchedule` in concurrence
+with other multithreaded workloads.
+
+!!! compat "Julia 1.6"
+ This function requires Julia 1.6 or later.
+"""
+function Threads.foreach(f, channel::Channel;
+ schedule::Threads.AbstractSchedule=Threads.FairSchedule(),
+ ntasks=Threads.nthreads())
+ apply = _apply_for_schedule(schedule)
+ stop = Threads.Atomic{Bool}(false)
+ @sync for _ in 1:ntasks
+ Threads.@spawn try
+ for item in channel
+ $apply(f, item)
+ # do `stop[] && break` after `f(item)` to avoid losing `item`.
+ # this isn't super comprehensive since a task could still get
+ # stuck on `take!` at `for item in channel`. We should think
+ # about a more robust mechanism to avoid dropping items. See also:
+ # https://github.com/JuliaLang/julia/pull/34543#discussion_r422695217
+ stop[] && break
+ end
+ catch
+ stop[] = true
+ rethrow()
+ end
+ end
+ return nothing
+end
+
+_apply_for_schedule(::Threads.StaticSchedule) = (f, x) -> f(x)
+_apply_for_schedule(::Threads.FairSchedule) = (f, x) -> wait(Threads.@spawn f(x))

test/threads_exec.jl

@@ -845,3 +845,38 @@ fib34666(x) =
  f(x)
  end
 @test fib34666(25) == 75025
+
+function jitter_channel(f, k, delay, ntasks, schedule)
+ x = Channel(ch -> foreach(i -> put!(ch, i), 1:k), 1)
+ y = Channel(k) do ch
+ g = i -> begin
+ iseven(i) && sleep(delay)
+ put!(ch, f(i))
+ end
+ Threads.foreach(g, x; schedule=schedule, ntasks=ntasks)
+ end
+ return y
+end
+
+@testset "Threads.foreach(f, ::Channel)" begin
+ k = 50
+ delay = 0.01
+ expected = sin.(1:k)
+ ordered_fair = collect(jitter_channel(sin, k, delay, 1, Threads.FairSchedule()))
+ ordered_static = collect(jitter_channel(sin, k, delay, 1, Threads.StaticSchedule()))
+ @test expected == ordered_fair
+ @test expected == ordered_static
+
+ unordered_fair = collect(jitter_channel(sin, k, delay, 10, Threads.FairSchedule()))
+ unordered_static = collect(jitter_channel(sin, k, delay, 10, Threads.StaticSchedule()))
+ @test expected != unordered_fair
+ @test expected != unordered_static
+ @test Set(expected) == Set(unordered_fair)
+ @test Set(expected) == Set(unordered_static)
+
+ ys = Channel() do ys
+ inner = Channel(xs -> foreach(i -> put!(xs, i), 1:3))
+ Threads.foreach(x -> put!(ys, x), inner)
+ end
+ @test sort!(collect(ys)) == 1:3
+end