Make sure workers don't outlive the manager

Project.toml

@@ -1,6 +1,6 @@
 name = "ParallelTestRunner"
 uuid = "d3525ed8-44d0-4b2c-a655-542cee43accc"
-version = "2.9.3"
+version = "2.10.0"
 authors = ["Valentin Churavy <v.churavy@gmail.com>"]
 
 [deps]

docs/src/index.md

@@ -147,6 +147,14 @@ Press `Ctrl+C` to interrupt the test run. The framework will:
 - Display a summary of completed tests
 - Exit gracefully
 
+### Workers Don't Outlive the Test Run
+
+Workers are stopped when the main process exits, also when it's terminated with `SIGTERM`, as
+CI services do when a job is cancelled. On Linux, FreeBSD and Windows they're killed
+by the operating system even when the main process is killed abruptly, e.g. with `SIGKILL`,
+including while they're busy running a test; on macOS that isn't possible, so a worker in that
+situation is left running until its test completes.
+
 ## Test File Structure
 
 Your test files should be standard Julia test files using the `Test` standard library:

src/ParallelTestRunner.jl

@@ -23,6 +23,9 @@ include("compatutils.jl")
 # PTRWorker, worker_id, test_exe, addworkers, addworker
 include("ptrworker.jl")
 
+# arm_manager_death!, bind_worker!, unbind_worker!
+include("lifeline.jl")
+
 # find_tests, ParsedArgs, extract_flag!, parse_args, filter_tests!, partition_tests
 include("args.jl")
 

src/lifeline.jl

@@ -0,0 +1,185 @@
+# Workers must not outlive the manager process. Malt spawns them detached: on Unix they get
+# their own session, so the signals CI runners send to the process group of the job when
+# it's cancelled don't reach them, and on Windows they break away from libuv's job which
+# kills its processes when the manager exits. A worker also notices the manager is gone
+# only when its event loop runs, which may never happen if it's busy running a test, and
+# if the manager is SIGKILLed it has no chance to stop its workers. So we ask the OS to
+# kill the workers when the manager dies, where that's possible:
+#
+# * on Linux and FreeBSD each worker asks for a SIGKILL when its parent process exits;
+# * on Windows the manager puts its workers in a job object which kills its processes when
+#   its last handle, held only by the manager, is closed, i.e. when the manager exits.
+#
+# On macOS there is no such facility, short of keeping a task always running in the worker
+# to watch the manager, so workers there are only covered when the manager exits normally
+# or on SIGTERM, by the `atexit` hook below.
+
+"""
+    arm_manager_death!(manager_pid)
+
+Run in a worker, make the OS kill it when the manager process, with PID `manager_pid`,
+exits. Returns whether this was successful.
+"""
+function arm_manager_death!(manager_pid::Integer)
+    @static if Sys.islinux() || Sys.isfreebsd()
+        # when `exename` is a wrapper, e.g. `compute-sanitizer julia`, our parent is the
+        # wrapper, not the manager. If the manager is gone already we are idle, waiting for
+        # the next request, so we'll exit anyway as soon as we see the socket is closed.
+        (@ccall getppid()::Cint) == manager_pid || return false
+        ok = @static if Sys.islinux()
+            # The signal is sent when the *thread* which spawned us exits, not the process,
+            # but the manager spawns workers from Julia's threads, which live until exit.
+            PR_SET_PDEATHSIG = Cint(1)
+            (@ccall prctl(PR_SET_PDEATHSIG::Cint; Base.SIGKILL::Culong)::Cint) == 0
+        else
+            P_PID = Cint(0)
+            PROC_PDEATHSIG_CTL = Cint(11)
+            signum = Ref(Cint(Base.SIGKILL))
+            (@ccall procctl(P_PID::Cint, 0::Int64, PROC_PDEATHSIG_CTL::Cint,
+                            signum::Ptr{Cint})::Cint) == 0
+        end
+        # the manager may have died before we asked for the signal
+        (@ccall getppid()::Cint) == manager_pid || @ccall _exit(1::Cint)::Cvoid
+        return ok
+    else
+        return false
+    end
+end
+
+@static if Sys.iswindows()
+
+# `JOBOBJECT_EXTENDED_LIMIT_INFORMATION`, 144 bytes on 64-bit Windows, 112 on 32-bit
+struct JobObjectExtendedLimitInformation
+    # JOBOBJECT_BASIC_LIMIT_INFORMATION
+    per_process_user_time_limit::Int64
+    per_job_user_time_limit::Int64
+    limit_flags::UInt32
+    minimum_working_set_size::UInt
+    maximum_working_set_size::UInt
+    active_process_limit::UInt32
+    affinity::UInt
+    priority_class::UInt32
+    scheduling_class::UInt32
+    # IO_COUNTERS
+    io_counters::NTuple{6,UInt64}
+    process_memory_limit::UInt
+    job_memory_limit::UInt
+    peak_process_memory_used::UInt
+    peak_job_memory_used::UInt
+end
+
+const JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = UInt32(0x2000)
+const JOB_OBJECT_EXTENDED_LIMIT_INFORMATION_CLASS = Cint(9)
+const PROCESS_TERMINATE = UInt32(0x0001)
+const PROCESS_SET_QUOTA = UInt32(0x0100)
+
+# The job object of the workers. Its handle is never closed, the OS does it when the
+# manager exits, killing the workers.
+const WORKERS_JOB = Lockable(Ref(C_NULL))
+
+function create_workers_job()
+    sizeof(JobObjectExtendedLimitInformation) == (Sys.WORD_SIZE == 64 ? 144 : 112) ||
+        return C_NULL
+    job = ccall((:CreateJobObjectW, "kernel32"), stdcall, Ptr{Cvoid},
+                (Ptr{Cvoid}, Ptr{Cvoid}), C_NULL, C_NULL)
+    job == C_NULL && return C_NULL
+    info = Ref(JobObjectExtendedLimitInformation(0, 0, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
+                                                 0, 0, 0, 0, 0, 0, ntuple(_ -> 0, 6),
+                                                 0, 0, 0, 0))
+    ok = ccall((:SetInformationJobObject, "kernel32"), stdcall, Cint,
+               (Ptr{Cvoid}, Cint, Ptr{JobObjectExtendedLimitInformation}, UInt32),
+               job, JOB_OBJECT_EXTENDED_LIMIT_INFORMATION_CLASS, info,
+               sizeof(JobObjectExtendedLimitInformation))
+    if ok == 0
+        ccall((:CloseHandle, "kernel32"), stdcall, Cint, (Ptr{Cvoid},), job)
+        return C_NULL
+    end
+    return job
+end
+
+function assign_to_workers_job(pid::Integer)
+    job = @lock WORKERS_JOB begin
+        if WORKERS_JOB[][] == C_NULL
+            WORKERS_JOB[][] = create_workers_job()
+        end
+        WORKERS_JOB[][]
+    end
+    job == C_NULL && return false
+    process = ccall((:OpenProcess, "kernel32"), stdcall, Ptr{Cvoid}, (UInt32, Cint, UInt32),
+                    PROCESS_TERMINATE | PROCESS_SET_QUOTA, false, pid)
+    process == C_NULL && return false
+    ok = ccall((:AssignProcessToJobObject, "kernel32"), stdcall, Cint,
+               (Ptr{Cvoid}, Ptr{Cvoid}), job, process)
+    ccall((:CloseHandle, "kernel32"), stdcall, Cint, (Ptr{Cvoid},), process)
+    return ok != 0
+end
+
+else
+
+assign_to_workers_job(pid::Integer) = false
+
+end
+
+# The workers not stopped yet, by `worker_id`
+const LIVE_WORKERS = Lockable(Dict{Int,PTRWorker}())
+
+# Stop the workers still running when the manager exits, e.g. on SIGTERM. Malt does this
+# too, but one worker at a time, waiting up to 30 seconds for each one before killing it,
+# while CI runners usually give a few seconds before killing the manager altogether.
+function stop_live_workers()
+    workers = @lock LIVE_WORKERS begin
+        workers = collect(values(LIVE_WORKERS[]))
+        empty!(LIVE_WORKERS[])
+        workers
+    end
+    filter!(Malt.isrunning, workers)
+    isempty(workers) && return
+    @static if VERSION < v"1.11" && !Sys.iswindows()
+        # On SIGTERM Julia v1.10 runs the `atexit` hooks while holding the I/O lock, so
+        # anything waiting on the event loop hangs: kill the workers without waiting.
+        for wrkr in workers
+            @ccall kill(wrkr.w.proc_pid::Cint, Base.SIGKILL::Cint)::Cint
+        end
+        return
+    end
+    # ask nicely first, so that they can e.g. write their coverage files
+    for wrkr in workers
+        try
+            Malt.remote_do(Base.exit, wrkr.w)
+        catch
+        end
+    end
+    timedwait(() -> !any(Malt.isrunning, workers), 2; pollint=0.05)
+    for wrkr in workers
+        Malt.isrunning(wrkr) && kill(wrkr.w, Base.SIGKILL)
+    end
+    # wait for them to be reaped, so that Malt doesn't try to stop them again
+    timedwait(() -> !any(Malt.isrunning, workers), 5; pollint=0.05)
+    return
+end
+
+"""
+    bind_worker!(wrkr::PTRWorker)
+
+Run in the manager, make sure `wrkr` doesn't outlive it. See also [`arm_manager_death!`](@ref).
+"""
+function bind_worker!(wrkr::PTRWorker)
+    try
+        @static if Sys.iswindows()
+            assign_to_workers_job(wrkr.w.proc_pid)
+        elseif Sys.islinux() || Sys.isfreebsd()
+            Malt.remote_eval_wait(Main, wrkr.w,
+                                  :(ParallelTestRunner.arm_manager_death!($(getpid()))))
+        end
+    catch err
+        # this is best-effort, and must never break the test run
+        @debug "Could not bind the worker to the manager" exception=(err, catch_backtrace())
+    end
+    @lock LIVE_WORKERS LIVE_WORKERS[][worker_id(wrkr)] = wrkr
+    # `atexit` hooks run in reverse order, and Malt registers one for each new worker, which
+    # would stop it slowly, so register ours again to run first.
+    atexit(stop_live_workers)
+    return
+end
+
+unbind_worker!(wrkr::PTRWorker) = @lock LIVE_WORKERS delete!(LIVE_WORKERS[], worker_id(wrkr))

src/ptrworker.jl

@@ -38,7 +38,10 @@ end
 
 worker_id(wrkr::PTRWorker) = wrkr.id
 Malt.isrunning(wrkr::PTRWorker) = Malt.isrunning(wrkr.w)
-Malt.stop(wrkr::PTRWorker) = Malt.stop(wrkr.w)
+function Malt.stop(wrkr::PTRWorker)
+    unbind_worker!(wrkr)
+    return Malt.stop(wrkr.w)
+end
 
 
 function test_exe(color::Bool=false)
@@ -103,6 +106,7 @@ function addworker(;
     # need it to reference `AbstractTestRecord`, `execute`, etc. when defining
     # custom record types.
     Malt.remote_eval_wait(Main, wrkr.w, :(import ParallelTestRunner))
+    bind_worker!(wrkr)
     if init_worker_code != :()
         Malt.remote_eval_wait(Main, wrkr.w, init_worker_code)
     end

test/workers.jl

@@ -272,6 +272,76 @@ end
     @test ParallelTestRunner.ID_COUNTER[] == old_id_counter + 3
 end
 
+# whether the process with PID `pid` is still running
+function pid_running(pid)
+    @static if Sys.iswindows()
+        SYNCHRONIZE = 0x00100000
+        WAIT_TIMEOUT = 0x00000102
+        handle = ccall((:OpenProcess, "kernel32"), stdcall, Ptr{Cvoid},
+                       (UInt32, Cint, UInt32), SYNCHRONIZE, false, pid)
+        handle == C_NULL && return false
+        running = ccall((:WaitForSingleObject, "kernel32"), stdcall, UInt32,
+                        (Ptr{Cvoid}, UInt32), handle, 0) == WAIT_TIMEOUT
+        ccall((:CloseHandle, "kernel32"), stdcall, Cint, (Ptr{Cvoid},), handle)
+        return running
+    else
+        (@ccall kill(pid::Cint, 0::Cint)::Cint) == 0 || return false
+        # a killed worker lingers as a zombie until it's reaped, which doesn't happen while
+        # the manager hangs on exit, as on Julia v1.10
+        stat = try
+            readchomp(`ps -o stat= -p $(pid)`)
+        catch
+            return false
+        end
+        return !startswith(stat, 'Z')
+    end
+end
+
+@testset "workers die with the manager" begin
+    script = """
+        using ParallelTestRunner
+        const Malt = ParallelTestRunner.Malt
+        workers = addworkers(2)
+        for w in workers
+            # keep the worker busy, so that it can't notice by itself that the manager is gone
+            Malt.remote_eval(Main, w.w, :(let x = 0; while true; x += 1; end; end))
+        end
+        println(join([w.w.proc_pid for w in workers], ' '))
+        sleep(600)
+        """
+    cmd = `$(Base.julia_cmd()) --startup-file=no --project=$(Base.active_project()) -e $script`
+    # SIGKILL can't be handled by the manager, so we rely on the OS to kill the workers
+    kill_supported = Sys.islinux() || Sys.isfreebsd() || Sys.iswindows()
+    @testset "$(signum == Base.SIGKILL ? "SIGKILL" : "SIGTERM")" for signum in (Base.SIGTERM, Base.SIGKILL)
+        if signum == Base.SIGKILL && !kill_supported
+            @test_skip false
+            continue
+        end
+        manager = open(pipeline(cmd; stderr=devnull))
+        pids = parse.(Int, split(readline(manager)))
+        try
+            @test length(pids) == 2
+            @test all(pid_running, pids)
+            kill(manager, signum)
+            # on SIGTERM the manager stops its workers without waiting for them too long,
+            # except on Julia v1.10, where it hangs after running its `atexit` hooks
+            if signum == Base.SIGKILL || VERSION >= v"1.11"
+                @test timedwait(() -> process_exited(manager), 15) === :ok
+            end
+            @test timedwait(() -> !any(pid_running, pids), 10) === :ok
+        finally
+            kill(manager, Base.SIGKILL)
+            for pid in filter(pid_running, pids)
+                @static if Sys.iswindows()
+                    run(ignorestatus(`taskkill /F /PID $(pid)`))
+                else
+                    @ccall kill(pid::Cint, Base.SIGKILL::Cint)::Cint
+                end
+            end
+        end
+    end
+end
+
 # All workers must have been stopped once `runtests` returns.
 @testset "no workers running" begin
     children = _count_child_pids()