Save thread context before yielding for GC

src/gc-interface.h

@@ -102,6 +102,9 @@ JL_DLLEXPORT int gc_is_collector_thread(int tid) JL_NOTSAFEPOINT;
 JL_DLLEXPORT unsigned char jl_gc_pin_object(void* obj);
 // Returns the version of which GC implementation is being used according to the list of supported GCs
 JL_DLLEXPORT const char* jl_active_gc_impl(void);
+// Notifies the GC that the given thread is about to yield for a GC. ctx is the ucontext for the thread
+// if it is already fetched by the caller, otherwise it is NULL.
+JL_DLLEXPORT void jl_gc_notify_thread_yield(jl_ptls_t ptls, void* ctx);
 
 // TODO: The preserve hook functions may be temporary. We should see the performance impact of the change.
 

src/gc-mmtk.c

@@ -270,6 +270,7 @@ JL_DLLEXPORT void jl_mmtk_prepare_to_collect(void)
     gc_num.total_time_to_safepoint += duration;
 
     if (!jl_atomic_load_acquire(&jl_gc_disable_counter)) {
+        jl_gc_notify_thread_yield(ptls, NULL);
         JL_LOCK_NOGC(&finalizers_lock); // all the other threads are stopped, so this does not make sense, right? otherwise, failing that, this seems like plausibly a deadlock
 #ifndef __clang_gcanalyzer__
         mmtk_block_thread_for_gc();
@@ -303,6 +304,19 @@ JL_DLLEXPORT unsigned char jl_gc_pin_object(void* obj) {
     return mmtk_pin_object(obj);
 }
 
+JL_DLLEXPORT void jl_gc_notify_thread_yield(jl_ptls_t ptls, void* ctx) {
+    if (ctx == NULL) {
+        // Save the context for the thread as it was running at the time of the call
+        int r = getcontext(&ptls->gc_tls.ctx_at_the_time_gc_started);
+        if (r == -1) {
+            jl_safe_printf("Failed to save context for conservative scanning\n");
+            abort();
+        }
+        return;
+    }
+    memcpy(&ptls->gc_tls.ctx_at_the_time_gc_started, ctx, sizeof(ucontext_t));
+}
+
 // ========================================================================= //
 // GC Statistics
 // ========================================================================= //

src/gc-stock.c

@@ -2930,6 +2930,7 @@ size_t jl_maxrss(void);
 // Only one thread should be running in this function
 static int _jl_gc_collect(jl_ptls_t ptls, jl_gc_collection_t collection)
 {
+    jl_gc_notify_thread_yield(ptls, NULL);
     combine_thread_gc_counts(&gc_num, 1);
 
     // We separate the update of the graph from the update of live_bytes here
@@ -4022,6 +4023,10 @@ JL_DLLEXPORT const char* jl_active_gc_impl(void) {
     return "";
 }
 
+JL_DLLEXPORT void jl_gc_notify_thread_yield(jl_ptls_t ptls, void* ctx) {
+    // Do nothing before a thread yields
+}
+
 #ifdef __cplusplus
 }
 #endif

src/gc-tls-mmtk.h

@@ -10,6 +10,7 @@ extern "C" {
 typedef struct {
     MMTkMutatorContext mmtk_mutator;
     size_t malloc_sz_since_last_poll;
+    ucontext_t ctx_at_the_time_gc_started;
 } jl_gc_tls_states_t;
 
 #ifdef __cplusplus

src/signals-unix.c

@@ -388,6 +388,8 @@ JL_NO_ASAN static void segv_handler(int sig, siginfo_t *info, void *context)
         return;
     }
     if (sig == SIGSEGV && info->si_code == SEGV_ACCERR && jl_addr_is_safepoint((uintptr_t)info->si_addr) && !is_write_fault(context)) {
+        // TODO: We should do the same for other platforms
+        jl_gc_notify_thread_yield(ct->ptls, context);
         jl_set_gc_and_wait();
         // Do not raise sigint on worker thread
         if (jl_atomic_load_relaxed(&ct->tid) != 0)