From 8dc55350ad1d770cab3983fb69dabb7dfe999a8c Mon Sep 17 00:00:00 2001
From: Yi Lin <qinsoon@gmail.com>
Date: Fri, 6 Dec 2024 13:09:26 +1300
Subject: [PATCH] Save thread context before yielding for GC (#78)

This PR ports https://github.com/mmtk/mmtk-julia/pull/159 to `dev`. The difference is that this PR adds a general call to the GC interface `jl_gc_notify_thread_yield`. In this case, each GC will do what they need in the call, and the context is saved in the GC specific TLS.
---
 src/gc-interface.h |  3 +++
 src/gc-mmtk.c      | 15 +++++++++++++++
 src/gc-stock.c     |  9 +++++++++
 src/gc-tls-mmtk.h  |  1 +
 src/signals-unix.c |  2 ++
 5 files changed, 30 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..9b83ed5bde28b 100644
--- a/src/gc-mmtk.c
+++ b/src/gc-mmtk.c
@@ -270,6 +270,8 @@ 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)) {
+        // This thread will yield.
+        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 +305,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..14780f8c1f6a3 100644
--- a/src/gc-stock.c
+++ b/src/gc-stock.c
@@ -3361,6 +3361,11 @@ JL_DLLEXPORT void jl_gc_collect(jl_gc_collection_t collection)
         gc_cblist_pre_gc, (collection));
 
     if (!jl_atomic_load_acquire(&jl_gc_disable_counter)) {
+        // This thread will yield.
+        // jl_gc_notify_thread_yield does nothing for the stock GC at the point, but it may be non empty in the future,
+        // and this is a place where we should call jl_gc_notify_thread_yield.
+        // TODO: This call can be removed if requested.
+        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__
         if (_jl_gc_collect(ptls, collection)) {
@@ -4022,6 +4027,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)