diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl index 9017f0d73..4c2abf205 100644 --- a/src/sch/Sch.jl +++ b/src/sch/Sch.jl @@ -11,7 +11,7 @@ import Base: @invokelatest import ..Dagger import ..Dagger: Context, Processor, Thunk, WeakThunk, ThunkFuture, ThunkFailedException, Chunk, WeakChunk, OSProc, AnyScope, DefaultScope, LockedObject import ..Dagger: order, dependents, noffspring, istask, inputs, unwrap_weak_checked, affinity, tochunk, timespan_start, timespan_finish, procs, move, chunktype, processor, default_enabled, get_processors, get_parent, execute!, rmprocs!, addprocs!, thunk_processor, constrain, cputhreadtime -import ..Dagger: @dagdebug +import ..Dagger: @dagdebug, @safe_lock_spin1 import DataStructures: PriorityQueue, enqueue!, dequeue_pair!, peek import ..Dagger diff --git a/src/sch/util.jl b/src/sch/util.jl index 4a1cddf41..c8fba7a61 100644 --- a/src/sch/util.jl +++ b/src/sch/util.jl @@ -1,7 +1,7 @@ "Like `errormonitor`, but tracks how many outstanding tasks are running." function errormonitor_tracked(t::Task) errormonitor(t) - lock(ERRORMONITOR_TRACKED) do tracked + @safe_lock_spin1 ERRORMONITOR_TRACKED tracked begin push!(tracked, t) end errormonitor(Threads.@spawn begin diff --git a/src/utils/locked-object.jl b/src/utils/locked-object.jl index e00eaac2d..291baf020 100644 --- a/src/utils/locked-object.jl +++ b/src/utils/locked-object.jl @@ -14,3 +14,29 @@ function Base.lock(f, x::LockedObject) unlock(x.lock) end end +Base.trylock(x::LockedObject) = trylock(x.lock) +Base.unlock(x::LockedObject) = unlock(x.lock) +payload(x::LockedObject) = x.payload + +# TODO: Move this back to MemPool +# If we actually want to acquire a lock from a finalizer, we can't cause a task +# switch. As a NonReentrantLock can only be taken by another thread that should +# be running, and not a concurrent task we'd need to switch to, we can safely +# spin. +macro safe_lock_spin1(l, o, ex) + quote + temp = $(esc(l)) + while !trylock(temp) + # we can't yield here + GC.safepoint() + end + MemPool.enable_finalizers(false) # retains compatibility with non-finalizer callers + try + $(esc(o)) = $payload(temp) + $(esc(ex)) + finally + unlock(temp) + MemPool.enable_finalizers(true) + end + end +end