Merge pull request #329 from JuliaParallel/jps/init-dynamic

Sch: Start dynamic listener in init_proc

Author: jpsamaroo

Diff for: src/sch/Sch.jl

@@ -286,6 +286,10 @@ function init_proc(state, p, log_sink)
     lock(state.lock) do
         state.worker_chans[p.pid] = (inp_chan, out_chan)
     end
+
+    # Setup dynamic listener
+    dynamic_listener!(ctx, state, p.pid)
+
     timespan_finish(ctx, :init_proc, p.pid, 0)
 end
 function _cleanup_proc(uid, log_sink)
@@ -405,9 +409,6 @@ function scheduler_init(ctx, state::ComputeState, d::Thunk, options, deps)
             @error "Error assigning workers" exception=(err,catch_backtrace())
         end
     end
-
-    # setup dynamic listeners
-    dynamic_listener!(ctx, state)
 end
 
 function scheduler_run(ctx, state::ComputeState, d::Thunk, options)

Diff for: src/sch/dynamic.jl

@@ -41,61 +41,56 @@ function safepoint(state)
 end
 
 "Processes dynamic messages from worker-executing thunks."
-function dynamic_listener!(ctx, state)
+function dynamic_listener!(ctx, state, wid)
     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]
-        push!(listener_tasks, @async begin
-            while isopen(inp_chan) && !state.halt.set
-                tid, f, data = try
-                    take!(inp_chan)
-                catch err
-                    if !(unwrap_nested_exception(err) isa Union{SchedulerHaltedException,
-                                                                ProcessExitedException,
-                                                                InvalidStateException})
-                        iob = IOContext(IOBuffer(), :color=>true)
-                        println(iob, "Error in sending dynamic request:")
-                        Base.showerror(iob, err)
-                        Base.show_backtrace(iob, catch_backtrace())
-                        println(iob)
-                        seek(iob.io, 0)
-                        write(stderr, iob)
-                    end
-                    break
+    inp_chan, out_chan = state.worker_chans[wid]
+    listener_task = @async begin
+        while isopen(inp_chan) && !state.halt.set
+            tid, f, data = try
+                take!(inp_chan)
+            catch err
+                if !(unwrap_nested_exception(err) isa Union{SchedulerHaltedException,
+                                                            ProcessExitedException,
+                                                            InvalidStateException})
+                    iob = IOContext(IOBuffer(), :color=>true)
+                    println(iob, "Error in sending dynamic request:")
+                    Base.showerror(iob, err)
+                    Base.show_backtrace(iob, catch_backtrace())
+                    println(iob)
+                    seek(iob.io, 0)
+                    write(stderr, iob)
                 end
-                res = try
-                    (false, lock(state.lock) do
-                        Base.invokelatest(f, ctx, state, task, tid, data)
-                    end)
-                catch err
-                    (true, RemoteException(CapturedException(err,catch_backtrace())))
-                end
-                try
-                    put!(out_chan, res)
-                catch err
-                    if !(unwrap_nested_exception(err) isa Union{SchedulerHaltedException,
-                                                                ProcessExitedException,
-                                                                InvalidStateException})
-                        iob = IOContext(IOBuffer(), :color=>true)
-                        println(iob, "Error in sending dynamic result from $f:")
-                        Base.showerror(iob, err)
-                        Base.show_backtrace(iob, catch_backtrace())
-                        println(iob)
-                        seek(iob.io, 0)
-                        write(stderr, iob)
-                    end
+                break
+            end
+            res = try
+                (false, lock(state.lock) do
+                    Base.invokelatest(f, ctx, state, task, tid, data)
+                end)
+            catch err
+                (true, RemoteException(CapturedException(err,catch_backtrace())))
+            end
+            try
+                put!(out_chan, res)
+            catch err
+                if !(unwrap_nested_exception(err) isa Union{SchedulerHaltedException,
+                                                            ProcessExitedException,
+                                                            InvalidStateException})
+                    iob = IOContext(IOBuffer(), :color=>true)
+                    println(iob, "Error in sending dynamic result from $f:")
+                    Base.showerror(iob, err)
+                    Base.show_backtrace(iob, catch_backtrace())
+                    println(iob)
+                    seek(iob.io, 0)
+                    write(stderr, iob)
                 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
+        # TODO: Not sure why we need the @async here, but otherwise we
+        # don't stop all the listener tasks
+        @async Base.throwto(listener_task, SchedulerHaltedException())
     end
 end