Merge pull request #98 from JuliaData/jps/faster-poolset

Optimize poolset and pool deletion

Author: jpsamaroo

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)
 

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)
 

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"))