Wait for all tasks before writing precompile files (#46571)

When the program goes to write out a precompile file, we would like for
the process to first reach a point where it is not still running
background tasks and work. This ensures that the precompile file is in a
consistent state, and isn't forgetting or delaying intended work.

In the future, we may want to add an `atexit` hook (after the other
`atexit` hooks) which optionally calls this function for regular code
too, probably under programmatic control and/or command line argument
control for the user to decide. And we would need to decide how to close
stdin first, so it doesn't continue to keep the program alive.

Add uv_ref and uv_unref internal hooks for this. You probably really
don't want to call these (as they may stop you from getting events on
these objects also), but very specific internal things will need them
for this to work.

Also (mostly unrelated) rewrite a Profile test to conform to best-practices.
Previously, the loop was expecting to observe the Profile test printing
even though nothing kept it alive (there was no reads on stdin).
We fix the design of that test, but also include a patch inside
`jl_process_events` to ensure the loop is alive and will handle events,
to avoid breaking anyone else who was relying on this pattern.

To assist package authors fix errors, we automatically print a note if
this new functionality is causing delays. They then need to ensure they
are calling close explicitly (not relying solely on finalizers), when
appropriate, and are cleaning up other resources (or calling the new
`Base.uv_unref`) also.

Fix #45170

Author: vtjnash

base/libuv.jl

@@ -103,6 +103,9 @@ uv_error(prefix::AbstractString, c::Integer) = c < 0 ? throw(_UVError(prefix, c)
 
 eventloop() = ccall(:jl_global_event_loop, Ptr{Cvoid}, ())
 
+uv_unref(h::Ptr{Cvoid}) = ccall(:uv_unref, Cvoid, (Ptr{Cvoid},), h)
+uv_ref(h::Ptr{Cvoid}) = ccall(:uv_ref, Cvoid, (Ptr{Cvoid},), h)
+
 function process_events()
     return ccall(:jl_process_events, Int32, ())
 end

contrib/generate_precompile.jl

@@ -316,6 +316,7 @@ generate_precompile_statements() = try # Make sure `ansi_enablecursor` is printe
                 current = current == 4 ? 1 : current + 1
             end
         end
+        close(t)
     end
 
     # Collect statements from running the script

src/gc.c

@@ -2958,6 +2958,7 @@ static void jl_gc_queue_thread_local(jl_gc_mark_cache_t *gc_cache, jl_gc_mark_sp
 }
 
 extern jl_value_t *cmpswap_names JL_GLOBALLY_ROOTED;
+extern jl_task_t *wait_empty JL_GLOBALLY_ROOTED;
 
 // mark the initial root set
 static void mark_roots(jl_gc_mark_cache_t *gc_cache, jl_gc_mark_sp_t *sp)
@@ -2996,6 +2997,8 @@ static void mark_roots(jl_gc_mark_cache_t *gc_cache, jl_gc_mark_sp_t *sp)
     gc_mark_queue_obj(gc_cache, sp, jl_emptytuple_type);
     if (cmpswap_names != NULL)
         gc_mark_queue_obj(gc_cache, sp, cmpswap_names);
+    if (wait_empty != NULL)
+        gc_mark_queue_obj(gc_cache, sp, wait_empty);
     gc_mark_queue_obj(gc_cache, sp, jl_global_roots_table);
 }
 

src/jl_uv.c

@@ -30,6 +30,68 @@ extern "C" {
 #endif
 
 static uv_async_t signal_async;
+static uv_timer_t wait_empty_worker;
+
+static void walk_print_cb(uv_handle_t *h, void *arg)
+{
+    if (!uv_is_active(h) || !uv_has_ref(h))
+        return;
+    const char *type = uv_handle_type_name(h->type);
+    if (!type)
+        type = "<unknown>";
+    uv_os_fd_t fd;
+    if (h->type == UV_PROCESS)
+        fd = uv_process_get_pid((uv_process_t*)h);
+    else if (uv_fileno(h, &fd))
+        fd = (uv_os_fd_t)-1;
+    const char *pad = "                "; // 16 spaces
+    int npad = fd == -1 ? 0 : snprintf(NULL, 0, "%zd", (size_t)fd);
+    if (npad < 0)
+        npad = 0;
+    npad += strlen(type);
+    pad += npad < strlen(pad) ? npad : strlen(pad);
+    if (fd == -1)
+        jl_safe_printf(" %s   %s@%p->%p\n", type,             pad, (void*)h, (void*)h->data);
+    else
+        jl_safe_printf(" %s[%zd] %s@%p->%p\n", type, (size_t)fd, pad, (void*)h, (void*)h->data);
+}
+
+static void wait_empty_func(uv_timer_t *t)
+{
+    // make sure this is hidden now, since we would auto-unref it later
+    uv_unref((uv_handle_t*)&signal_async);
+    if (!uv_loop_alive(t->loop))
+        return;
+    jl_safe_printf("\n[pid %zd] waiting for IO to finish:\n"
+                   " TYPE[FD/PID]       @UV_HANDLE_T->DATA\n",
+                   (size_t)uv_os_getpid());
+    uv_walk(jl_io_loop, walk_print_cb, NULL);
+    jl_gc_collect(JL_GC_FULL);
+}
+
+void jl_wait_empty_begin(void)
+{
+    JL_UV_LOCK();
+    if (wait_empty_worker.type != UV_TIMER && jl_io_loop) {
+        // try to purge anything that is just waiting for cleanup
+        jl_io_loop->stop_flag = 0;
+        uv_run(jl_io_loop, UV_RUN_NOWAIT);
+        uv_timer_init(jl_io_loop, &wait_empty_worker);
+        uv_update_time(jl_io_loop);
+        uv_timer_start(&wait_empty_worker, wait_empty_func, 10, 15000);
+        uv_unref((uv_handle_t*)&wait_empty_worker);
+    }
+    JL_UV_UNLOCK();
+}
+
+void jl_wait_empty_end(void)
+{
+    JL_UV_LOCK();
+    uv_close((uv_handle_t*)&wait_empty_worker, NULL);
+    JL_UV_UNLOCK();
+}
+
+
 
 static void jl_signal_async_cb(uv_async_t *hdl)
 {
@@ -49,6 +111,7 @@ jl_mutex_t jl_uv_mutex;
 void jl_init_uv(void)
 {
     uv_async_init(jl_io_loop, &signal_async, jl_signal_async_cb);
+    uv_unref((uv_handle_t*)&signal_async);
     JL_MUTEX_INIT(&jl_uv_mutex); // a file-scope initializer can be used instead
 }
 
@@ -110,7 +173,7 @@ static void jl_uv_closeHandle(uv_handle_t *handle)
         ct->world_age = last_age;
         return;
     }
-    if (handle == (uv_handle_t*)&signal_async)
+    if (handle == (uv_handle_t*)&signal_async || handle == (uv_handle_t*)&wait_empty_worker)
         return;
     free(handle);
 }
@@ -213,7 +276,9 @@ JL_DLLEXPORT int jl_process_events(void)
         if (jl_atomic_load_relaxed(&jl_uv_n_waiters) == 0 && jl_mutex_trylock(&jl_uv_mutex)) {
             JL_PROBE_RT_START_PROCESS_EVENTS(ct);
             loop->stop_flag = 0;
+            uv_ref((uv_handle_t*)&signal_async); // force the loop alive
             int r = uv_run(loop, UV_RUN_NOWAIT);
+            uv_unref((uv_handle_t*)&signal_async);
             JL_PROBE_RT_FINISH_PROCESS_EVENTS(ct);
             JL_UV_UNLOCK();
             return r;

src/julia.h

@@ -1760,6 +1760,7 @@ JL_DLLEXPORT void jl_init_with_image(const char *julia_bindir,
 JL_DLLEXPORT const char *jl_get_default_sysimg_path(void);
 JL_DLLEXPORT int jl_is_initialized(void);
 JL_DLLEXPORT void jl_atexit_hook(int status);
+JL_DLLEXPORT void jl_task_wait_empty(void);
 JL_DLLEXPORT void jl_postoutput_hook(void);
 JL_DLLEXPORT void JL_NORETURN jl_exit(int status);
 JL_DLLEXPORT void JL_NORETURN jl_raise(int signo);

src/partr.c

@@ -281,6 +281,27 @@ static int check_empty(jl_value_t *checkempty)
     return jl_apply_generic(checkempty, NULL, 0) == jl_true;
 }
 
+jl_task_t *wait_empty JL_GLOBALLY_ROOTED;
+void jl_wait_empty_begin(void);
+void jl_wait_empty_end(void);
+
+void jl_task_wait_empty(void)
+{
+    jl_task_t *ct = jl_current_task;
+    if (jl_atomic_load_relaxed(&ct->tid) == 0 && jl_base_module) {
+        jl_wait_empty_begin();
+        jl_value_t *f = jl_get_global(jl_base_module, jl_symbol("wait"));
+        wait_empty = ct;
+        size_t lastage = ct->world_age;
+        ct->world_age = jl_atomic_load_acquire(&jl_world_counter);
+        if (f)
+            jl_apply_generic(f, NULL, 0);
+        ct->world_age = lastage;
+        wait_empty = NULL;
+        jl_wait_empty_end();
+    }
+}
+
 static int may_sleep(jl_ptls_t ptls) JL_NOTSAFEPOINT
 {
     // sleep_check_state is only transitioned from not_sleeping to sleeping
@@ -312,7 +333,7 @@ JL_DLLEXPORT jl_task_t *jl_task_get_next(jl_value_t *trypoptask, jl_value_t *q,
 
         jl_cpu_pause();
         jl_ptls_t ptls = ct->ptls;
-        if (sleep_check_after_threshold(&start_cycles) || (!jl_atomic_load_relaxed(&_threadedregion) && ptls->tid == 0)) {
+        if (sleep_check_after_threshold(&start_cycles) || (ptls->tid == 0 && (!jl_atomic_load_relaxed(&_threadedregion) || wait_empty))) {
             // acquire sleep-check lock
             jl_atomic_store_relaxed(&ptls->sleep_check_state, sleeping);
             jl_fence(); // [^store_buffering_1]
@@ -409,6 +430,14 @@ JL_DLLEXPORT jl_task_t *jl_task_get_next(jl_value_t *trypoptask, jl_value_t *q,
             int8_t gc_state = jl_gc_safe_enter(ptls);
             uv_mutex_lock(&ptls->sleep_lock);
             while (may_sleep(ptls)) {
+                if (ptls->tid == 0 && wait_empty) {
+                    task = wait_empty;
+                    if (jl_atomic_load_relaxed(&ptls->sleep_check_state) != not_sleeping) {
+                        jl_atomic_store_relaxed(&ptls->sleep_check_state, not_sleeping); // let other threads know they don't need to wake us
+                        JL_PROBE_RT_SLEEP_CHECK_TASK_WAKE(ptls);
+                    }
+                    break;
+                }
                 uv_cond_wait(&ptls->wake_signal, &ptls->sleep_lock);
                 // TODO: help with gc work here, if applicable
             }
@@ -417,6 +446,11 @@ JL_DLLEXPORT jl_task_t *jl_task_get_next(jl_value_t *trypoptask, jl_value_t *q,
             JULIA_DEBUG_SLEEPWAKE( ptls->sleep_leave = cycleclock() );
             jl_gc_safe_leave(ptls, gc_state); // contains jl_gc_safepoint
             start_cycles = 0;
+            if (task) {
+                assert(task == wait_empty);
+                wait_empty = NULL;
+                return task;
+            }
         }
         else {
             // maybe check the kernel for new messages too

src/precompile.c

@@ -75,6 +75,8 @@ JL_DLLEXPORT void jl_write_compiler_output(void)
         return;
     }
 
+    jl_task_wait_empty();
+
     if (!jl_module_init_order) {
         jl_printf(JL_STDERR, "WARNING: --output requested, but no modules defined during run\n");
         return;

stdlib/Profile/src/Profile.jl

@@ -156,7 +156,9 @@ function __init__()
     # used, if not manually initialized before that.
     @static if !Sys.iswindows()
         # triggering a profile via signals is not implemented on windows
-        PROFILE_PRINT_COND[] = Base.AsyncCondition()
+        cond = Base.AsyncCondition()
+        Base.uv_unref(cond.handle)
+        PROFILE_PRINT_COND[] = cond
         ccall(:jl_set_peek_cond, Cvoid, (Ptr{Cvoid},), PROFILE_PRINT_COND[].handle)
         errormonitor(Threads.@spawn(profile_printing_listener()))
     end

stdlib/Profile/test/runtests.jl

@@ -198,41 +198,48 @@ if Sys.isbsd() || Sys.islinux()
     @testset "SIGINFO/SIGUSR1 profile triggering" begin
         let cmd = Base.julia_cmd()
             script = """
-                x = rand(1000, 1000)
-                println(stderr, "started")
-                while true
-                    x * x
-                    yield()
-                end
+                print(stderr, "started\n")
+                eof(stdin)
+                close(t)
                 """
             iob = Base.BufferStream()
-            p = run(pipeline(`$cmd -e $script`, stderr = iob, stdout = devnull), wait = false)
+            notify_exit = Base.PipeEndpoint()
+            p = run(pipeline(`$cmd -e $script`, stdin=notify_exit, stderr=iob, stdout=devnull), wait=false)
             t = Timer(120) do t
                 # should be under 10 seconds, so give it 2 minutes then report failure
                 println("KILLING BY PROFILE TEST WATCHDOG\n")
                 kill(p, Base.SIGTERM)
                 sleep(10)
                 kill(p, Base.SIGKILL)
-                close(iob)
+                close(p)
             end
             try
-                s = readuntil(iob, "started", keep = true)
+                s = readuntil(iob, "started", keep=true)
                 @assert occursin("started", s)
                 @assert process_running(p)
-                for _ in 1:2
-                    sleep(2.5)
+                for i in 1:2
+                    i > 1 && sleep(5)
                     if Sys.isbsd()
                         kill(p, 29) # SIGINFO
                     elseif Sys.islinux()
                         kill(p, 10) # SIGUSR1
                     end
-                    s = readuntil(iob, "Overhead ╎", keep = true)
+                    s = readuntil(iob, "Overhead ╎", keep=true)
                     @test process_running(p)
+                    readavailable(iob)
                     @test occursin("Overhead ╎", s)
                 end
-            finally
-                kill(p, Base.SIGKILL)
+                close(notify_exit) # notify test finished
+                s = read(iob, String) # consume test output
+                wait(p) # wait for test completion
+                close(t)
+            catch
+                close(notify_exit)
+                errs = read(iob, String) # consume test output
+                isempty(errs) || println("CHILD STDERR after test failure: ", errs)
+                wait(p) # wait for test completion
                 close(t)
+                rethrow()
             end
         end
     end

test/precompile.jl

@@ -261,6 +261,10 @@ precompile_test_harness(false) do dir
 
               # check that @ccallable works from precompiled modules
               Base.@ccallable Cint f35014(x::Cint) = x+Cint(1)
+
+              # check that Tasks work from serialized state
+              ch1 = Channel(x -> nothing)
+              ch2 = Channel(x -> (push!(x, 2); nothing), Inf)
           end
           """)
     # Issue #12623
@@ -310,6 +314,13 @@ precompile_test_harness(false) do dir
         @test Foo.layout2 == Any[Ptr{Int8}(0), Ptr{Int16}(0), Ptr{Int32}(-1)]
         @test typeof.(Foo.layout2) == [Ptr{Int8}, Ptr{Int16}, Ptr{Int32}]
         @test Foo.layout3 == ["ab", "cd", "ef", "gh", "ij"]
+
+        @test !isopen(Foo.ch1)
+        @test !isopen(Foo.ch2)
+        @test !isready(Foo.ch1)
+        @test isready(Foo.ch2)
+        @test take!(Foo.ch2) === 2
+        @test !isready(Foo.ch2)
     end
 
     @eval begin function ccallable_test()