Ensure dynamic listeners terminate (fixes memory leak)

Author: jpsamaroo

src/sch/Sch.jl

@@ -60,8 +60,8 @@ Fields:
 - `worker_chans::Dict{Int, Tuple{RemoteChannel,RemoteChannel}}` - Communication channels between the scheduler and each worker
 - `procs_cache_list::Base.RefValue{Union{ProcessorCacheEntry,Nothing}}` - Cached linked list of processors ready to be used
 - `function_cost_cache::Dict{Type{<:Tuple},UInt}` - Cache of estimated CPU time required to compute the given signature
-- `halt::Base.RefValue{Bool}` - Flag indicating, when set, that the scheduler should halt immediately
-- `lock::ReentrantLock()` - Lock around operations which modify the state
+- `halt::Base.Event` - Event indicating that the scheduler is halting
+- `lock::ReentrantLock` - Lock around operations which modify the state
 - `futures::Dict{Thunk, Vector{ThunkFuture}}` - Futures registered for waiting on the result of a thunk.
 - `errored::Set{Thunk}` - Thunks that threw an error
 - `chan::RemoteChannel{Channel{Any}}` - Channel for receiving completed thunks
@@ -84,7 +84,7 @@ struct ComputeState
     worker_chans::Dict{Int, Tuple{RemoteChannel,RemoteChannel}}
     procs_cache_list::Base.RefValue{Union{ProcessorCacheEntry,Nothing}}
     function_cost_cache::Dict{Type{<:Tuple},UInt}
-    halt::Base.RefValue{Bool}
+    halt::Base.Event
     lock::ReentrantLock
     futures::Dict{Thunk, Vector{ThunkFuture}}
     errored::Set{Thunk}
@@ -109,7 +109,7 @@ function start_state(deps::Dict, node_order, chan)
                          Dict{Int, Tuple{RemoteChannel,RemoteChannel}}(),
                          Ref{Union{ProcessorCacheEntry,Nothing}}(nothing),
                          Dict{Type{<:Tuple},UInt}(),
-                         Ref{Bool}(false),
+                         Base.Event(),
                          ReentrantLock(),
                          Dict{Thunk, Vector{ThunkFuture}}(),
                          Set{Thunk}(),
@@ -441,7 +441,9 @@ function compute_dag(ctx, d::Thunk; options=SchedulerOptions())
 
         safepoint(state)
     end
-    state.halt[] = true
+    @assert !isready(state.chan)
+    close(state.chan)
+    notify(state.halt)
     @sync for p in procs_to_use(ctx)
         @async cleanup_proc(state, p)
     end

src/sch/dynamic.jl

@@ -25,7 +25,7 @@ struct DynamicThunkException <: Exception
 end
 
 function safepoint(state)
-    if state.halt[]
+    if state.halt.set
         # Force dynamic thunks and listeners to terminate
         for (inp_chan,out_chan) in values(state.worker_chans)
             close(inp_chan)
@@ -39,10 +39,11 @@ end
 "Processes dynamic messages from worker-executing thunks."
 function dynamic_listener!(ctx, state)
     task = current_task() # The scheduler's main task
+    listener_tasks = Task[]
     for tid in keys(state.worker_chans)
         inp_chan, out_chan = state.worker_chans[tid]
-        @async begin
-            while isopen(inp_chan)
+        push!(listener_tasks, @async begin
+            while isopen(inp_chan) && !state.halt.set
                 tid, f, data = try
                     take!(inp_chan)
                 catch err
@@ -70,6 +71,14 @@ function dynamic_listener!(ctx, state)
                     end
                 end
             end
+        end)
+    end
+    @async begin
+        wait(state.halt)
+        for ltask in listener_tasks
+            # TODO: Not sure why we need the @async here, but otherwise we
+            # don't stop all the listener tasks
+            @async Base.throwto(ltask, SchedulerHaltedException())
         end
     end
 end
@@ -91,7 +100,7 @@ end
 "Commands the scheduler to halt execution immediately."
 halt!(h::SchedulerHandle) = exec!(_halt, h, nothing)
 function _halt(ctx, state, task, tid, _)
-    state.halt[] = true
+    notify(state.halt)
     put!(state.chan, (1, nothing, SchedulerHaltedException(), nothing))
     Base.throwto(task, SchedulerHaltedException())
 end