Skip to content

Commit

Permalink
Lock the atexit_hooks during execution of the hooks on shutdown.
Browse files Browse the repository at this point in the history
Fixes #49841.

Follow-up to #49774.

This PR makes two changes:
1. It locks `atexit_hooks` while iterating the hooks during execution of `_atexit()` at shutdown.
    - This prevents any data races if another Task is registering a new atexit hook while the hooks are being evaluated.
2. It defines semantics for what happens if another Task attempts to register another atexit hook _after all the hooks have finished_, and we've proceeded on to the rest of shutdown.
    - Previously, those atexit hooks would be _ignored,_ which violates the user's expectations and violates the "atexit" contract.
    - Now, the attempt to register the atexit hook will **throw an exception,** which ensures that we never break our promise, since the user was never able to register the atexit hook at all.
    - This does mean that users will need to handle the thrown exception and likely do now whatever tear down they were hoping to delay until exit.
  • Loading branch information
NHDaly committed Jul 19, 2023
1 parent 8c739b1 commit 54ae3b6
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 6 deletions.
37 changes: 33 additions & 4 deletions base/initdefs.jl
Original file line number Diff line number Diff line change
Expand Up @@ -351,6 +351,7 @@ const atexit_hooks = Callable[
() -> Filesystem.temp_cleanup_purge(force=true)
]
const _atexit_hooks_lock = ReentrantLock()
global _atexit_hooks_finished::Bool = false

"""
atexit(f)
Expand All @@ -363,12 +364,40 @@ exit code `n` (instead of the original exit code). If more than one exit hook
calls `exit(n)`, then Julia will exit with the exit code corresponding to the
last called exit hook that calls `exit(n)`. (Because exit hooks are called in
LIFO order, "last called" is equivalent to "first registered".)
Note: Once all exit hooks have been called, no more exit hooks can be registered,
and any call to `atexit(f)` after all hooks have completed will throw an exception.
This situation may occur if you are registering exit hooks from background Tasks that
may still be executing concurrently during shutdown.
"""
atexit(f::Function) = Base.@lock _atexit_hooks_lock (pushfirst!(atexit_hooks, f); nothing)
function atexit(f::Function)
Base.@lock _atexit_hooks_lock begin
_atexit_hooks_finished && error("cannot register new atexit hook; already exiting.")
pushfirst!(atexit_hooks, f)
return nothing
end
end

function _atexit()
while !isempty(atexit_hooks)
f = popfirst!(atexit_hooks)
function _atexit(exitcode::Cint)
# Don't hold the lock around the iteration, just in case any other thread executing in
# parallel tries to register a new atexit hook while this is running. We don't want to
# block that thread from proceeding, and we can allow it to register its hook which we
# will immediately run here.
while true
local f
Base.@lock _atexit_hooks_lock begin
# If this is the last iteration, atomically disable atexit hooks to prevent
# someone from registering a hook that will never be run.
# (We do this inside the loop, so that it is atomic: no one can have registered
# a hook that never gets run, and we run all the hooks we know about until
# the vector is empty.)
if isempty(atexit_hooks)
global _atexit_hooks_finished = true
break
end

f = popfirst!(atexit_hooks)
end
try
f()
catch ex
Expand Down
99 changes: 97 additions & 2 deletions test/atexit.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,19 @@ using Test

@testset "atexit.jl" begin
function _atexit_tests_gen_cmd_eval(expr::String)
# We run the atexit tests with 2 threads, for the parallelism tests at the end.
cmd_eval = ```
$(Base.julia_cmd()) -e $(expr)
$(Base.julia_cmd()) -t2 -e $(expr)
```
return cmd_eval
end
function _atexit_tests_gen_cmd_script(temp_dir::String, expr::String)
script, io = mktemp(temp_dir)
println(io, expr)
close(io)
# We run the atexit tests with 2 threads, for the parallelism tests at the end.
cmd_script = ```
$(Base.julia_cmd()) $(script)
$(Base.julia_cmd()) -t2 $(script)
```
return cmd_script
end
Expand Down Expand Up @@ -150,5 +152,98 @@ using Test
@test p_script.exitcode == expected_exit_code
end
end
@testset "test calling atexit() in parallel with running atexit hooks." begin
# These tests cover 3 parallelism cases, as described by the following comments.
julia_expr_list = Dict(
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# 1. registering a hook from inside a hook
"""
atexit() do
atexit() do
exit(11)
end
end
# This will attempt to exit 0, but the execution of the atexit hook will
# register another hook, which will exit 11.
exit(0)
""" => 11,
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# 2. registering a hook from another thread while hooks are running
"""
c = Channel()
# This hook must execute _last_. (Execution is LIFO.)
atexit() do
put!(c, nothing)
put!(c, nothing)
end
atexit() do
# This will run in a concurrent task, testing that we can register atexit
# hooks from another task while running atexit hooks.
Threads.@spawn begin
Core.println("INSIDE")
take!(c) # block on c
Core.println("go")
atexit() do
Core.println("exit11")
exit(11)
end
take!(c) # keep the _atexit() loop alive until we've added another item.
Core.println("done")
end
end
exit(0)
""" => 11,
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# 3. attempting to register a hook after all hooks have finished (disallowed)
"""
const atexit_has_finished = Threads.Atomic{Bool}(false)
atexit() do
Threads.@spawn begin
# Block until the atexit hooks have all finished. We use a manual "spin
# lock" because task switch is disallowed inside the finalizer, below.
while !atexit_has_finished[] end
Core.println("done")
try
# By the time this runs, all the atexit hooks will be done.
# So this will throw.
atexit() do
exit(11)
end
catch
# Meaning we _actually_ exit 22.
exit(22)
end
end
end
# Finalizers run after the atexit hooks, so this blocks exit until the spawned
# task above gets a chance to run.
x = []
finalizer(x) do x
Core.println("FINALIZER")
# Allow the spawned task to finish
atexit_has_finished[] = true
Core.println("ready")
# Then spin forever to prevent exit.
while atexit_has_finished[] end
Core.println("exiting")
end
exit(0)
""" => 22,
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
)
for julia_expr in keys(julia_expr_list)
cmd_eval = _atexit_tests_gen_cmd_eval(julia_expr)
cmd_script = _atexit_tests_gen_cmd_script(atexit_temp_dir, julia_expr)
expected_exit_code = julia_expr_list[julia_expr]
@test_throws(ProcessFailedException, run(cmd_eval))
@test_throws(ProcessFailedException, run(cmd_script))
p_eval = run(cmd_eval; wait = false)
p_script = run(cmd_script; wait = false)
wait(p_eval)
wait(p_script)
@test p_eval.exitcode == expected_exit_code
@test p_script.exitcode == expected_exit_code
end
end
rm(atexit_temp_dir; force = true, recursive = true)
end

0 comments on commit 54ae3b6

Please sign in to comment.