Fix GC stack scrub when the GC is running on a worker thread

Also scrub the stacks of all tasks instead of only the currently running one.

src/gc-debug.c

@@ -257,7 +257,6 @@ JL_DLLEXPORT jl_gc_debug_env_t jl_gc_debug_env = {
     {0, UINT64_MAX, 0, 0, 0, {0, 0, 0}},
     {0, UINT64_MAX, 0, 0, 0, {0, 0, 0}}
 };
-static char *gc_stack_lo;
 
 static void gc_debug_alloc_setnext(jl_alloc_num_t *num)
 {
@@ -345,11 +344,26 @@ void gc_debug_print(void)
     gc_debug_print_status();
 }
 
-static void gc_scrub_range(char *stack_lo, char *stack_hi)
+// a list of tasks for conservative stack scan during gc_scrub
+static arraylist_t jl_gc_debug_tasks;
+
+void gc_scrub_record_task(jl_task_t *t)
+{
+    arraylist_push(&jl_gc_debug_tasks, t);
+}
+
+static void gc_scrub_range(char *low, char *high)
 {
-    stack_lo = (char*)((uintptr_t)stack_lo & ~(uintptr_t)15);
-    for (char **stack_p = (char**)stack_lo;
-         stack_p > (char**)stack_hi;stack_p--) {
+    jl_ptls_t ptls = jl_get_ptls_states();
+    jl_jmp_buf *old_buf = ptls->safe_restore;
+    jl_jmp_buf buf;
+    if (jl_setjmp(buf, 0)) {
+        ptls->safe_restore = old_buf;
+        return;
+    }
+    ptls->safe_restore = &buf;
+    low = (char*)((uintptr_t)low & ~(uintptr_t)15);
+    for (char **stack_p = ((char**)high) - 1; stack_p > (char**)low; stack_p--) {
         char *p = *stack_p;
         size_t osize;
         jl_taggedvalue_t *tag = jl_gc_find_taggedvalue_pool(p, &osize);
@@ -371,11 +385,42 @@ static void gc_scrub_range(char *stack_lo, char *stack_hi)
         // set mark to GC_MARKED (young and marked)
         tag->bits.gc = GC_MARKED;
     }
+    ptls->safe_restore = old_buf;
+}
+
+static void gc_scrub_task(jl_task_t *ta)
+{
+    int16_t tid = ta->tid;
+    jl_ptls_t ptls = jl_get_ptls_states();
+    jl_ptls_t ptls2 = jl_all_tls_states[tid];
+    if (ptls == ptls2 && ta == ptls2->current_task) {
+        // scan up to current `sp` for current thread and task
+        char *low = (char*)jl_get_frame_addr();
+#ifdef COPY_STACKS
+        gc_scrub_range(low, ptls2->stack_hi);
+#else
+        gc_scrub_range(low, (char*)ta->stkbuf + ta->ssize);
+#endif
+        return;
+    }
+    // The task that owns/is running on the threads's stack.
+#ifdef COPY_STACKS
+    jl_task_t *thread_task = ptls2->current_task;
+#else
+    jl_task_t *thread_task = ptls2->root_task;
+#endif
+    if (ta == thread_task)
+        gc_scrub_range(ptls2->stack_lo, ptls2->stack_hi);
+    if (ta->stkbuf == (void*)(intptr_t)(-1) || !ta->stkbuf)
+        return;
+    gc_scrub_range((char*)ta->stkbuf, (char*)ta->stkbuf + ta->ssize);
 }
 
-void gc_scrub(char *stack_hi)
+void gc_scrub(void)
 {
-    gc_scrub_range(gc_stack_lo, stack_hi);
+    for (size_t i = 0; i < jl_gc_debug_tasks.len; i++)
+        gc_scrub_task((jl_task_t*)jl_gc_debug_tasks.items[i]);
+    jl_gc_debug_tasks.len = 0;
 }
 #else
 void gc_debug_critical_error(void)
@@ -694,7 +739,6 @@ void gc_time_sweep_pause(uint64_t gc_end_t, int64_t actual_allocd,
 void gc_debug_init(void)
 {
 #ifdef GC_DEBUG_ENV
-    gc_stack_lo = (char*)gc_get_stack_ptr();
     char *env = getenv("JULIA_GC_NO_GENERATIONAL");
     if (env && strcmp(env, "0") != 0)
         jl_gc_debug_env.always_full = 1;
@@ -703,6 +747,7 @@ void gc_debug_init(void)
     gc_debug_alloc_init(&jl_gc_debug_env.pool, "POOL");
     gc_debug_alloc_init(&jl_gc_debug_env.other, "OTHER");
     gc_debug_alloc_init(&jl_gc_debug_env.print, "PRINT");
+    arraylist_new(&jl_gc_debug_tasks, 0);
 #endif
 
 #ifdef GC_VERIFY

src/gc.c

@@ -1255,6 +1255,7 @@ static void gc_mark_stack(jl_ptls_t ptls, jl_value_t *ta, jl_gcframe_t *s,
 
 static void gc_mark_task_stack(jl_ptls_t ptls, jl_task_t *ta, int d)
 {
+    gc_scrub_record_task(ta);
     int stkbuf = (ta->stkbuf != (void*)(intptr_t)-1 && ta->stkbuf != NULL);
     int16_t tid = ta->tid;
     jl_ptls_t ptls2 = jl_all_tls_states[tid];
@@ -1649,7 +1650,7 @@ void jl_gc_sync_total_bytes(void) {last_gc_total_bytes = jl_gc_total_bytes();}
 #define MIN_SCAN_BYTES 1024*1024
 
 // Only one thread should be running in this function
-static void _jl_gc_collect(jl_ptls_t ptls, int full, char *stack_hi)
+static void _jl_gc_collect(jl_ptls_t ptls, int full)
 {
     JL_TIMING(GC);
     uint64_t t0 = jl_hrtime();
@@ -1785,7 +1786,7 @@ static void _jl_gc_collect(jl_ptls_t ptls, int full, char *stack_hi)
     // 5. start sweeping
     sweep_weak_refs();
     gc_sweep_other(ptls, sweep_full);
-    gc_scrub(stack_hi);
+    gc_scrub();
     gc_sweep_pool(sweep_full);
     // sweeping is over
     // 6. if it is a quick sweep, put back the remembered objects in queued state
@@ -1822,7 +1823,7 @@ static void _jl_gc_collect(jl_ptls_t ptls, int full, char *stack_hi)
     gc_num.freed = 0;
 
     if (recollect) {
-        _jl_gc_collect(ptls, 0, stack_hi);
+        _jl_gc_collect(ptls, 0);
     }
 }
 
@@ -1834,7 +1835,6 @@ JL_DLLEXPORT void jl_gc_collect(int full)
         gc_num.allocd = -(int64_t)gc_num.interval;
         return;
     }
-    char *stack_hi = (char*)gc_get_stack_ptr();
     gc_debug_print();
 
     int8_t old_state = jl_gc_state(ptls);
@@ -1851,7 +1851,7 @@ JL_DLLEXPORT void jl_gc_collect(int full)
 
     if (!jl_gc_disable_counter) {
         JL_LOCK_NOGC(&finalizers_lock);
-        _jl_gc_collect(ptls, full, stack_hi);
+        _jl_gc_collect(ptls, full);
         JL_UNLOCK_NOGC(&finalizers_lock);
     }
 

src/gc.h

@@ -392,7 +392,8 @@ JL_DLLEXPORT extern jl_gc_debug_env_t jl_gc_debug_env;
 int gc_debug_check_other(void);
 int gc_debug_check_pool(void);
 void gc_debug_print(void);
-void gc_scrub(char *stack_hi);
+void gc_scrub_record_task(jl_task_t *ta);
+void gc_scrub(void);
 #else
 #define gc_sweep_always_full 0
 static inline int gc_debug_check_other(void)
@@ -406,9 +407,12 @@ static inline int gc_debug_check_pool(void)
 static inline void gc_debug_print(void)
 {
 }
-static inline void gc_scrub(char *stack_hi)
+static inline void gc_scrub_record_task(jl_task_t *ta)
+{
+    (void)ta;
+}
+static inline void gc_scrub(void)
 {
-    (void)stack_hi;
 }
 #endif
 
@@ -431,8 +435,8 @@ static inline void objprofile_reset(void)
 #endif
 
 #ifdef MEMPROFILE
-static void gc_stats_all_pool(void);
-static void gc_stats_big_obj(void);
+void gc_stats_all_pool(void);
+void gc_stats_big_obj(void);
 #else
 #define gc_stats_all_pool()
 #define gc_stats_big_obj()