Fix stack scanning

Author: peterzhu2118

Diff for: gc.c

@@ -196,6 +196,12 @@ rb_gc_get_ractor_newobj_cache(void)
     return GET_RACTOR()->newobj_cache;
 }
 
+void *
+rb_gc_get_current_execution_context(void)
+{
+    return GET_EC();
+}
+
 void
 rb_gc_ractor_newobj_cache_foreach(void (*func)(void *cache, void *data), void *data)
 {
@@ -2273,6 +2279,12 @@ rb_mark_locations(void *begin, void *end)
 
 # if defined(__EMSCRIPTEN__)
 
+void
+rb_gc_save_machine_context(void)
+{
+    // no-op
+}
+
 static void
 mark_current_machine_context(void *objspace, rb_execution_context_t *ec)
 {
@@ -2300,35 +2312,18 @@ mark_current_machine_context(void *objspace, rb_execution_context_t *ec)
 
 #else // !defined(__wasm__)
 
-static void
-mark_current_machine_context(void *objspace, rb_execution_context_t *ec)
+void
+rb_gc_save_machine_context(void)
 {
-    union {
-        rb_jmp_buf j;
-        VALUE v[sizeof(rb_jmp_buf) / (sizeof(VALUE))];
-    } save_regs_gc_mark;
-    VALUE *stack_start, *stack_end;
+    rb_thread_t *thread = GET_THREAD();
 
-    FLUSH_REGISTER_WINDOWS;
-    memset(&save_regs_gc_mark, 0, sizeof(save_regs_gc_mark));
-    /* This assumes that all registers are saved into the jmp_buf (and stack) */
-    rb_setjmp(save_regs_gc_mark.j);
-
-    /* SET_STACK_END must be called in this function because
-     * the stack frame of this function may contain
-     * callee save registers and they should be marked. */
-    SET_STACK_END;
-    GET_STACK_BOUNDS(stack_start, stack_end, 1);
-
-    struct mark_machine_stack_location_maybe_data data = {
-        .objspace = objspace,
-#ifdef RUBY_ASAN_ENABLED
-        .ec = ec
-#endif
-    };
+    RB_VM_SAVE_MACHINE_CONTEXT(thread);
+}
 
-    each_location((void *)&data, save_regs_gc_mark.v, numberof(save_regs_gc_mark.v), gc_mark_machine_stack_location_maybe);
-    each_stack_location((void *)&data, ec, stack_start, stack_end, gc_mark_machine_stack_location_maybe);
+static void
+mark_current_machine_context(void *objspace, const rb_execution_context_t *ec)
+{
+    rb_gc_mark_machine_context(ec);
 }
 #endif
 
@@ -2460,7 +2455,7 @@ mark_const_table_i(VALUE value, void *objspace)
 }
 
 void
-rb_gc_mark_roots(void *objspace, const char **categoryp)
+rb_gc_mark_roots(void *objspace, const void *ec, const char **categoryp)
 {
     rb_vm_t *vm = GET_VM();
 
@@ -2487,18 +2482,11 @@ rb_gc_mark_roots(void *objspace, const char **categoryp)
     }
 #endif
 
-    MARK_CHECKPOINT("finish");
-}
-
-void
-rb_gc_mark_thread_roots(void *objspace, void *ractor, const char **categoryp)
-{
-    if (ractor == NULL) ractor = GET_RACTOR();
-
-    rb_execution_context_t *ec = ((rb_ractor_t *)ractor)->threads.running_ec;
-
     MARK_CHECKPOINT("machine_context");
     mark_current_machine_context(objspace, ec);
+
+    MARK_CHECKPOINT("finish");
+
 #undef MARK_CHECKPOINT
 }
 
@@ -3617,7 +3605,8 @@ rb_objspace_reachable_objects_from_root(void (func)(const char *category, VALUE,
     }, *prev_mfd = cr->mfd;
 
     cr->mfd = &mfd;
-    rb_gc_mark_roots(rb_gc_get_objspace(), &data.category);
+    rb_gc_save_machine_context();
+    rb_gc_mark_roots(rb_gc_get_objspace(), rb_gc_get_current_execution_context(), &data.category);
     cr->mfd = prev_mfd;
 }
 

Diff for: gc/default.c

@@ -4838,9 +4838,8 @@ mark_roots(rb_objspace_t *objspace, const char **categoryp)
 
     if (stress_to_class) rb_gc_mark(stress_to_class);
 
-    rb_gc_mark_roots(objspace, categoryp);
-
-    rb_gc_mark_thread_roots(objspace, NULL, categoryp);
+    rb_gc_save_machine_context();
+    rb_gc_mark_roots(objspace, rb_gc_get_current_execution_context(), categoryp);
 }
 
 static inline void

Diff for: gc/gc.h

@@ -27,13 +27,14 @@ void rb_gc_reachable_objects_from_callback(VALUE obj);
 void rb_gc_event_hook(VALUE obj, rb_event_flag_t event);
 void *rb_gc_get_objspace(void);
 void *rb_gc_get_ractor_newobj_cache(void);
+void *rb_gc_get_current_execution_context(void);
 size_t rb_size_mul_or_raise(size_t x, size_t y, VALUE exc);
 void rb_gc_run_obj_finalizer(VALUE objid, long count, VALUE (*callback)(long i, void *data), void *data);
 void rb_gc_set_pending_interrupt(void);
 void rb_gc_unset_pending_interrupt(void);
 bool rb_gc_obj_free(void *objspace, VALUE obj);
-void rb_gc_mark_roots(void *objspace, const char **categoryp);
-void rb_gc_mark_thread_roots(void *objspace, void *ractor, const char **categoryp);
+void rb_gc_save_machine_context(void);
+void rb_gc_mark_roots(void *objspace, const void *ec, const char **categoryp);
 void rb_gc_ractor_newobj_cache_foreach(void (*func)(void *cache, void *data), void *data);
 bool rb_gc_multi_ractor_p(void);
 void rb_objspace_reachable_objects_from_root(void (func)(const char *category, VALUE, void *), void *passing_data);

Diff for: gc/mmtk.c

@@ -11,24 +11,6 @@
 #include "gc/gc_impl.h"
 #include "gc/mmtk.h"
 
-struct MMTk_final_job {
-    struct MMTk_final_job *next;
-    enum {
-        MMTK_FINAL_JOB_DFREE,
-        MMTK_FINAL_JOB_FINALIZE,
-    } kind;
-    union {
-        struct {
-            void (*dfree)(void *);
-            void *data;
-        } dfree;
-        struct {
-            VALUE object_id;
-            VALUE finalizer_array;
-        } finalize;
-    } as;
-};
-
 struct objspace {
     st_table *id_to_obj_tbl;
     st_table *obj_to_id_tbl;
@@ -38,6 +20,7 @@ struct objspace {
     struct MMTk_final_job *finalizer_jobs;
     rb_postponed_job_handle_t finalizer_postponed_job;
 
+    struct MMTk_ractor_cache *ractor_caches;
     unsigned long live_ractor_cache_count;
 
     int lock_lev;
@@ -49,6 +32,31 @@ struct objspace {
     size_t start_the_world_count;
 };
 
+struct MMTk_ractor_cache {
+    struct MMTk_ractor_cache *next_ractor_cache;
+
+    MMTk_Mutator *mutator;
+    void *execution_context;
+};
+
+struct MMTk_final_job {
+    struct MMTk_final_job *next;
+    enum {
+        MMTK_FINAL_JOB_DFREE,
+        MMTK_FINAL_JOB_FINALIZE,
+    } kind;
+    union {
+        struct {
+            void (*dfree)(void *);
+            void *data;
+        } dfree;
+        struct {
+            VALUE object_id;
+            VALUE finalizer_array;
+        } finalize;
+    } as;
+};
+
 #ifdef RB_THREAD_LOCAL_SPECIFIER
 RB_THREAD_LOCAL_SPECIFIER struct MMTk_GCThreadTLS *rb_mmtk_gc_thread_tls;
 #else
@@ -113,7 +121,7 @@ rb_mmtk_resume_mutators(void)
 }
 
 static void
-rb_mmtk_block_for_gc(MMTk_VMMutatorThread tls)
+rb_mmtk_block_for_gc(MMTk_VMMutatorThread mutator)
 {
     struct objspace *objspace = rb_gc_get_objspace();
 
@@ -125,6 +133,8 @@ rb_mmtk_block_for_gc(MMTk_VMMutatorThread tls)
     // Increment the stopped ractor count
     objspace->stopped_ractors++;
     if (objspace->stopped_ractors == 1) {
+        rb_gc_save_machine_context();
+        mutator->execution_context = rb_gc_get_current_execution_context();
         pthread_cond_broadcast(&objspace->cond_world_stopped);
     }
 
@@ -146,23 +156,28 @@ rb_mmtk_block_for_gc(MMTk_VMMutatorThread tls)
 static size_t
 rb_mmtk_number_of_mutators(void)
 {
-    // TODO
-    return 1;
+    struct objspace *objspace = rb_gc_get_objspace();
+    return objspace->live_ractor_cache_count;
 }
 
-static void *mutator;
-
 static void
 rb_mmtk_get_mutators(void (*visit_mutator)(MMTk_Mutator *mutator, void *data), void *data)
 {
-    // TODO
-    visit_mutator(mutator, data);
+    struct objspace *objspace = rb_gc_get_objspace();
+
+    struct MMTk_ractor_cache *ractor_cache = objspace->ractor_caches;
+    RUBY_ASSERT(ractor_cache != NULL);
+
+    while (ractor_cache != NULL) {
+        visit_mutator(ractor_cache->mutator, data);
+        ractor_cache = ractor_cache->next_ractor_cache;
+    }
 }
 
 static void
 rb_mmtk_scan_gc_roots(void)
 {
-    rb_gc_mark_roots(rb_gc_get_objspace(), NULL);
+    // rb_gc_mark_roots(rb_gc_get_objspace(), NULL);
 }
 
 static int
@@ -188,9 +203,7 @@ rb_mmtk_scan_objspace(void)
 static void
 rb_mmtk_scan_roots_in_mutator_thread(MMTk_VMMutatorThread mutator, MMTk_VMWorkerThread worker)
 {
-    void *ractor = mutator;
-
-    rb_gc_mark_thread_roots(rb_gc_get_objspace(), ractor, NULL);
+    rb_gc_mark_roots(rb_gc_get_objspace(), mutator->execution_context, NULL);
 }
 
 static void
@@ -457,10 +470,13 @@ rb_gc_impl_ractor_cache_alloc(void *objspace_ptr, void *ractor)
     }
     objspace->live_ractor_cache_count++;
 
-    mutator = mmtk_bind_mutator(ractor);
-    return mutator;
+    struct MMTk_ractor_cache *cache = malloc(sizeof(struct MMTk_ractor_cache));
+    cache->next_ractor_cache = objspace->ractor_caches;
+    objspace->ractor_caches = cache;
 
-    // return mmtk_bind_mutator(ractor);
+    cache->mutator = mmtk_bind_mutator(cache);
+
+    return cache;
 }
 
 void
@@ -572,7 +588,9 @@ rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags
         }
     }
 
-    VALUE *alloc_obj = mmtk_alloc(cache_ptr, alloc_size + 8, MMTk_MIN_OBJ_ALIGN, 0, MMTK_ALLOCATION_SEMANTICS_DEFAULT);
+    struct MMTk_ractor_cache *ractor_cache = cache_ptr;
+
+    VALUE *alloc_obj = mmtk_alloc(ractor_cache->mutator, alloc_size + 8, MMTk_MIN_OBJ_ALIGN, 0, MMTK_ALLOCATION_SEMANTICS_DEFAULT);
     alloc_obj++;
     alloc_obj[-1] = alloc_size;
     alloc_obj[0] = flags;
@@ -581,7 +599,7 @@ rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags
     if (alloc_size > 24) alloc_obj[3] = v2;
     if (alloc_size > 32) alloc_obj[4] = v3;
 
-    mmtk_post_alloc(cache_ptr, (void*)alloc_obj, alloc_size + 8, MMTK_ALLOCATION_SEMANTICS_DEFAULT);
+    mmtk_post_alloc(ractor_cache->mutator, (void*)alloc_obj, alloc_size + 8, MMTK_ALLOCATION_SEMANTICS_DEFAULT);
 
     // TODO: only add when object needs obj_free to be called
     mmtk_add_obj_free_candidate(alloc_obj);
@@ -708,7 +726,9 @@ rb_gc_impl_location(void *objspace_ptr, VALUE value)
 void
 rb_gc_impl_writebarrier(void *objspace_ptr, VALUE a, VALUE b)
 {
-    mmtk_object_reference_write_post(rb_gc_get_ractor_newobj_cache(), (MMTk_ObjectReference)a);
+    struct MMTk_ractor_cache *cache = rb_gc_get_ractor_newobj_cache();
+
+    mmtk_object_reference_write_post(cache->mutator, (MMTk_ObjectReference)a);
 }
 
 void
@@ -720,7 +740,7 @@ rb_gc_impl_writebarrier_unprotect(void *objspace_ptr, VALUE obj)
 void
 rb_gc_impl_writebarrier_remember(void *objspace_ptr, VALUE obj)
 {
-    mmtk_object_reference_write_post(rb_gc_get_ractor_newobj_cache(), (MMTk_ObjectReference)obj);
+    rb_gc_impl_writebarrier(objspace_ptr, obj, Qundef);
 }
 
 // Heap walking

Diff for: gc/mmtk.h

@@ -12,10 +12,8 @@
 typedef struct MMTk_Builder MMTk_Builder;
 typedef struct MMTk_Mutator MMTk_Mutator;
 
-struct rb_thread_struct;
-typedef struct rb_thread_struct rb_thread_t;
-typedef rb_thread_t *MMTk_VMThread;
-typedef rb_thread_t *MMTk_VMMutatorThread;
+typedef struct MMTk_ractor_cache *MMTk_VMThread;
+typedef struct MMTk_ractor_cache *MMTk_VMMutatorThread;
 typedef struct MMTk_GCThreadTLS *MMTk_VMWorkerThread;
 typedef void *MMTk_Address;
 typedef void *MMTk_ObjectReference;

Diff for: gc/mmtk/cbindgen.toml

@@ -13,10 +13,8 @@ after_includes = """
 typedef struct MMTk_Builder MMTk_Builder;
 typedef struct MMTk_Mutator MMTk_Mutator;
 
-struct rb_thread_struct;
-typedef struct rb_thread_struct rb_thread_t;
-typedef rb_thread_t *MMTk_VMThread;
-typedef rb_thread_t *MMTk_VMMutatorThread;
+typedef struct MMTk_ractor_cache *MMTk_VMThread;
+typedef struct MMTk_ractor_cache *MMTk_VMMutatorThread;
 typedef struct MMTk_GCThreadTLS *MMTk_VMWorkerThread;
 typedef void *MMTk_Address;
 typedef void *MMTk_ObjectReference;