From a6a37142f61937eabde80f2208a0ab12cf665f0b Mon Sep 17 00:00:00 2001
From: Jake Hughes <jh@jakehughes.uk>
Date: Wed, 14 Feb 2024 11:05:26 +0000
Subject: [PATCH] Add support for scanning compiler thread locals

Rust's fast TLS implementation uses LLVM generated thread-locals where
thread local data is stored in the PT_TLS segment [1].

To allow the BDWGC to scan this, we must locate the PT_TLS block for
each thread which is suspended during GC, and then add this to the range
of values which need marking.
---
 include/private/pthread_support.h |  6 ++++
 pthread_stop_world.c              | 54 +++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/include/private/pthread_support.h b/include/private/pthread_support.h
index 97ef42479..4ce2ac35d 100644
--- a/include/private/pthread_support.h
+++ b/include/private/pthread_support.h
@@ -44,6 +44,11 @@
 
 EXTERN_C_BEGIN
 
+typedef struct GC_ThreadLocalRoots {
+    ptr_t start;
+    ptr_t end;
+} tlr;
+
 typedef struct GC_StackContext_Rep {
 # if defined(THREAD_SANITIZER) && defined(SIGNAL_BASED_STOP_WORLD)
     char dummy[sizeof(oh)];     /* A dummy field to avoid TSan false    */
@@ -113,6 +118,7 @@ typedef struct GC_StackContext_Rep {
                                 /* stack (thread); may be NULL.         */
 
   ptr_t tls_rootset;
+  tlr compiler_thread_roots;
 } *GC_stack_context_t;
 
 #ifdef GC_WIN32_THREADS
diff --git a/pthread_stop_world.c b/pthread_stop_world.c
index 36065bb2c..a77de32dc 100644
--- a/pthread_stop_world.c
+++ b/pthread_stop_world.c
@@ -33,6 +33,8 @@
 # include <alloca.h>
 #endif
 
+# include <link.h>
+
 GC_INLINE void GC_usleep(unsigned us)
 {
 #   if defined(LINT2) || defined(THREAD_SANITIZER)
@@ -325,10 +327,53 @@ GC_INLINE void GC_store_stack_ptr(GC_stack_context_t crtn)
 # endif
 }
 
+static int find_segment(struct dl_phdr_info *info, size_t size, void *data) {
+    UNUSED_ARG(size);
+    tlr * roots = (tlr *) data;
+
+    for (size_t i = 0; i < info->dlpi_phnum; i ++) {
+        if ( info -> dlpi_phdr[i].p_type != PT_TLS )
+            continue;
+
+        if ( info -> dlpi_tls_data == NULL )
+            /* This SO has no thread locals */
+            return 0;
+
+        size_t memsz = info -> dlpi_phdr[i].p_memsz;
+        ptr_t start = info -> dlpi_tls_data;
+
+        GC_ASSERT(memsz > 0);
+
+        roots -> start = start;
+        roots -> end = start + memsz;
+        return 1;
+    }
+    return 0;
+}
+
+
+/* Get the TLS roots for the current thread.                                  */
+/*                                                                            */
+/* This works because a Rust program (and any shared objects) use the PT_TLS  */
+/* segment in the binary to store every thread's local instance of each       */
+/* thread-local variable. For each thread, these instances are stored at a    */
+/* fixed offset inside the same PT_TLS segment [1].                           */
+/*                                                                            */
+/* This returns the ranges inside the `PT_TLS` segment which contains the     */
+/* thread-local instances for the current thread only.                        */
+/*                                                                            */
+/* [1]: https://www.akkadia.org/drepper/tls.pdf                               */
+void get_thread_local_roots(tlr * roots)
+{
+    dl_iterate_phdr(find_segment, roots);
+    return;
+}
+
 STATIC void GC_suspend_handler_inner(ptr_t dummy, void *context)
 {
   GC_thread me;
   GC_stack_context_t crtn;
+  struct GC_ThreadLocalRoots tlr;
 # ifdef E2K
     ptr_t bs_lo;
     size_t stack_size;
@@ -371,6 +416,9 @@ STATIC void GC_suspend_handler_inner(ptr_t dummy, void *context)
   crtn = me -> crtn;
   GC_store_stack_ptr(crtn);
   crtn -> tls_rootset = GC_tls_rootset();
+  get_thread_local_roots(&tlr);
+  crtn -> compiler_thread_roots = tlr;
+
 # ifdef E2K
     GC_ASSERT(NULL == crtn -> backing_store_end);
     GET_PROCEDURE_STACK_LOCAL(&bs_lo, &stack_size);
@@ -801,7 +849,10 @@ GC_INNER void GC_push_all_stacks(void)
     // perform null check.
     GC_thread me = GC_self_thread_inner();
     if (me != NULL) {
+        struct GC_ThreadLocalRoots tlr;
+        get_thread_local_roots(&tlr);
         me -> crtn -> tls_rootset = GC_tls_rootset();
+        me -> crtn -> compiler_thread_roots = tlr;
     }
 
     GC_ASSERT(I_HOLD_LOCK());
@@ -887,6 +938,9 @@ GC_INNER void GC_push_all_stacks(void)
           if (GC_sp_corrector != 0)
             GC_sp_corrector((void **)&lo, (void *)(p -> id));
 #       endif
+        /* Scan the TLS roots */
+        GC_push_all_eager(crtn->compiler_thread_roots.start, crtn->compiler_thread_roots.end);
+
         GC_push_all_stack_sections(lo, hi, traced_stack_sect);
 #       ifdef DEBUG_THREADS
         GC_log_printf("Pushing TLS rootset from thread %p\n",