poolset: Add optional destructor function

JuliaData · Jan 3, 2024 · 809c6de · 809c6de
1 parent 7a6447d
commit 809c6de
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/src/datastore.jl b/src/datastore.jl
@@ -351,7 +351,8 @@ isondisk(x::DRef) = isondisk(x.id)
 function poolset(@nospecialize(x), pid=myid(); size=approx_size(x),
                  retain=false, restore=false,
                  device=GLOBAL_DEVICE[], leaf_device=initial_leaf_device(device),
-                 tag=nothing, leaf_tag=Tag())
+                 tag=nothing, leaf_tag=Tag(),
+                 destructor=nothing)
     if pid == myid()
         id = atomic_add!(id_counter, 1)
         sstate = if !restore
@@ -368,7 +369,8 @@ function poolset(@nospecialize(x), pid=myid(); size=approx_size(x),
         state = RefState(sstate,
                          size,
                          tag,
-                         leaf_tag)
+                         leaf_tag,
+                         destructor)
         with_lock(datastore_counters_lock) do
             datastore_counters[(pid, id)] = RefCounters()
         end
@@ -451,6 +453,10 @@ function datastore_delete(id)
     @safe_lock_spin datastore_lock begin
         haskey(datastore, id) && delete!(datastore, id)
     end
+    dtor = state.destructor
+    if dtor !== nothing
+        errormonitor(Threads.@spawn dtor())
+    end
     return
 end
 

diff --git a/src/storage.jl b/src/storage.jl
@@ -309,9 +309,11 @@ mutable struct RefState
     # Metadata to associate with the reference
     tag::Any
     leaf_tag::Tag
+    # Destructor, if any
+    destructor::Any
 end
 RefState(storage::StorageState, size::Integer) =
-    RefState(storage, size, nothing, Tag())
+    RefState(storage, size, nothing, Tag(), nothing)
 function Base.getproperty(state::RefState, field::Symbol)
     if field === :storage
         throw(ArgumentError("Cannot directly read `:storage` field of `RefState`\nUse `storage_read(state)` instead"))

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -246,6 +246,16 @@ end
     end
 end
 
+@testset "Destructors" begin
+    ref_del = Ref(false)
+    # @eval because testsets retain values in weird ways
+    x = @eval Ref{Any}(poolset(123; destructor=()->(@assert !$ref_del[]; $ref_del[]=true;)))
+    @test !ref_del[]
+    x[] = nothing
+    GC.gc(); yield()
+    @test ref_del[]
+end
+
 @testset "StorageState" begin
     sstate1 = MemPool.StorageState(nothing,
                                    MemPool.StorageLeaf[],
@@ -305,7 +315,7 @@ end
                                    CPURAMDevice(),
                                    Base.Event())
     notify(sstate1)
-    state = MemPool.RefState(sstate1, 64, "abc", MemPool.Tag(SerializationFileDevice=>123))
+    state = MemPool.RefState(sstate1, 64, "abc", MemPool.Tag(SerializationFileDevice=>123), nothing)
     @test state.size == 64
     @test MemPool.storage_size(state) == 64
     @test_throws ArgumentError state.storage