From 8f296a1d1c71077fab6e6d05a1f09aa30e70fdb7 Mon Sep 17 00:00:00 2001
From: Yi Lin <qinsoon@gmail.com>
Date: Thu, 5 Dec 2024 03:33:39 +0000
Subject: [PATCH] Save thread context before yielding for GC

---
 src/gc-interface.h |  3 +++
 src/gc-mmtk.c      | 14 ++++++++++++++
 src/gc-stock.c     |  5 +++++
 src/gc-tls-mmtk.h  |  1 +
 src/signals-unix.c |  2 ++
 5 files changed, 25 insertions(+)

diff --git a/src/gc-interface.h b/src/gc-interface.h
index cc9ee947798fe..9931bec35f2f6 100644
--- a/src/gc-interface.h
+++ b/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.
 
diff --git a/src/gc-mmtk.c b/src/gc-mmtk.c
index 663995563d6af..19ba77da9da93 100644
--- a/src/gc-mmtk.c
+++ b/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
 // ========================================================================= //
diff --git a/src/gc-stock.c b/src/gc-stock.c
index 01ad8ea1a725a..33edf09b5787d 100644
--- a/src/gc-stock.c
+++ b/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
diff --git a/src/gc-tls-mmtk.h b/src/gc-tls-mmtk.h
index 01791345b6984..309ab64f3b86b 100644
--- a/src/gc-tls-mmtk.h
+++ b/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
diff --git a/src/signals-unix.c b/src/signals-unix.c
index caf0e977929c5..e92056cf23845 100644
--- a/src/signals-unix.c
+++ b/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)