fix rmprocs/addprocs locking issue. (JuliaLang#19716)

Author: amitmurthy

NEWS.md

@@ -60,6 +60,10 @@ This section lists changes that do not have deprecation warnings.
   * The array-scalar operations `div`, `mod`, `rem`, `&`, `|`, `xor`, `/`, `\`, `*`, `+`, and `-`
     now follow broadcast promotion rules ([#19692]).
 
+  * `rmprocs` now throws an exception if requested workers have not been completely
+    removed before `waitfor` seconds. With a `waitfor=0`, `rmprocs` returns immediately
+    without waiting for worker exits.
+
 Library improvements
 --------------------
 

base/multi.jl

@@ -505,52 +505,72 @@ function workers()
 end
 
 """
-    rmprocs(pids...; waitfor=0.0)
+    rmprocs(pids...; waitfor=typemax(Int))
 
-Removes the specified workers. Note that only
-process 1 can add or remove workers - if another
-worker tries to call `rmprocs`, an error will be
-thrown. The optional argument `waitfor` determines
-how long the first process will wait for the workers
-to shut down.
+Removes the specified workers. Note that only process 1 can add or remove
+workers.
+
+Argument `waitfor` specifies how long to wait for the workers to shut down:
+    - If unspecified, `rmprocs` will wait until all requested `pids` are removed.
+    - An `ErrorException` is raised if all workers cannot be terminated before
+      the requested `waitfor` seconds.
+    - With a `waitfor` value of 0, the call returns immediately with the workers
+      scheduled for removal in a different task. The scheduled `Task` object is
+      returned. The user should call `wait` on the task before invoking any other
+      parallel calls.
 """
-function rmprocs(pids...; waitfor = 0.0)
+function rmprocs(pids...; waitfor=typemax(Int))
     # Only pid 1 can add and remove processes
     if myid() != 1
-        error("only process 1 can add and remove processes")
+        throw(ErrorException("only process 1 can add and remove processes"))
     end
 
+    pids = vcat(pids...)
+    if waitfor == 0
+        t = @schedule _rmprocs(pids, typemax(Int))
+        yield()
+        return t
+    else
+        _rmprocs(pids, waitfor)
+        # return a dummy task object that user code can wait on.
+        return @schedule nothing
+    end
+end
+
+function _rmprocs(pids, waitfor)
     lock(worker_lock)
     try
         rmprocset = []
-        for i in vcat(pids...)
-            if i == 1
+        for p in vcat(pids...)
+            if p == 1
                 warn("rmprocs: process 1 not removed")
             else
-                if haskey(map_pid_wrkr, i)
-                    w = map_pid_wrkr[i]
+                if haskey(map_pid_wrkr, p)
+                    w = map_pid_wrkr[p]
                     set_worker_state(w, W_TERMINATING)
-                    kill(w.manager, i, w.config)
+                    kill(w.manager, p, w.config)
                     push!(rmprocset, w)
                 end
             end
         end
 
         start = time()
         while (time() - start) < waitfor
-            if all(w -> w.state == W_TERMINATED, rmprocset)
-                break
-            else
-                sleep(0.1)
-            end
+            all(w -> w.state == W_TERMINATED, rmprocset) && break
+            sleep(min(0.1, waitfor - (time() - start)))
         end
 
-        ((waitfor > 0) && any(w -> w.state != W_TERMINATED, rmprocset)) ? :timed_out : :ok
+        unremoved = [wrkr.id for wrkr in filter(w -> w.state != W_TERMINATED, rmprocset)]
+        if length(unremoved) > 0
+            estr = string("rmprocs: pids ", unremoved, " not terminated after ", waitfor, " seconds.")
+            throw(ErrorException(estr))
+        end
     finally
         unlock(worker_lock)
     end
 end
 
+
 """
     ProcessExitedException()
 
@@ -2238,18 +2258,18 @@ function check_same_host(pids)
 end
 
 function terminate_all_workers()
-    if myid() != 1
-        return
-    end
+    myid() != 1 && return
 
     if nprocs() > 1
-        ret = rmprocs(workers(); waitfor=0.5)
-        if ret !== :ok
+        try
+            rmprocs(workers(); waitfor=5.0)
+        catch _ex
             warn("Forcibly interrupting busy workers")
             # Might be computation bound, interrupt them and try again
             interrupt(workers())
-            ret = rmprocs(workers(); waitfor=0.5)
-            if ret !== :ok
+            try
+                rmprocs(workers(); waitfor=5.0)
+            catch _ex2
                 warn("Unable to terminate all workers")
             end
         end

test/parallel_exec.jl

@@ -953,9 +953,7 @@ if is_unix() # aka have ssh
             end
         end
 
-        @test :ok == remotecall_fetch(1, new_pids) do p
-            rmprocs(p; waitfor=5.0)
-        end
+        remotecall_fetch(plst->rmprocs(plst; waitfor=5.0), 1, new_pids)
     end
 
     print("\n\nTesting SSHManager. A minimum of 4GB of RAM is recommended.\n")
@@ -1250,3 +1248,30 @@ function test_add_procs_threaded_blas()
     rmprocs(processes_added)
 end
 test_add_procs_threaded_blas()
+
+#19687
+# ensure no race conditions between rmprocs and addprocs
+for i in 1:5
+    p = addprocs(1)[1]
+    @spawnat p sleep(5)
+    rmprocs(p; waitfor=0)
+end
+
+# Test if a wait has been called on rmprocs(...;waitfor=0), further remotecalls
+# don't throw errors.
+for i in 1:5
+    p = addprocs(1)[1]
+    np = nprocs()
+    @spawnat p sleep(5)
+    wait(rmprocs(p; waitfor=0))
+    for pid in procs()
+        @test pid == remotecall_fetch(myid, pid)
+    end
+    @test nprocs() == np - 1
+end
+
+# Test that an exception is thrown if workers are unable to be removed within requested time.
+if DoFullTest
+    pids=addprocs(4);
+    @test_throws ErrorException rmprocs(pids; waitfor=0.001);
+end

test/topology.jl

@@ -2,6 +2,8 @@
 
 include("testdefs.jl")
 addprocs(4; topology="master_slave")
+using Base.Test
+
 @test_throws RemoteException remotecall_fetch(()->remotecall_fetch(myid, 3), 2)
 
 function test_worker_counts()
@@ -19,7 +21,7 @@ end
 
 function remove_workers_and_test()
     while nworkers() > 0
-        @test :ok == rmprocs(workers()[1]; waitfor=2.0)
+        rmprocs(workers()[1]; waitfor=2.0)
         test_worker_counts()
         if nworkers() == nprocs()
             break