Join GC threads

Author: wks

Diff for: mmtk/src/api.rs

@@ -157,6 +157,7 @@ pub extern "C" fn mmtk_initialize_collection(tls: VMThread) {
 #[no_mangle]
 pub extern "C" fn mmtk_prepare_to_fork() {
     mmtk().prepare_to_fork();
+    binding().join_all_gc_threads();
 }
 
 #[no_mangle]

Diff for: mmtk/src/binding.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::ffi::CString;
 use std::sync::Mutex;
+use std::thread::JoinHandle;
 
 use libc::c_void;
 use mmtk::util::ObjectReference;
@@ -36,6 +37,7 @@ pub struct RubyBinding {
     pub weak_proc: WeakProcessor,
     pub ppp_registry: PPPRegistry,
     pub(crate) moved_givtbl: Mutex<HashMap<ObjectReference, MovedGIVTblEntry>>,
+    pub gc_thread_join_handles: Mutex<Vec<JoinHandle<()>>>
 }
 
 unsafe impl Sync for RubyBinding {}
@@ -58,6 +60,7 @@ impl RubyBinding {
             weak_proc: WeakProcessor::new(),
             ppp_registry: PPPRegistry::new(),
             moved_givtbl: Default::default(),
+            gc_thread_join_handles: Default::default(),
         }
     }
 
@@ -75,4 +78,20 @@ impl RubyBinding {
         }
         plan_name.as_deref().unwrap().as_ptr()
     }
+
+    pub fn join_all_gc_threads(&self) {
+        let handles = {
+            let mut guard = self.gc_thread_join_handles.lock().unwrap();
+            std::mem::take(&mut *guard)
+        };
+
+        debug!("Joining GC threads...");
+        let total = handles.len();
+        let mut joined = 0;
+        for handle in handles {
+            handle.join().unwrap();
+            joined += 1;
+            debug!("{joined}/{total} GC threads joined.");
+        }
+    }
 }

Diff for: mmtk/src/collection.rs

@@ -32,34 +32,37 @@ impl Collection<Ruby> for VMCollection {
     }
 
     fn spawn_gc_thread(_tls: VMThread, ctx: GCThreadContext<Ruby>) {
-        match ctx {
-            GCThreadContext::Worker(mut worker) => {
-                thread::Builder::new()
-                    .name("MMTk Worker Thread".to_string())
-                    .spawn(move || {
-                        let ordinal = worker.ordinal;
-                        debug!(
-                            "Hello! This is MMTk Worker Thread running! ordinal: {}",
-                            ordinal
-                        );
-                        crate::register_gc_thread(thread::current().id());
-                        let ptr_worker = &mut *worker as *mut GCWorker<Ruby>;
-                        let gc_thread_tls =
-                            Box::into_raw(Box::new(GCThreadTLS::for_worker(ptr_worker)));
-                        (upcalls().init_gc_worker_thread)(gc_thread_tls);
-                        memory_manager::start_worker(
-                            mmtk(),
-                            GCThreadTLS::to_vwt(gc_thread_tls),
-                            worker,
-                        );
-                        debug!(
-                            "An MMTk Worker Thread is quitting. Good bye! ordinal: {}",
-                            ordinal
-                        );
-                        crate::unregister_gc_thread(thread::current().id());
-                    })
-                    .unwrap();
-            }
+        let join_handle = match ctx {
+            GCThreadContext::Worker(mut worker) => thread::Builder::new()
+                .name("MMTk Worker Thread".to_string())
+                .spawn(move || {
+                    let ordinal = worker.ordinal;
+                    debug!(
+                        "Hello! This is MMTk Worker Thread running! ordinal: {}",
+                        ordinal
+                    );
+                    crate::register_gc_thread(thread::current().id());
+                    let ptr_worker = &mut *worker as *mut GCWorker<Ruby>;
+                    let gc_thread_tls =
+                        Box::into_raw(Box::new(GCThreadTLS::for_worker(ptr_worker)));
+                    (upcalls().init_gc_worker_thread)(gc_thread_tls);
+                    memory_manager::start_worker(
+                        mmtk(),
+                        GCThreadTLS::to_vwt(gc_thread_tls),
+                        worker,
+                    );
+                    debug!(
+                        "An MMTk Worker Thread is quitting. Good bye! ordinal: {}",
+                        ordinal
+                    );
+                    crate::unregister_gc_thread(thread::current().id());
+                })
+                .unwrap(),
+        };
+
+        {
+            let mut handles = crate::binding().gc_thread_join_handles.lock().unwrap();
+            handles.push(join_handle);
         }
     }