make threaded exit and cache file loading safer by pausing other threads (JuliaLang#55105)

Author: vtjnash

src/init.c

@@ -334,15 +334,15 @@ JL_DLLEXPORT void jl_atexit_hook(int exitcode) JL_NOTSAFEPOINT_ENTER
                          // we would like to guarantee this, but cannot currently, so there is still a small race window
                          // that needs to be fixed in libuv
     }
-    if (ct)
-        (void)jl_gc_safe_enter(ct->ptls); // park in gc-safe
     if (loop != NULL) {
         // TODO: consider uv_loop_close(loop) here, before shutdown?
         uv_library_shutdown();
         // no JL_UV_UNLOCK(), since it is now torn down
     }
-
-    // TODO: Destroy threads?
+    if (ct)
+        jl_safepoint_suspend_all_threads(ct); // Destroy other threads, so that they don't segfault
+    if (ct)
+        (void)jl_gc_safe_enter(ct->ptls); // park in gc-safe
 
     jl_destroy_timing(); // cleans up the current timing_stack for noreturn
 #ifdef USE_TIMING_COUNTS

src/julia.h

@@ -1107,6 +1107,8 @@ STATIC_INLINE void jl_gc_multi_wb(const void *parent, const jl_value_t *ptr) JL_
 JL_DLLEXPORT void *jl_gc_managed_malloc(size_t sz);
 JL_DLLEXPORT void jl_gc_safepoint(void);
 JL_DLLEXPORT int jl_safepoint_suspend_thread(int tid, int waitstate);
+JL_DLLEXPORT void jl_safepoint_suspend_all_threads(struct _jl_task_t *ct);
+JL_DLLEXPORT void jl_safepoint_resume_all_threads(struct _jl_task_t *ct);
 JL_DLLEXPORT int jl_safepoint_resume_thread(int tid) JL_NOTSAFEPOINT;
 
 void *mtarraylist_get(small_arraylist_t *_a, size_t idx) JL_NOTSAFEPOINT;

src/julia_internal.h

@@ -959,6 +959,7 @@ extern JL_DLLEXPORT ssize_t jl_tls_offset;
 extern JL_DLLEXPORT const int jl_tls_elf_support;
 void jl_init_threading(void);
 void jl_start_threads(void);
+extern uv_mutex_t safepoint_lock;
 
 // Whether the GC is running
 extern char *jl_safepoint_pages;

src/safepoint.c

@@ -267,12 +267,12 @@ void jl_safepoint_wait_thread_resume(void)
         uv_cond_broadcast(&safepoint_cond_begin);
         uv_mutex_unlock(&safepoint_lock);
         uv_mutex_lock(&ct->ptls->sleep_lock);
+        while (jl_atomic_load_relaxed(&ct->ptls->suspend_count))
+            uv_cond_wait(&ct->ptls->wake_signal, &ct->ptls->sleep_lock);
     }
-    while (jl_atomic_load_relaxed(&ct->ptls->suspend_count))
-        uv_cond_wait(&ct->ptls->wake_signal, &ct->ptls->sleep_lock);
-    // must while still holding the mutex_unlock, so we know other threads in
-    // jl_safepoint_suspend_thread will observe this thread in the correct GC
-    // state, and not still stuck in JL_GC_STATE_WAITING
+    // must exit gc while still holding the mutex_unlock, so we know other
+    // threads in jl_safepoint_suspend_thread will observe this thread in the
+    // correct GC state, and not still stuck in JL_GC_STATE_WAITING
     jl_atomic_store_release(&ct->ptls->gc_state, state);
     uv_mutex_unlock(&ct->ptls->sleep_lock);
 }
@@ -290,12 +290,20 @@ int jl_safepoint_suspend_thread(int tid, int waitstate)
     if (0 > tid || tid >= jl_atomic_load_acquire(&jl_n_threads))
         return 0;
     jl_ptls_t ptls2 = jl_atomic_load_relaxed(&jl_all_tls_states)[tid];
+    jl_task_t *ct2 = ptls2 ? jl_atomic_load_relaxed(&ptls2->current_task) : NULL;
+    if (ct2 == NULL) {
+        // this thread is not alive yet or already dead
+        return 0;
+    }
+    uv_mutex_lock(&safepoint_lock);
     uv_mutex_lock(&ptls2->sleep_lock);
     int16_t suspend_count = jl_atomic_load_relaxed(&ptls2->suspend_count) + 1;
     jl_atomic_store_relaxed(&ptls2->suspend_count, suspend_count);
     if (suspend_count == 1) { // first to suspend
         jl_safepoint_enable(3);
         jl_atomic_store_relaxed(&ptls2->safepoint, (size_t*)(jl_safepoint_pages + jl_page_size * 3 + sizeof(void*)));
+        if (jl_atomic_load(&_threadedregion) != 0 || tid == jl_atomic_load_relaxed(&io_loop_tid))
+            jl_wake_libuv(); // our integration with libuv right now doesn't handle except by waking it
     }
     uv_mutex_unlock(&ptls2->sleep_lock);
     if (waitstate) {
@@ -305,17 +313,20 @@ int jl_safepoint_suspend_thread(int tid, int waitstate)
             // not, so assume it is running GC and wait for GC to finish first.
             // It will be unable to reenter helping with GC because we have
             // changed its safepoint page.
+            uv_mutex_unlock(&safepoint_lock);
             jl_set_gc_and_wait();
+            uv_mutex_lock(&safepoint_lock);
         }
         while (jl_atomic_load_acquire(&ptls2->suspend_count) != 0) {
             int8_t state2 = jl_atomic_load_acquire(&ptls2->gc_state);
             if (waitstate <= 2 && state2 != JL_GC_STATE_UNSAFE)
                 break;
             if (waitstate == 3 && state2 == JL_GC_STATE_WAITING)
                 break;
-            jl_cpu_pause(); // yield (wait for safepoint_cond_begin, for example)?
+            uv_cond_wait(&safepoint_cond_begin, &safepoint_lock);
         }
     }
+    uv_mutex_unlock(&safepoint_lock);
     return suspend_count;
 }
 
@@ -326,6 +337,11 @@ int jl_safepoint_resume_thread(int tid) JL_NOTSAFEPOINT
     if (0 > tid || tid >= jl_atomic_load_acquire(&jl_n_threads))
         return 0;
     jl_ptls_t ptls2 = jl_atomic_load_relaxed(&jl_all_tls_states)[tid];
+    jl_task_t *ct2 = ptls2 ? jl_atomic_load_relaxed(&ptls2->current_task) : NULL;
+    if (ct2 == NULL) {
+        // this thread is not alive yet or already dead
+        return 0;
+    }
     uv_mutex_lock(&safepoint_lock);
     uv_mutex_lock(&ptls2->sleep_lock);
     int16_t suspend_count = jl_atomic_load_relaxed(&ptls2->suspend_count);
@@ -338,6 +354,7 @@ int jl_safepoint_resume_thread(int tid) JL_NOTSAFEPOINT
 #ifdef _OS_DARWIN_
         jl_safepoint_resume_thread_mach(ptls2, tid);
 #endif
+        uv_cond_broadcast(&safepoint_cond_begin);
     }
     if (suspend_count != 0) {
         jl_atomic_store_relaxed(&ptls2->suspend_count, suspend_count - 1);

src/signals-mach.c

@@ -45,7 +45,6 @@ static void attach_exception_port(thread_port_t thread, int segv_only);
 
 // low 16 bits are the thread id, the next 8 bits are the original gc_state
 static arraylist_t suspended_threads;
-extern uv_mutex_t safepoint_lock;
 extern uv_cond_t safepoint_cond_begin;
 
 #define GC_STATE_SHIFT 8*sizeof(int16_t)

src/staticdata.c

@@ -3132,6 +3132,7 @@ static void jl_restore_system_image_from_stream_(ios_t *f, jl_image_t *image, jl
                                                  jl_array_t **ext_targets, jl_array_t **edges,
                                                  char **base, arraylist_t *ccallable_list, pkgcachesizes *cachesizes) JL_GC_DISABLED
 {
+    jl_task_t *ct = jl_current_task;
     int en = jl_gc_enable(0);
     ios_t sysimg, const_data, symbols, relocs, gvar_record, fptr_record;
     jl_serializer_state s = {0};
@@ -3143,7 +3144,7 @@ static void jl_restore_system_image_from_stream_(ios_t *f, jl_image_t *image, jl
     s.relocs = &relocs;
     s.gvar_record = &gvar_record;
     s.fptr_record = &fptr_record;
-    s.ptls = jl_current_task->ptls;
+    s.ptls = ct->ptls;
     jl_value_t **const*const tags = get_tags();
     htable_t new_dt_objs;
     htable_new(&new_dt_objs, 0);
@@ -3316,6 +3317,7 @@ static void jl_restore_system_image_from_stream_(ios_t *f, jl_image_t *image, jl
     arraylist_new(&cleanup_list, 0);
     arraylist_t delay_list;
     arraylist_new(&delay_list, 0);
+    JL_LOCK(&typecache_lock); // Might GC--prevent other threads from changing any type caches while we inspect them all
     for (size_t i = 0; i < s.uniquing_types.len; i++) {
         uintptr_t item = (uintptr_t)s.uniquing_types.items[i];
         // check whether we are operating on the typetag
@@ -3443,6 +3445,7 @@ static void jl_restore_system_image_from_stream_(ios_t *f, jl_image_t *image, jl
     }
     arraylist_grow(&cleanup_list, -cleanup_list.len);
     // finally cache all our new types now
+    jl_safepoint_suspend_all_threads(ct); // past this point, it is now not safe to observe the intermediate states on other threads via reflection, so temporarily pause those
     for (size_t i = 0; i < new_dt_objs.size; i += 2) {
         void *dt = table[i + 1];
         if (dt != HT_NOTFOUND) {
@@ -3456,6 +3459,8 @@ static void jl_restore_system_image_from_stream_(ios_t *f, jl_image_t *image, jl
         assert(jl_is_datatype(obj));
         jl_cache_type_((jl_datatype_t*)obj);
     }
+    JL_UNLOCK(&typecache_lock); // Might GC
+    jl_safepoint_resume_all_threads(ct); // TODO: move this later to also protect MethodInstance allocations, but we would need to acquire all jl_specializations_get_linfo and jl_module_globalref locks, which is hard
     // Perform fixups: things like updating world ages, inserting methods & specializations, etc.
     for (size_t i = 0; i < s.uniquing_objs.len; i++) {
         uintptr_t item = (uintptr_t)s.uniquing_objs.items[i];

src/threading.c

@@ -432,6 +432,30 @@ JL_DLLEXPORT jl_gcframe_t **jl_adopt_thread(void)
     return &ct->gcstack;
 }
 
+
+void jl_safepoint_suspend_all_threads(jl_task_t *ct)
+{
+    // TODO: prevent jl_n_threads changing or jl_safepoint_resume_thread calls on another thread
+    //uv_mutex_lock(&tls_lock);
+    //disallow_resume = ct->tid;
+    //uv_mutex_unlock(&tls_lock);
+    for (int16_t tid = 0; tid < jl_atomic_load_relaxed(&jl_n_threads); tid++) {
+        if (tid != jl_atomic_load_relaxed(&ct->tid))
+            jl_safepoint_suspend_thread(tid, 1);
+    };
+}
+
+void jl_safepoint_resume_all_threads(jl_task_t *ct)
+{
+    //uv_mutex_lock(&tls_lock);
+    //if (disallow_resume != ct->tid) return;
+    //uv_mutex_unlock(&tls_lock);
+    for (int16_t tid = 0; tid < jl_atomic_load_relaxed(&jl_n_threads); tid++) {
+        if (tid != jl_atomic_load_relaxed(&ct->tid))
+            jl_safepoint_resume_thread(tid);
+    };
+}
+
 void jl_task_frame_noreturn(jl_task_t *ct) JL_NOTSAFEPOINT;
 void scheduler_delete_thread(jl_ptls_t ptls) JL_NOTSAFEPOINT;
 

test/atexit.jl

@@ -220,7 +220,7 @@ using Test
                     # Block until the atexit hooks have all finished. We use a manual "spin
                     # lock" because task switch is disallowed inside the finalizer, below.
                     atexit_has_finished[] = 1
-                    while atexit_has_finished[] == 1 end
+                    while atexit_has_finished[] == 1; GC.safepoint(); end
                     try
                         # By the time this runs, all the atexit hooks will be done.
                         # So this will throw.
@@ -232,7 +232,7 @@ using Test
                         exit(22)
                     end
                 end
-                while atexit_has_finished[] == 0 end
+                while atexit_has_finished[] == 0; GC.safepoint(); end
             end
             # Finalizers run after the atexit hooks, so this blocks exit until the spawned
             # task above gets a chance to run.
@@ -241,7 +241,7 @@ using Test
                 # Allow the spawned task to finish
                 atexit_has_finished[] = 2
                 # Then spin forever to prevent exit.
-                while atexit_has_finished[] == 2 end
+                while atexit_has_finished[] == 2; GC.safepoint(); end
             end
             exit(0)
             """ => 22,