Add forking support

Author: peterzhu2118

Diff for: gc.c

@@ -680,6 +680,9 @@ typedef struct gc_function_map {
     // Object ID
     VALUE (*object_id)(void *objspace_ptr, VALUE obj);
     VALUE (*object_id_to_ref)(void *objspace_ptr, VALUE object_id);
+    // Forking
+    void (*before_fork)(void *objspace_ptr);
+    void (*after_fork)(void *objspace_ptr, rb_pid_t pid);
     // Statistics
     VALUE (*set_measure_total_time)(void *objspace_ptr, VALUE flag);
     VALUE (*get_measure_total_time)(void *objspace_ptr);
@@ -810,6 +813,9 @@ ruby_external_gc_init(void)
     // Object ID
     load_external_gc_func(object_id);
     load_external_gc_func(object_id_to_ref);
+    // Forking
+    load_external_gc_func(before_fork);
+    load_external_gc_func(after_fork);
     // Statistics
     load_external_gc_func(set_measure_total_time);
     load_external_gc_func(get_measure_total_time);
@@ -888,6 +894,9 @@ ruby_external_gc_init(void)
 // Object ID
 # define rb_gc_impl_object_id rb_gc_functions.object_id
 # define rb_gc_impl_object_id_to_ref rb_gc_functions.object_id_to_ref
+// Forking
+# define rb_gc_impl_before_fork rb_gc_functions.before_fork
+# define rb_gc_impl_after_fork rb_gc_functions.after_fork
 // Statistics
 # define rb_gc_impl_set_measure_total_time rb_gc_functions.set_measure_total_time
 # define rb_gc_impl_get_measure_total_time rb_gc_functions.get_measure_total_time
@@ -4374,6 +4383,18 @@ rb_obj_info_dump_loc(VALUE obj, const char *file, int line, const char *func)
     fprintf(stderr, "<OBJ_INFO:%s@%s:%d> %s\n", func, file, line, rb_raw_obj_info(buff, 0x100, obj));
 }
 
+void
+rb_gc_before_fork(void)
+{
+    rb_gc_impl_before_fork(rb_gc_get_objspace());
+}
+
+void
+rb_gc_after_fork(rb_pid_t pid)
+{
+    rb_gc_impl_after_fork(rb_gc_get_objspace(), pid);
+}
+
 /*
  * Document-module: ObjectSpace
  *

Diff for: gc/default.c

@@ -9347,6 +9347,9 @@ gc_malloc_allocations(VALUE self)
 }
 #endif
 
+void rb_gc_impl_before_fork(void *objspace_ptr) { /* no-op */ }
+void rb_gc_impl_after_fork(void *objspace_ptr, rb_pid_t pid) { /* no-op */ }
+
 void *
 rb_gc_impl_objspace_alloc(void)
 {

Diff for: gc/gc_impl.h

@@ -84,6 +84,9 @@ GC_IMPL_FN void rb_gc_impl_shutdown_call_finalizer(void *objspace_ptr);
 // Object ID
 GC_IMPL_FN VALUE rb_gc_impl_object_id(void *objspace_ptr, VALUE obj);
 GC_IMPL_FN VALUE rb_gc_impl_object_id_to_ref(void *objspace_ptr, VALUE object_id);
+// Forking
+GC_IMPL_FN void rb_gc_impl_before_fork(void *objspace_ptr);
+GC_IMPL_FN void rb_gc_impl_after_fork(void *objspace_ptr, rb_pid_t pid);
 // Statistics
 GC_IMPL_FN VALUE rb_gc_impl_set_measure_total_time(void *objspace_ptr, VALUE flag);
 GC_IMPL_FN VALUE rb_gc_impl_get_measure_total_time(void *objspace_ptr);

Diff for: gc/mmtk.c

@@ -748,13 +748,13 @@ rb_gc_impl_mark_maybe(void *objspace_ptr, VALUE obj)
 void
 rb_gc_impl_mark_weak(void *objspace_ptr, VALUE *ptr)
 {
-    mmtk_mark_weak((const MMTk_ObjectReference *)ptr);
+    mmtk_mark_weak((MMTk_ObjectReference *)ptr);
 }
 
 void
 rb_gc_impl_remove_weak(void *objspace_ptr, VALUE parent_obj, VALUE *ptr)
 {
-    mmtk_remove_weak((const MMTk_ObjectReference *)ptr);
+    mmtk_remove_weak((MMTk_ObjectReference *)ptr);
 }
 
 // Compaction
@@ -1080,6 +1080,20 @@ rb_gc_impl_object_id_to_ref(void *objspace_ptr, VALUE object_id)
     }
 }
 
+// Forking
+
+void
+rb_gc_impl_before_fork(void *objspace_ptr)
+{
+    mmtk_before_fork();
+}
+
+void
+rb_gc_impl_after_fork(void *objspace_ptr, rb_pid_t pid)
+{
+    mmtk_after_fork(rb_gc_get_ractor_newobj_cache());
+}
+
 // Statistics
 VALUE rb_gc_impl_set_measure_total_time(void *objspace_ptr, VALUE flag) { }
 VALUE rb_gc_impl_get_measure_total_time(void *objspace_ptr) { }

Diff for: gc/mmtk.h

@@ -127,7 +127,7 @@ void mmtk_post_alloc(MMTk_Mutator *mutator,
 
 void mmtk_add_obj_free_candidate(MMTk_ObjectReference object);
 
-void mmtk_mark_weak(const MMTk_ObjectReference *ptr);
+void mmtk_mark_weak(MMTk_ObjectReference *ptr);
 
 void mmtk_remove_weak(const MMTk_ObjectReference *ptr);
 
@@ -141,6 +141,10 @@ struct MMTk_RawVecOfObjRef mmtk_get_all_obj_free_candidates(void);
 
 void mmtk_free_raw_vec_of_obj_ref(struct MMTk_RawVecOfObjRef raw_vec);
 
+void mmtk_before_fork(void);
+
+void mmtk_after_fork(MMTk_VMThread tls);
+
 size_t mmtk_total_bytes(void);
 
 size_t mmtk_used_bytes(void);

Diff for: gc/mmtk/src/api.rs

@@ -193,6 +193,19 @@ pub extern "C" fn mmtk_free_raw_vec_of_obj_ref(raw_vec: RawVecOfObjRef) {
     unsafe { raw_vec.into_vec() };
 }
 
+// =============== Forking ===============
+
+#[no_mangle]
+pub extern "C" fn mmtk_before_fork() {
+    mmtk().prepare_to_fork();
+    binding().join_all_gc_threads();
+}
+
+#[no_mangle]
+pub extern "C" fn mmtk_after_fork(tls: VMThread) {
+    mmtk().after_fork(tls);
+}
+
 // =============== Statistics ===============
 
 #[no_mangle]

Diff for: internal/gc.h

@@ -213,6 +213,9 @@ void rb_gc_ref_update_table_values_only(st_table *tbl);
 
 void rb_gc_initial_stress_set(VALUE flag);
 
+void rb_gc_before_fork(void);
+void rb_gc_after_fork(rb_pid_t pid);
+
 #define rb_gc_mark_and_move_ptr(ptr) do { \
     VALUE _obj = (VALUE)*(ptr); \
     rb_gc_mark_and_move(&_obj); \

Diff for: process.c

@@ -1676,12 +1676,15 @@ after_exec(void)
 static void
 before_fork_ruby(void)
 {
+    rb_gc_before_fork();
     before_exec();
 }
 
 static void
 after_fork_ruby(rb_pid_t pid)
 {
+    rb_gc_after_fork(pid);
+
     if (pid == 0) {
         // child
         clear_pid_cache();