diff --git a/src/datastore.jl b/src/datastore.jl index f729c41..b4169e1 100644 --- 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 index 8531389..03d8061 100644 --- 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 index 4c0e6c3..034786b 100644 --- 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