errormonitor_tracked: Add task name

Author: jpsamaroo

src/precompile.jl

@@ -6,6 +6,8 @@
     t1 = @spawn 1+1
     t2 = spawn(+, 1, t1)
     fetch(t2)
+
+    # Shutdown scheduler and clean up
     spawn() do
         Sch.halt!(sch_handle())
     end
@@ -16,12 +18,15 @@
     GC.gc()
     yield()
     lock(Sch.ERRORMONITOR_TRACKED) do tracked
-        if all(t->istaskdone(t) || istaskfailed(t), tracked)
+        if all(t->istaskdone(t) || istaskfailed(t), map(last, tracked))
             empty!(tracked)
             return
         end
-        for t in tracked
-            Base.throwto(t, InterruptException())
+        for (name, t) in tracked
+            @warn "Waiting on $name"
+            if t.state == :runnable
+                Base.throwto(t, InterruptException())
+            end
         end
     end
     MemPool.exit_hook()

src/sch/Sch.jl

@@ -366,7 +366,7 @@ function init_proc(state, p, log_sink)
                         end
                     end
                 end
-                errormonitor_tracked(t)
+                errormonitor_tracked("worker monitor $wid", t)
                 WORKER_MONITOR_TASKS[wid] = t
                 WORKER_MONITOR_CHANS[wid] = Dict{UInt64,RemoteChannel}()
             end
@@ -1290,7 +1290,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
                 else
                     t.sticky = false
                 end
-                tasks[thunk_id] = errormonitor_tracked(schedule(t))
+                tasks[thunk_id] = errormonitor_tracked("thunk $thunk_id", schedule(t))
                 proc_occupancy[] += task_occupancy
                 time_pressure[] += time_util
             end
@@ -1302,7 +1302,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
     else
         proc_run_task.sticky = false
     end
-    return errormonitor_tracked(schedule(proc_run_task))
+    return errormonitor_tracked("processor $to_proc", schedule(proc_run_task))
 end
 
 """

src/sch/dynamic.jl

@@ -87,8 +87,8 @@ function dynamic_listener!(ctx, state, wid)
             end
         end
     end
-    errormonitor_tracked(listener_task)
-    errormonitor_tracked(@async begin
+    errormonitor_tracked("dynamic_listener! $wid", listener_task)
+    errormonitor_tracked("dynamic_listener! (halt+throw) $wid", @async begin
         wait(state.halt)
         # TODO: Not sure why we need the @async here, but otherwise we
         # don't stop all the listener tasks

src/sch/eager.jl

@@ -20,7 +20,7 @@ function init_eager()
         return
     end
     ctx = eager_context()
-    errormonitor_tracked(Threads.@spawn try
+    errormonitor_tracked("eager compute()", Threads.@spawn try
         sopts = SchedulerOptions(;allow_errors=true)
         opts = Dagger.Options((;scope=Dagger.ExactScope(Dagger.ThreadProc(1, 1)),
                                 occupancy=Dict(Dagger.ThreadProc=>0)))
@@ -101,7 +101,7 @@ function thunk_yield(f)
 end
 
 eager_cleanup(t::Dagger.EagerThunkFinalizer) =
-    errormonitor_tracked(Threads.@spawn eager_cleanup(EAGER_STATE[], t.uid))
+    errormonitor_tracked("eager_cleanup $(t.uid)", Threads.@spawn eager_cleanup(EAGER_STATE[], t.uid))
 function eager_cleanup(state, uid)
     tid = nothing
     lock(EAGER_ID_MAP) do id_map

src/sch/util.jl

@@ -1,15 +1,15 @@
 "Like `errormonitor`, but tracks how many outstanding tasks are running."
-function errormonitor_tracked(t::Task)
+function errormonitor_tracked(name::String, t::Task)
     errormonitor(t)
     @safe_lock_spin1 ERRORMONITOR_TRACKED tracked begin
-        push!(tracked, t)
+        push!(tracked, name => t)
     end
     errormonitor(Threads.@spawn begin
         try
             wait(t)
         finally
             lock(ERRORMONITOR_TRACKED) do tracked
-                idx = findfirst(o->o===t, tracked)
+                idx = findfirst(o->o[2]===t, tracked)
                 # N.B. This may be nothing if precompile emptied these
                 if idx !== nothing
                     deleteat!(tracked, idx)
@@ -18,7 +18,7 @@ function errormonitor_tracked(t::Task)
         end
     end)
 end
-const ERRORMONITOR_TRACKED = LockedObject(Task[])
+const ERRORMONITOR_TRACKED = LockedObject(Pair{String,Task}[])
 
 """
     unwrap_nested_exception(err::Exception) -> Bool