Cache chunks per-process

Author: jpsamaroo

src/chunks.jl

@@ -63,6 +63,9 @@ function unrelease(c::Chunk{<:Any,DRef})
 end
 unrelease(c::Chunk) = c
 
+Base.:(==)(c1::Chunk, c2::Chunk) = c1.handle == c2.handle
+Base.hash(c::Chunk, x::UInt64) = hash(c.handle, x)
+
 collect_remote(chunk::Chunk, network) =
     move(chunk.processor, OSProc(), poolget(chunk.handle; network=network); network=network)
 function collect(ctx::Context, chunk::Chunk; options=nothing)

src/sch/Sch.jl

@@ -441,6 +441,8 @@ function Base.show(io::IO, se::SchedulingException)
     print(io, "SchedulingException ($(se.reason))")
 end
 
+const CHUNK_CACHE = Dict{Chunk,Any}()
+
 function schedule!(ctx, state, procs=procs_to_use(ctx))
     lock(state.lock) do
         safepoint(state)
@@ -787,7 +789,17 @@ function do_task(to_proc, extra_util, thunk_id, f, data, send_result, persist, c
         fetch_report.(map(Iterators.zip(data,ids)) do (x, id)
             @async begin
                 @dbg timespan_start(ctx, :move, (thunk_id, id), (f, id))
-                x = move(to_proc, x; network=network)
+                x = if x isa Chunk
+                    if haskey(CHUNK_CACHE, x)
+                        CHUNK_CACHE[x]
+                    else
+                        _x = move(to_proc, x; network=network)
+                        CHUNK_CACHE[x] = _x
+                        _x
+                    end
+                else
+                    move(to_proc, x; network=network)
+                end
                 @dbg timespan_end(ctx, :move, (thunk_id, id), (f, id))
                 return x
             end