thread safety: lock around atexit_hooks global (#49868)

Ensure the lock is precise, so that we are allowed to register new
atexit hooks from inside an atexit hook. But then disable `atexit()`
when shutting down after it finishes running.

Add tests that cover all the cases:
1. registering a hook from inside a hook
2. registering a hook from another thread while hooks are running
3. attempting to register a hook after all hooks have finished
   (disallowed)

Fixes #49841
Co-authored-by: Jameson Nash <vtjnash@gmail.com>

base/initdefs.jl

@@ -354,6 +354,7 @@ const atexit_hooks = Callable[
     () -> Filesystem.temp_cleanup_purge(force=true)
 ]
 const _atexit_hooks_lock = ReentrantLock()
+global _atexit_hooks_finished::Bool = false
 
 """
     atexit(f)
@@ -374,12 +375,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(exitcode::Cint)
-    while !isempty(atexit_hooks)
-        f = popfirst!(atexit_hooks)
+    # 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
             if hasmethod(f, (Cint,))
                 f(exitcode)

test/atexit.jl

@@ -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
@@ -172,5 +174,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