diff --git a/src/MemPool.jl b/src/MemPool.jl index 12dc75b..ab3e309 100644 --- a/src/MemPool.jl +++ b/src/MemPool.jl @@ -137,6 +137,10 @@ end function __init__() SESSION[] = "sess-" * randstring(6) + # Ensure the shared born-ready Event is set at runtime, independent of how + # precompilation serialized it. + notify(ALWAYS_READY) + DISKCACHE_CONFIG[] = diskcache_config = DiskCacheConfig() setup_global_device!(diskcache_config) diff --git a/src/datastore.jl b/src/datastore.jl index a58b999..d85f0f6 100644 --- a/src/datastore.jl +++ b/src/datastore.jl @@ -463,9 +463,13 @@ function poolset(@nospecialize(x), pid=myid(); size=approx_size(x), id = atomic_add!(id_counter, 1) sstate = if !restore + # Data is already in memory, so this state is born ready: share the + # set `ALWAYS_READY` Event and the empty `NO_LEAVES` vector instead + # of allocating fresh ones. StorageState(Some{Any}(x), - Vector{StorageLeaf}(), - device) + NO_LEAVES, + device, + ALWAYS_READY) else @assert !isa(leaf_device, CPURAMDevice) "Cannot use `CPURAMDevice()` as leaf device when `restore=true`" StorageState(nothing, @@ -562,7 +566,7 @@ end function _getlocal(f, id, remote, args...; local_only::Bool, from::Int) state = with_lock(()->datastore[id], datastore_lock) - lock_read(state.lock) do + lock_read(getlock!(state)) do if state.redirect !== nothing return RedirectTo(state.redirect) end @@ -574,6 +578,13 @@ function _getlocal(f, id, remote, args...; local_only::Bool, from::Int) end end +"Deferred device-deletion work item drained by the shared `SEND_QUEUE` task." +function _delete_device_work(state::RefState, id::Int) + device = storage_read(state).root + device !== nothing && delete_from_device!(device, state, id) + return +end + function datastore_delete(id) @safe_lock_spin datastore_counters_lock begin DEBUG_REFCOUNTING[] && _enqueue_work(Core.print, "-- (", myid(), ", ", id, ") with ", string(datastore_counters[(myid(), id)]), "\n"; gc_context=true) @@ -585,12 +596,13 @@ function datastore_delete(id) haskey(datastore, id) ? datastore[id] : nothing end (state === nothing) && return - errormonitor(Threads.@spawn begin - device = storage_read(state).root - if device !== nothing - delete_from_device!(device, state, id) - end - end) + # Defer device deletion onto the shared serial work queue rather than + # spawning a fresh task per teardown. datastore_delete frequently runs from + # a GC/finalizer context where task switches are illegal, so the device read + # and `delete_from_device!` (which may wait on an in-flight storage + # transition) must not run inline here. `_enqueue_work` is finalizer-safe and + # reuses one long-lived task, avoiding a Task allocation per ref teardown. + _enqueue_work(_delete_device_work, state, id; gc_context=true) @safe_lock_spin datastore_lock begin haskey(datastore, id) && delete!(datastore, id) end @@ -621,7 +633,7 @@ function migrate!(ref::DRef, to::Integer; pre_migration=nothing, dest_post_migra # Lock the ref against further accesses # FIXME: Below is racey w.r.t data mutation local new_ref - @lock state.lock begin + @lock getlock!(state) begin # Read the current value of the ref data = read_from_device(state, ref, true) diff --git a/src/storage.jl b/src/storage.jl index 18d3741..77c019f 100644 --- a/src/storage.jl +++ b/src/storage.jl @@ -304,6 +304,17 @@ end Base.notify(sstate::StorageState) = notify(sstate.ready) Base.wait(sstate::StorageState) = wait(sstate.ready) +# Shared sentinels to avoid per-ref allocation on the common `poolset` path. +# `ALWAYS_READY`: a permanently-set Event — refs whose data is in memory at +# creation never need readers to block on `.ready`, so they can share one set +# Event instead of allocating a fresh one. It is (re)notified in `__init__` so +# it is set regardless of how precompilation serializes it. +# `NO_LEAVES`: an empty leaves vector — `leaves` are only ever replaced via +# copy-on-write (`vcat`/`filter`/`copy_leaves`), never mutated in place, so a +# single shared empty vector is safe for any ref that has not been spilled. +const ALWAYS_READY = Base.Event() +const NO_LEAVES = StorageLeaf[] + mutable struct RefState # The storage state associated with the reference and its values @atomic storage::StorageState @@ -314,8 +325,11 @@ mutable struct RefState leaf_tag::Tag # Destructor, if any destructor::Any - # A Reader-Writer lock to protect access to this struct - lock::ReadWriteLock + # A Reader-Writer lock to protect access to this struct, created lazily on + # first use (see `getlock!`): most refs are never read-locked or migrated, + # so allocating the lock (a ReentrantLock + two Conditions) eagerly is pure + # overhead on the ref-creation hot path. + @atomic lock::Union{ReadWriteLock,Nothing} # The DRef that this value may be redirecting to redirect::Union{DRef,Nothing} end @@ -325,7 +339,24 @@ RefState(storage::StorageState, size::Integer; RefState(storage, size, tag, leaf_tag, destructor, - ReadWriteLock(), nothing) + nothing, nothing) + +""" + getlock!(state::RefState) -> ReadWriteLock + +Returns the reader-writer lock guarding `state`, creating and atomically +installing it on first use. The lock is allocated lazily because most refs +(small in-memory values, task handles, etc.) are never read-locked via +`_getlocal`/`poolget` nor migrated, so eager allocation just wastes time and +memory during ref creation. +""" +function getlock!(state::RefState) + lk = @atomic :acquire state.lock + lk !== nothing && return lk + newlk = ReadWriteLock() + repl = @atomicreplace :acquire_release :acquire state.lock nothing => newlk + return repl.success ? newlk : (repl.old::ReadWriteLock) +end function Base.getproperty(state::RefState, field::Symbol) if field === :storage throw(ArgumentError("Cannot directly read `:storage` field of `RefState`\nUse `storage_read(state)` instead"))