Add DEvent and DFuture

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)

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)

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[],