Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/MemPool.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
32 changes: 22 additions & 10 deletions src/datastore.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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)

Expand Down
37 changes: 34 additions & 3 deletions src/storage.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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"))
Expand Down
Loading