Implement rb_gc_impl_mark_weak

Author: peterzhu2118

Diff for: gc/mmtk.c

@@ -450,7 +450,7 @@ void *
 rb_gc_impl_objspace_alloc(void)
 {
     MMTk_Builder *builder = mmtk_builder_default();
-    mmtk_init_binding(builder, NULL, &ruby_upcalls);
+    mmtk_init_binding(builder, NULL, &ruby_upcalls, (MMTk_ObjectReference)Qundef);
 
     return calloc(1, sizeof(struct objspace));
 }
@@ -738,21 +738,23 @@ rb_gc_impl_mark_and_pin(void *objspace_ptr, VALUE obj)
 }
 
 void
-rb_gc_impl_mark_maybe(void *objspace_ptr, VALUE obj) {
+rb_gc_impl_mark_maybe(void *objspace_ptr, VALUE obj)
+{
     if (rb_gc_impl_pointer_to_heap_p(objspace_ptr, (const void *)obj)) {
         rb_gc_impl_mark_and_pin(objspace_ptr, obj);
     }
 }
 
 void
-rb_gc_impl_mark_weak(void *objspace_ptr, VALUE *ptr) {
-    rb_bug("unimplemented");
+rb_gc_impl_mark_weak(void *objspace_ptr, VALUE *ptr)
+{
+    mmtk_mark_weak((const MMTk_ObjectReference *)ptr);
 }
 
 void
 rb_gc_impl_remove_weak(void *objspace_ptr, VALUE parent_obj, VALUE *ptr)
 {
-    rb_bug("unimplemented");
+    mmtk_remove_weak((const MMTk_ObjectReference *)ptr);
 }
 
 // Compaction

Diff for: gc/mmtk.h

@@ -101,7 +101,8 @@ MMTk_Builder *mmtk_builder_default(void);
 
 void mmtk_init_binding(MMTk_Builder *builder,
                        const struct MMTk_RubyBindingOptions *_binding_options,
-                       const struct MMTk_RubyUpcalls *upcalls);
+                       const struct MMTk_RubyUpcalls *upcalls,
+                       MMTk_ObjectReference weak_reference_dead_value);
 
 void mmtk_initialize_collection(MMTk_VMThread tls);
 
@@ -126,6 +127,10 @@ 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_remove_weak(const MMTk_ObjectReference *ptr);
+
 void mmtk_object_reference_write_post(MMTk_Mutator *mutator, MMTk_ObjectReference object);
 
 void mmtk_register_wb_unprotected_object(MMTk_ObjectReference object);

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

@@ -55,6 +55,7 @@ pub extern "C" fn mmtk_init_binding(
     builder: *mut MMTKBuilder,
     _binding_options: *const RubyBindingOptions,
     upcalls: *const RubyUpcalls,
+    weak_reference_dead_value: ObjectReference,
 ) {
     crate::set_panic_hook();
 
@@ -63,7 +64,7 @@ pub extern "C" fn mmtk_init_binding(
     let mmtk_boxed = mmtk_init(&builder);
     let mmtk_static = Box::leak(Box::new(mmtk_boxed));
 
-    let binding = RubyBinding::new(mmtk_static, &binding_options, upcalls);
+    let binding = RubyBinding::new(mmtk_static, &binding_options, upcalls, weak_reference_dead_value);
 
     crate::BINDING
         .set(binding)
@@ -133,6 +134,18 @@ pub extern "C" fn mmtk_add_obj_free_candidate(object: ObjectReference) {
     binding().weak_proc.add_obj_free_candidate(object)
 }
 
+// =============== Marking ===============
+
+#[no_mangle]
+pub extern "C" fn mmtk_mark_weak(ptr: &'static mut ObjectReference) {
+    binding().weak_proc.add_weak_reference(ptr);
+}
+
+#[no_mangle]
+pub extern "C" fn mmtk_remove_weak(ptr: &ObjectReference) {
+    binding().weak_proc.remove_weak_reference(ptr);
+}
+
 // =============== Write barriers ===============
 
 #[no_mangle]

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

@@ -1,6 +1,5 @@
 use std::collections::{HashMap, HashSet};
 use std::ffi::CString;
-use std::str::FromStr;
 use std::sync::atomic::AtomicBool;
 use std::sync::Mutex;
 use std::thread::JoinHandle;
@@ -66,39 +65,24 @@ pub struct RubyBinding {
     pub(crate) moved_givtbl: Mutex<HashMap<ObjectReference, MovedGIVTblEntry>>,
     pub gc_thread_join_handles: Mutex<Vec<JoinHandle<()>>>,
     pub wb_unprotected_objects: Mutex<HashSet<ObjectReference>>,
-    pub st_entries_chunk_size: usize,
-    pub st_bins_chunk_size: usize,
+
+    pub weak_reference_dead_value: ObjectReference,
 }
 
 unsafe impl Sync for RubyBinding {}
 unsafe impl Send for RubyBinding {}
 
-fn env_default<T>(name: &str, default: T) -> T
-where
-    T: FromStr,
-{
-    std::env::var(name)
-        .ok()
-        .and_then(|x| x.parse::<T>().ok())
-        .unwrap_or(default)
-}
-
 impl RubyBinding {
     pub fn new(
         mmtk: &'static MMTK<Ruby>,
         binding_options: &RubyBindingOptions,
         upcalls: *const abi::RubyUpcalls,
+        weak_reference_dead_value: ObjectReference,
     ) -> Self {
         unsafe {
             crate::BINDING_FAST.suffix_size = binding_options.suffix_size;
         }
 
-        let st_entries_chunk_size = env_default::<usize>("RUBY_MMTK_ENTRIES_CHUNK_SIZE", 1024);
-        let st_bins_chunk_size = env_default::<usize>("RUBY_MMTK_BINS_CHUNK_SIZE", 4096);
-
-        debug!("st_entries_chunk_size: {st_entries_chunk_size}");
-        debug!("st_bins_chunk_size: {st_bins_chunk_size}");
-
         Self {
             mmtk,
             options: binding_options.clone(),
@@ -109,8 +93,8 @@ impl RubyBinding {
             moved_givtbl: Default::default(),
             gc_thread_join_handles: Default::default(),
             wb_unprotected_objects: Default::default(),
-            st_entries_chunk_size,
-            st_bins_chunk_size,
+
+            weak_reference_dead_value
         }
     }
 

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

@@ -18,6 +18,7 @@ pub struct WeakProcessor {
     /// If it is a bottleneck, replace it with a lock-free data structure,
     /// or add candidates in batch.
     obj_free_candidates: Mutex<Vec<ObjectReference>>,
+    weak_references: Mutex<Vec<&'static mut ObjectReference>>,
 }
 
 impl Default for WeakProcessor {
@@ -30,6 +31,7 @@ impl WeakProcessor {
     pub fn new() -> Self {
         Self {
             obj_free_candidates: Mutex::new(Vec::new()),
+            weak_references: Mutex::new(Vec::new()),
         }
     }
 
@@ -56,12 +58,28 @@ impl WeakProcessor {
         std::mem::take(obj_free_candidates.as_mut())
     }
 
+    pub fn add_weak_reference(&self, ptr: &'static mut ObjectReference) {
+        let mut weak_references = self.weak_references.lock().unwrap();
+        weak_references.push(ptr);
+    }
+
+    pub fn remove_weak_reference(&self, ptr: &ObjectReference) {
+        let mut weak_references = self.weak_references.lock().unwrap();
+        for (i, curr_ptr) in weak_references.iter().enumerate() {
+            if *curr_ptr == ptr {
+                weak_references.swap_remove(i);
+                break;
+            }
+        }
+    }
+
     pub fn process_weak_stuff(
         &self,
         worker: &mut GCWorker<Ruby>,
         _tracer_context: impl ObjectTracerContext<Ruby>,
     ) {
         worker.add_work(WorkBucketStage::VMRefClosure, ProcessObjFreeCandidates);
+        worker.add_work(WorkBucketStage::VMRefClosure, ProcessWeakReferences);
 
         worker.scheduler().work_buckets[WorkBucketStage::VMRefClosure].bulk_add(vec![
             Box::new(UpdateGenericIvTbl) as _,
@@ -141,6 +159,26 @@ impl GCWork<Ruby> for ProcessObjFreeCandidates {
     }
 }
 
+struct ProcessWeakReferences;
+
+impl GCWork<Ruby> for ProcessWeakReferences {
+    fn do_work(&mut self, _worker: &mut GCWorker<Ruby>, _mmtk: &'static mmtk::MMTK<Ruby>) {
+        let mut weak_references = crate::binding()
+            .weak_proc
+            .weak_references
+            .try_lock()
+            .expect("Mutators should not be holding the lock.");
+
+            for ptr_ptr in weak_references.iter_mut() {
+                if !(**ptr_ptr).is_reachable() {
+                    **ptr_ptr = crate::binding().weak_reference_dead_value;
+                }
+            }
+
+            weak_references.clear();
+    }
+}
+
 trait GlobalTableProcessingWork {
     fn process_table(&mut self);