From d23b8215469a242d87fa7f5faeee783c2b667e40 Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Sun, 12 Jul 2026 13:13:02 -0700 Subject: [PATCH] Add DEvent and DFuture --- src/MemPool.jl | 2 + src/devent.jl | 209 +++++++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 148 +++++++++++++++++++++++++++++++++ 3 files changed, 359 insertions(+) create mode 100644 src/devent.jl diff --git a/src/MemPool.jl b/src/MemPool.jl index ab3e309..a694ed9 100644 --- a/src/MemPool.jl +++ b/src/MemPool.jl @@ -4,6 +4,7 @@ import Preferences: @load_preference, @set_preferences! using Serialization, Sockets, Random import Serialization: serialize, deserialize export DRef, FileRef, poolset, poolget, mmwrite, mmread, cleanup +export DEvent, DFuture import .Threads: ReentrantLock using ScopedValues @@ -72,6 +73,7 @@ include("read_write_lock.jl") include("stack.jl") include("clock.jl") include("datastore.jl") +include("devent.jl") """ approx_size(d) diff --git a/src/devent.jl b/src/devent.jl new file mode 100644 index 0000000..5a27ad5 --- /dev/null +++ b/src/devent.jl @@ -0,0 +1,209 @@ +# Distributed event (`DEvent`) and a `Future`-alike (`DFuture`) built on top of +# MemPool's `DRef` machinery. +# +# Motivation: `Distributed.Future`/`RemoteChannel` are unsafe under concurrent +# access from multiple threads within the same process. Every `put!`/`fetch`/ +# `wait` mutates process-global tables (`client_refs`, `PGRP.refs`) via +# `lookup_ref`/`del_client`, and `fetch` auto-deletes the backing ref. Under +# heavy multithreaded access (e.g. Dagger's parallel datadeps scheduling) this +# races: a `lookup_ref` that runs after the ref has been deleted silently +# fabricates a fresh, never-fulfilled `RemoteValue`, and the waiter blocks +# forever. +# +# `DEvent`/`DFuture` avoid that entirely: +# - Readiness is signalled by a plain `Base.Event` living on the owner worker, +# which is fully thread-safe. Same-process waits/notifies (the common case) +# never touch any Distributed machinery. +# - Lifetime is managed by a backing `DRef`, reusing MemPool's battle-tested +# distributed refcounting + serialization. The owner-side state is dropped via +# the `DRef`'s `destructor` once the last reference (local or remote) is gone. +# - Cross-worker operations are stateless RPCs (`remotecall_fetch`/ +# `remotecall_wait`) against the owner; there is no write-once ref that gets +# deleted out from under a concurrent reader. +# - `DFuture`'s local value cache is never serialized: it's dropped (reset to +# `nothing`) whenever a `DFuture` crosses the wire, so that an +# already-fetched value isn't needlessly retransmitted to a process that may +# not even need it. + +"Owner-side state backing a `DEvent`/`DFuture`. Never serialized or spilled." +mutable struct DEventBox + @atomic set::Bool + @atomic value::Union{Some{Any},Nothing} + const event::Base.Event # autoreset=false: stays signalled once notified +end +DEventBox() = DEventBox(false, nothing, Base.Event()) + +# Owner-side registry: backing-`DRef` id => box. Accessed under a +# `NonReentrantLock` via spin-locking so it is safe to touch from the `DRef` +# destructor, which may run in a GC/finalizer context where task switches (and +# hence a blocking `lock`) are illegal. +const DEVENT_REGISTRY = Dict{Int,DEventBox}() +const DEVENT_REGISTRY_LOCK = NonReentrantLock() +# Tiny sentinel stored in the datastore for each backing `DRef`; we only use the +# ref for its identity + refcounting, not its payload. +const DEVENT_SENTINEL = :__mempool_devent__ + +_devent_box(id::Int) = @safe_lock_spin DEVENT_REGISTRY_LOCK begin + get(DEVENT_REGISTRY, id, nothing) +end +function _devent_register!(id::Int, box::DEventBox) + @safe_lock_spin DEVENT_REGISTRY_LOCK begin + DEVENT_REGISTRY[id] = box + end + return +end +function _devent_delete!(id::Int) + @safe_lock_spin DEVENT_REGISTRY_LOCK begin + delete!(DEVENT_REGISTRY, id) + end + return +end + +""" + DEvent() + DEvent(pid::Integer) + +A distributed, one-shot event. `notify` sets it (idempotently) and `wait` blocks +until it is set. Safe under concurrent multithreaded access, and serializable to +other workers (all operations are then performed against the owning worker). +""" +struct DEvent + ref::DRef +end +function DEvent() + box = DEventBox() + idbox = Ref{Int}(0) + ref = poolset(DEVENT_SENTINEL; destructor = () -> _devent_delete!(idbox[])) + idbox[] = ref.id + _devent_register!(ref.id, box) + return DEvent(ref) +end +function DEvent(pid::Integer) + pid == myid() && return DEvent() + return remotecall_fetch(DEvent, pid) +end + +owner(de::DEvent) = de.ref.owner + +function _devent_notify_local(id::Int) + box = _devent_box(id) + box === nothing && return + @atomic box.set = true + notify(box.event) + return +end +function Base.notify(de::DEvent) + o = owner(de) + if o == myid() + _devent_notify_local(de.ref.id) + else + remotecall_wait(_devent_notify_local, o, de.ref.id) + end + return de +end + +function _devent_wait_local(id::Int) + box = _devent_box(id) + box === nothing && return # already cleaned up => must have fired + wait(box.event) + return +end +function Base.wait(de::DEvent) + o = owner(de) + if o == myid() + _devent_wait_local(de.ref.id) + else + remotecall_fetch(_devent_wait_local, o, de.ref.id) + end + return de +end + +function _devent_isset_local(id::Int) + box = _devent_box(id) + box === nothing && return true + return @atomic box.set +end +function isset(de::DEvent) + o = owner(de) + if o == myid() + return _devent_isset_local(de.ref.id) + else + return remotecall_fetch(_devent_isset_local, o, de.ref.id) + end +end + +""" + DFuture() + DFuture(pid::Integer) + +A write-once, `Future`-like value cell built on `DEvent`. Supports `put!`, +`fetch`, `wait`, and `isready`, and is safe under concurrent multithreaded +access (unlike `Distributed.Future`). Serializable to other workers. +""" +mutable struct DFuture + event::DEvent + @atomic cache::Union{Some{Any},Nothing} # local value cache +end +DFuture() = DFuture(DEvent(), nothing) +DFuture(pid::Integer) = DFuture(DEvent(pid), nothing) + +# `cache` is deliberately dropped on serialization: it's a local convenience +# copy, and re-sending it would waste bandwidth re-transmitting (possibly +# large) data that the destination may never even `fetch`. The destination +# just re-populates its own cache lazily, via the backing `DEvent`, on its +# first `fetch`. +function Serialization.serialize(io::AbstractSerializer, f::DFuture) + Serialization.serialize_cycle_header(io, f) && return + serialize(io, f.event) +end +function Serialization.deserialize(io::AbstractSerializer, ::Type{DFuture}) + f = ccall(:jl_new_struct_uninit, Any, (Any,), DFuture) + Serialization.deserialize_cycle(io, f) + event = deserialize(io) + ccall(:jl_set_nth_field, Cvoid, (Any, Csize_t, Any), f, 0, event) + ccall(:jl_set_nth_field, Cvoid, (Any, Csize_t, Any), f, 1, nothing) + return f +end + +function _devent_put_local(id::Int, @nospecialize(v)) + box = _devent_box(id) + box === nothing && return + _, ok = @atomicreplace box.value nothing => Some{Any}(v) + ok || return # write-once: ignore double-puts leniently + @atomic box.set = true + notify(box.event) + return +end +function Base.put!(f::DFuture, @nospecialize(v)) + de = f.event + o = owner(de) + if o == myid() + _devent_put_local(de.ref.id, v) + else + remotecall_wait(_devent_put_local, o, de.ref.id, v) + end + return f +end + +function _devent_fetch_local(id::Int) + box = _devent_box(id) + box === nothing && error("DFuture value is unavailable (already cleaned up)") + wait(box.event) + return something(@atomic box.value) +end +function Base.fetch(f::DFuture) + c = @atomic f.cache + c !== nothing && return something(c) + de = f.event + o = owner(de) + v = if o == myid() + _devent_fetch_local(de.ref.id) + else + remotecall_fetch(_devent_fetch_local, o, de.ref.id) + end + @atomicreplace f.cache nothing => Some{Any}(v) + return v +end + +Base.wait(f::DFuture) = (wait(f.event); f) +Base.isready(f::DFuture) = ((@atomic f.cache) !== nothing) || isset(f.event) diff --git a/test/runtests.jl b/test/runtests.jl index 52d30d1..4b5e059 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -275,6 +275,154 @@ end =# end +@testset "DEvent" begin + @testset "Same-process" begin + e = MemPool.DEvent() + @test MemPool.owner(e) == myid() + @test !MemPool.isset(e) + + notify(e) + @test MemPool.isset(e) + # wait returns immediately once set, and notify is idempotent + wait(e) + notify(e) + @test MemPool.isset(e) + + # Concurrent multithreaded waiters are all released by a single notify + e2 = MemPool.DEvent() + ntasks = 8 + released = zeros(Bool, ntasks) + waiters = [Threads.@spawn begin + wait(e2) + released[i] = MemPool.isset(e2) + end for i in 1:ntasks] + sleep(0.2) + @test !any(released) + notify(e2) + foreach(wait, waiters) + @test all(released) + + # Serializes/deserializes (roundtrips locally) without error + io = IOBuffer() + serialize(io, e) + e3 = deserialize(seekstart(io)) + @test e3 isa MemPool.DEvent + @test MemPool.isset(e3) + end + + @testset "Cross-process" begin + # Created on worker 2, checked/notified from worker 1 (us) + e1 = MemPool.DEvent(2) + @test MemPool.owner(e1) == 2 + @test !MemPool.isset(e1) + @test fetch(@spawnat 2 !MemPool.isset(e1)) + + notify(e1) + @test MemPool.isset(e1) + wait(e1) + @test fetch(@spawnat 2 MemPool.isset(e1)) + + # Created locally, notified from a remote worker + e2 = MemPool.DEvent() + @test MemPool.owner(e2) == myid() + @test !MemPool.isset(e2) + fetch(@spawnat 2 notify(e2)) + @test MemPool.isset(e2) + + # Created on worker 2, notified from worker 3, observed from us + e3 = MemPool.DEvent(2) + fetch(@spawnat 3 notify(e3)) + @test MemPool.isset(e3) + @test fetch(@spawnat 3 MemPool.isset(e3)) + + # wait() from a remote worker blocks until notified from another remote worker + e4 = MemPool.DEvent(2) + waiter = @spawnat 3 wait(e4) + sleep(0.2) + @test !isready(waiter) + fetch(@spawnat 2 notify(e4)) + fetch(waiter) + @test MemPool.isset(e4) + end +end + +@testset "DFuture" begin + @testset "Same-process" begin + f = MemPool.DFuture() + @test MemPool.owner(f.event) == myid() + @test !isready(f) + + put!(f, "hello") + @test isready(f) + @test fetch(f) == "hello" + + # write-once: a second put! is ignored leniently + put!(f, "world") + @test fetch(f) == "hello" + + # wait() blocks until put!, and is released by it + f2 = MemPool.DFuture() + waiter = Threads.@spawn wait(f2) + sleep(0.2) + @test !istaskdone(waiter) + put!(f2, 1) + wait(waiter) + @test istaskdone(waiter) + + # cache is populated after the first fetch + f3 = MemPool.DFuture() + put!(f3, [1, 2, 3]) + @test (@atomic f3.cache) === nothing + @test fetch(f3) == [1, 2, 3] + @test something(@atomic f3.cache) == [1, 2, 3] + end + + @testset "Cross-process" begin + # Created on worker 2, put!/fetch'd from us + f1 = MemPool.DFuture(2) + @test MemPool.owner(f1.event) == 2 + @test !isready(f1) + @test fetch(@spawnat 2 !isready(f1)) + + put!(f1, [1, 2, 3]) + @test isready(f1) + @test fetch(f1) == [1, 2, 3] + @test fetch(@spawnat 2 fetch(f1)) == [1, 2, 3] + + # Created locally, put! from a remote worker + f2 = MemPool.DFuture() + fetch(@spawnat 2 put!(f2, :remote)) + @test fetch(f2) == :remote + + # Created on worker 2, put! from worker 3, fetch'd from us and worker 3 + f3 = MemPool.DFuture(2) + fetch(@spawnat 3 put!(f3, "abc")) + @test fetch(f3) == "abc" + @test fetch(@spawnat 3 fetch(f3)) == "abc" + + # The cached value is not re-serialized: sending an already-fetched + # DFuture to another process resets its cache (verified via a + # roundtrip through the same serializer used for RPC), so the + # destination must (transparently) re-fetch from the owner rather + # than receive our locally-cached copy + f4 = MemPool.DFuture() + put!(f4, "cached-locally") + @test fetch(f4) == "cached-locally" # populates our local cache + io = IOBuffer() + serialize(io, f4) + f4_copy = deserialize(seekstart(io)) + @test (@atomic f4_copy.cache) === nothing # cache was dropped + @test something(@atomic f4.cache) == "cached-locally" # original is untouched + @test fetch(f4_copy) == "cached-locally" # transparently re-fetched from the owner + + # And this also holds when actually crossing process boundaries + f5 = MemPool.DFuture() + put!(f5, "cached-locally-too") + @test fetch(f5) == "cached-locally-too" + @test fetch(@spawnat 2 fetch(f5)) == "cached-locally-too" + end +end + @testset "StorageState" begin sstate1 = MemPool.StorageState(nothing, MemPool.StorageLeaf[],