fixup! Add support for worker state callbacks

Author: JamesWrigley

src/cluster.jl

@@ -1009,13 +1009,18 @@ remove_worker_exiting_callback(key) = _remove_callback(key, worker_exiting_callb
 
 Register a callback to be called on the master process when a worker has exited
 for any reason (i.e. not only because of [`rmprocs()`](@ref) but also the worker
-segfaulting etc). The callback will be called with the worker ID,
-e.g. `f(w::Int)`. Chooses and returns a unique key for the callback if `key` is
+segfaulting etc). Chooses and returns a unique key for the callback if `key` is
 not specified.
 
+The callback will be called with the worker ID and the final
+`Distributed.WorkerState` of the worker, e.g. `f(w::Int, state)`. `state` is an
+enum, a value of `WorkerState_terminated` means a graceful exit and a value of
+`WorkerState_exterminated` means the worker died unexpectedly.
+
 If the callback throws an exception it will be caught and printed.
 """
-add_worker_exited_callback(f::Base.Callable; key=nothing) = _add_callback(f, key, worker_exited_callbacks)
+add_worker_exited_callback(f::Base.Callable; key=nothing) = _add_callback(f, key, worker_exited_callbacks;
+                                                                          arg_types=Tuple{Int, WorkerState})
 
 """
     remove_worker_exited_callback(key)
@@ -1426,7 +1431,7 @@ function deregister_worker(pg, pid)
     if myid() == 1
         for (name, callback) in worker_exited_callbacks
             try
-                callback(pid)
+                callback(pid, w.state)
             catch ex
                 @error "Error when running worker-exited callback '$(name)'" exception=(ex, catch_backtrace())
             end

test/distributed_exec.jl

@@ -1946,11 +1946,11 @@ end
     starting_managers = []
     started_workers = Int[]
     exiting_workers = Int[]
-    exited_workers = Int[]
+    exited_workers = []
     starting_key = DistributedNext.add_worker_starting_callback((manager, kwargs) -> push!(starting_managers, manager))
     started_key = DistributedNext.add_worker_started_callback(pid -> (push!(started_workers, pid); error("foo")))
     exiting_key = DistributedNext.add_worker_exiting_callback(pid -> push!(exiting_workers, pid))
-    exited_key = DistributedNext.add_worker_exited_callback(pid -> push!(exited_workers, pid))
+    exited_key = DistributedNext.add_worker_exited_callback((pid, state) -> push!(exited_workers, (pid, state)))
 
     # Test that the worker-started exception bubbles up
     @test_throws TaskFailedException addprocs(1)
@@ -1960,7 +1960,7 @@ end
     @test started_workers == [pid]
     rmprocs(workers())
     @test exiting_workers == [pid]
-    @test exited_workers == [pid]
+    @test exited_workers == [(pid, DistributedNext.WorkerState_terminated)]
 
     # Trying to reset an existing callback should fail
     @test_throws ArgumentError DistributedNext.add_worker_started_callback(Returns(nothing); key=started_key)
@@ -1992,6 +1992,17 @@ end
     @test length(started_workers) == 1
     @test length(exiting_workers) == 1
     @test length(exited_workers) == 1
+
+    # Test that workers that were killed forcefully are detected as such
+    exit_state = nothing
+    DistributedNext.add_worker_exited_callback((pid, state) -> exit_state = state)
+    pid = only(addprocs(1))
+
+    redirect_stderr(devnull) do
+        remote_do(exit, pid)
+        timedwait(() -> !isnothing(exit_state), 10)
+    end
+    @test exit_state == DistributedNext.WorkerState_exterminated
 end
 
 # This is a simplified copy of a test from Revise.jl's tests