[ruby/mmtk] Keep Ruby's GC disable state across a forced GC.start

GC.start has to run even when GC is disabled, so we enable collection,
run the GC and then disable it again. #134 fixed the crash we saw when
another Ractor requested a GC before we could re-disable, but we still
restore whatever state we saw before the GC.

That's a problem when another Ractor changes the state in the meantime.
If it enables GC while we're collecting, we disable it again on its
behalf and GC stays off for good. Ractor.new disables and enables GC
around thread allocation, so this isn't hard to hit. GC.disable and the
internal gc_disable also share a single flag in mmtk-core, so one can
undo the other.

In this commit we track the internal flag, the user flag and any forced
GC.start requests in the binding, and mirror the combined state into
mmtk-core. A forced GC.start only enables collection until the next
pause ends, and we reapply the state in resume_mutators before any
mutator can request another GC.

GC workers also toggle the internal flag during a pause, through
DURING_GC_COULD_MALLOC_REGION. That reads the flag and then sets it, so
two workers running in parallel can leave it set. mmtk-core used to
ignore these because it can't disable collection mid-pause, so we do the
same and drop internal changes between stop_all_mutators and
resume_mutators.

This replaces the retry loop from #134, if we fail to disable because a
GC is pending, the end of that GC puts the state right.

AtomicU64::update needs Rust 1.95, so we declare that in Cargo.toml.

https://github.com/ruby/mmtk/commit/7430ebeaa6

gc/mmtk/Cargo.toml

@@ -3,6 +3,7 @@ name = "mmtk_ruby"
 version = "0.1.0"
 authors = []
 edition = "2024"
+rust-version = "1.95"
 
 [lib]
 name = "mmtk_ruby"

gc/mmtk/mmtk.c

@@ -22,7 +22,6 @@
 #endif
 
 struct objspace {
-    bool user_gc_disabled;
     bool measure_gc_time;
     bool gc_stress;
 
@@ -839,23 +838,13 @@ rb_gc_impl_gc_disable(void *objspace_ptr, bool finish_current_gc)
 bool
 rb_gc_impl_user_gc_disabled_set(void *objspace_ptr, bool disable)
 {
-    struct objspace *objspace = objspace_ptr;
-    const bool was = objspace->user_gc_disabled;
-    objspace->user_gc_disabled = disable;
-
-    if (was != disable) {
-        mmtk_set_gc_enabled(!disable);
-        objspace->user_gc_disabled = disable;
-    }
-
-    return was;
+    return mmtk_set_user_gc_disabled(disable);
 }
 
 bool
 rb_gc_impl_user_gc_disabled_p(void *objspace_ptr)
 {
-    struct objspace *objspace = objspace_ptr;
-    return objspace->user_gc_disabled;
+    return mmtk_user_gc_disabled_p();
 }
 
 bool

gc/mmtk/mmtk.h

@@ -88,6 +88,10 @@ void mmtk_set_gc_enabled(bool enable);
 
 bool mmtk_gc_enabled_p(void);
 
+bool mmtk_set_user_gc_disabled(bool disable);
+
+bool mmtk_user_gc_disabled_p(void);
+
 size_t mmtk_max_non_los_default_alloc_bytes(void);
 
 bool mmtk_plan_has_write_barrier(void);

gc/mmtk/src/api.rs

@@ -250,7 +250,8 @@ pub extern "C" fn mmtk_get_vo_bit_base_addr() -> usize {
 
 #[unsafe(no_mangle)]
 pub extern "C" fn mmtk_initialize_collection(tls: VMThread) {
-    memory_manager::initialize_collection(mmtk(), tls)
+    memory_manager::initialize_collection(mmtk(), tls);
+    crate::binding().gc_disable.on_collection_initialized();
 }
 
 #[unsafe(no_mangle)]
@@ -292,51 +293,29 @@ pub extern "C" fn mmtk_handle_user_collection_request(
     force: bool,
     exhaustive: bool,
 ) {
-    // GC.start must run even when GC is disabled. However, GC.stress=true does
-    // not run when GC is disabled. force is true when GC.start and false when
-    // GC.stress=true.
-    // TODO: This is not Ractor-safe. Another Ractor could disable GC, which
-    //       would prevent the GC from running. Another Ractor could also enable
-    //       GC, which would be lost.
-    let gc_was_disabled = force && !crate::mmtk().is_collection_enabled();
-    if gc_was_disabled {
-        crate::mmtk().enable_collection();
-    }
-
-    crate::mmtk().handle_user_collection_request(tls, force, exhaustive);
-
-    if gc_was_disabled {
-        // While GC was enabled, another Ractor may have requested a GC that hasn't run yet.
-        // MMTk can't disable GC until that GC runs, and it needs this thread to stop first.
-        // So block for it and try again. Since a GC is already requested, this waits for that
-        // GC instead of starting a new one.
-        while crate::mmtk().disable_collection().is_err() {
-            crate::mmtk().handle_user_collection_request(tls, true, false);
-        }
-    }
+    crate::binding()
+        .gc_disable
+        .handle_user_collection_request(tls, force, exhaustive);
 }
 
 #[unsafe(no_mangle)]
 pub extern "C" fn mmtk_set_gc_enabled(enable: bool) {
-    if enable {
-        crate::mmtk().enable_collection();
-    } else if crate::mmtk().is_collection_enabled() {
-        // We must not call disable_collection if GC is already disabled because
-        // MMTk keeps track of the depth of disable_collection. This would require
-        // enable_collection calls equal to the number of disable_collection calls
-        // before GC is enabled.
-
-        // We cannot panic on failing to disable GC because
-        // DURING_GC_COULD_MALLOC_REGION_START disables GC during GC. It's
-        // fine to ignore the failure because MMTk's malloc does not trigger
-        // GC.
-        let _ = crate::mmtk().disable_collection();
-    }
+    crate::binding().gc_disable.set_internal_disabled(!enable);
 }
 
 #[unsafe(no_mangle)]
 pub extern "C" fn mmtk_gc_enabled_p() -> bool {
-    crate::mmtk().is_collection_enabled()
+    !crate::binding().gc_disable.internal_disabled()
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn mmtk_set_user_gc_disabled(disable: bool) -> bool {
+    crate::binding().gc_disable.set_user_disabled(disable)
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn mmtk_user_gc_disabled_p() -> bool {
+    crate::binding().gc_disable.user_disabled()
 }
 
 // =============== Object allocation ===============

gc/mmtk/src/binding.rs

@@ -8,6 +8,7 @@ use mmtk::util::ObjectReference;
 
 use crate::Ruby;
 use crate::abi::RubyBindingOptions;
+use crate::gc_disable::GcDisableState;
 use crate::pinning_registry::PinningRegistry;
 use crate::weak_proc::WeakProcessor;
 
@@ -35,6 +36,7 @@ pub struct RubyBinding {
     pub pinning_registry: PinningRegistry,
     pub gc_thread_join_handles: Mutex<Vec<JoinHandle<()>>>,
     pub wb_unprotected_objects: Mutex<HashSet<ObjectReference>>,
+    pub gc_disable: GcDisableState,
 }
 
 unsafe impl Sync for RubyBinding {}
@@ -54,6 +56,7 @@ impl RubyBinding {
             pinning_registry: PinningRegistry::new(),
             gc_thread_join_handles: Default::default(),
             wb_unprotected_objects: Default::default(),
+            gc_disable: Default::default(),
         }
     }
 

gc/mmtk/src/collection.rs

@@ -30,6 +30,7 @@ impl Collection<Ruby> for VMCollection {
         F: FnMut(&'static mut mmtk::Mutator<Ruby>),
     {
         upcalls::stop_the_world();
+        crate::binding().gc_disable.on_pause_start();
 
         if crate::mmtk().get_plan().current_gc_may_move_object() {
             CURRENT_GC_MAY_MOVE.store(true, Ordering::Relaxed);
@@ -47,6 +48,8 @@ impl Collection<Ruby> for VMCollection {
     }
 
     fn resume_mutators(_tls: VMWorkerThread) {
+        crate::binding().gc_disable.on_pause_end();
+
         let current_gc_may_move = CURRENT_GC_MAY_MOVE.load(Ordering::Relaxed);
 
         if current_gc_may_move {

gc/mmtk/src/gc_disable.rs

@@ -0,0 +1,231 @@
+use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+
+use mmtk::util::VMMutatorThread;
+
+/// Ruby's GC-disable state, owned by the binding instead of mmtk-core.
+///
+/// mmtk-core cannot run a GC while collection is disabled, and every pause ends in `NotInGC`.
+/// A forced `GC.start` therefore has to enable collection in mmtk-core, and mmtk-core alone
+/// would lose the disabled state across that GC. The binding keeps the state here and mirrors
+/// it into mmtk-core, where it holds at most one `disable_collection` once all callers return.
+///
+/// CRuby's fatal-signal handlers call `rb_gc_impl_gc_disable` and `rb_gc_impl_gc_enable`, and a
+/// handler can interrupt a thread anywhere in this module. The state is therefore an atomic word.
+/// `apply` runs one caller at a time. A caller that read an old state corrects mmtk-core before
+/// it returns, so a stale step never undoes the result of an `apply` that already returned. A
+/// handler that interrupts the applying thread does not wait for it, because it would wait
+/// forever.
+#[derive(Default)]
+pub struct GcDisableState {
+    state: AtomicU64,
+    /// The thread that runs `apply`, as returned by `current_thread_id`, or 0.
+    applier: AtomicUsize,
+}
+
+/// A nonzero value that is unique to the calling thread. It must be safe to call in a signal
+/// handler. A Rust thread-local is not: this library is loaded with `dlopen`, so the first access
+/// on a thread can allocate in `__tls_get_addr`. `pthread_self` only reads the thread pointer.
+fn current_thread_id() -> usize {
+    unsafe { libc::pthread_self() as usize }
+}
+
+#[derive(Clone, Copy)]
+struct State {
+    /// Set by `rb_gc_impl_gc_disable` and cleared by `rb_gc_impl_gc_enable`.
+    internal_disabled: bool,
+    /// Set and cleared by `GC.disable` and `GC.enable`.
+    user_disabled: bool,
+    /// Forced `GC.start` calls that enabled collection and wait for a pause.
+    forced_gc_requests: u32,
+    /// Set between `stop_all_mutators` and `resume_mutators`.
+    in_pause: bool,
+    /// Incremented when a pause ends. A pause serves every forced request made before it ends.
+    pause_epoch: u32,
+}
+
+impl State {
+    const INTERNAL_DISABLED: u64 = 1 << 0;
+    const USER_DISABLED: u64 = 1 << 1;
+    const IN_PAUSE: u64 = 1 << 2;
+    const FORCED_SHIFT: u32 = 3;
+    const FORCED_MAX: u32 = (1 << 29) - 1;
+    const EPOCH_SHIFT: u32 = 32;
+
+    fn decode(bits: u64) -> Self {
+        Self {
+            internal_disabled: bits & Self::INTERNAL_DISABLED != 0,
+            user_disabled: bits & Self::USER_DISABLED != 0,
+            in_pause: bits & Self::IN_PAUSE != 0,
+            forced_gc_requests: (bits >> Self::FORCED_SHIFT) as u32 & Self::FORCED_MAX,
+            pause_epoch: (bits >> Self::EPOCH_SHIFT) as u32,
+        }
+    }
+
+    fn encode(self) -> u64 {
+        debug_assert!(self.forced_gc_requests <= Self::FORCED_MAX);
+        (self.internal_disabled as u64)
+            | (self.user_disabled as u64) << 1
+            | (self.in_pause as u64) << 2
+            | (self.forced_gc_requests as u64) << Self::FORCED_SHIFT
+            | (self.pause_epoch as u64) << Self::EPOCH_SHIFT
+    }
+
+    fn want_disabled(self) -> bool {
+        (self.internal_disabled || self.user_disabled) && self.forced_gc_requests == 0
+    }
+}
+
+impl GcDisableState {
+    fn load(&self) -> State {
+        State::decode(self.state.load(Ordering::SeqCst))
+    }
+
+    /// Returns the state before the update.
+    fn update(&self, f: impl Fn(&mut State)) -> State {
+        let bits = self
+            .state
+            .update(Ordering::SeqCst, Ordering::SeqCst, |bits| {
+                let mut state = State::decode(bits);
+                f(&mut state);
+                state.encode()
+            });
+        State::decode(bits)
+    }
+
+    fn apply(&self) {
+        let me = current_thread_id();
+
+        // A signal handler interrupted this thread in `apply`. CRuby's handlers never return to
+        // the interrupted code: they abort or longjmp. Apply the state and release `applier`,
+        // so that a longjmp does not leave it held.
+        if self.applier.load(Ordering::SeqCst) == me {
+            self.mirror_into_mmtk();
+            let _ = self
+                .applier
+                .compare_exchange(me, 0, Ordering::SeqCst, Ordering::SeqCst);
+            return;
+        }
+
+        while self
+            .applier
+            .compare_exchange_weak(0, me, Ordering::SeqCst, Ordering::SeqCst)
+            .is_err()
+        {
+            std::thread::yield_now();
+        }
+        self.mirror_into_mmtk();
+        let _ = self
+            .applier
+            .compare_exchange(me, 0, Ordering::SeqCst, Ordering::SeqCst);
+    }
+
+    /// Brings mmtk-core to the state that `want_disabled` asks for. The binding is the only
+    /// caller of `disable_collection` and `enable_collection`, so mmtk-core's status is the
+    /// mirror.
+    fn mirror_into_mmtk(&self) {
+        let mmtk = crate::mmtk();
+        loop {
+            let want_disabled = self.load().want_disabled();
+            let disabled = !mmtk.is_collection_enabled();
+            if want_disabled == disabled {
+                return;
+            }
+
+            if want_disabled {
+                match mmtk.disable_collection() {
+                    Ok(true) => {}
+                    // Only if a signal handler ran `apply` on this thread and then returned.
+                    // Remove the extra nesting level.
+                    Ok(false) => {
+                        mmtk.enable_collection();
+                    }
+                    // A pause is requested or running, or collection is not initialized. The end
+                    // of that pause, or the initialization of collection, applies the state again.
+                    Err(_) => return,
+                }
+            } else {
+                mmtk.enable_collection();
+            }
+        }
+    }
+
+    /// GC workers run CRuby's `DURING_GC_COULD_MALLOC_REGION` in parallel. Its read-then-set of
+    /// this flag is not atomic, so two workers can leave it set. Mutators are stopped during a
+    /// pause, so writes made then are dropped, as mmtk-core dropped them before this state existed.
+    pub fn set_internal_disabled(&self, disabled: bool) {
+        let before = self.update(|state| {
+            if !state.in_pause {
+                state.internal_disabled = disabled;
+            }
+        });
+        if !before.in_pause {
+            self.apply();
+        }
+    }
+
+    pub fn internal_disabled(&self) -> bool {
+        self.load().internal_disabled
+    }
+
+    /// Returns the previous value.
+    pub fn set_user_disabled(&self, disabled: bool) -> bool {
+        let was = self
+            .update(|state| state.user_disabled = disabled)
+            .user_disabled;
+        self.apply();
+        was
+    }
+
+    pub fn user_disabled(&self) -> bool {
+        self.load().user_disabled
+    }
+
+    /// `GC.start` (`force`) runs even when GC is disabled. `GC.stress` does not.
+    pub fn handle_user_collection_request(
+        &self,
+        tls: VMMutatorThread,
+        force: bool,
+        exhaustive: bool,
+    ) {
+        let mmtk = crate::mmtk();
+
+        if !force {
+            mmtk.handle_user_collection_request(tls, false, exhaustive);
+            return;
+        }
+
+        let epoch = self
+            .update(|state| state.forced_gc_requests += 1)
+            .pause_epoch;
+
+        self.apply();
+        mmtk.handle_user_collection_request(tls, true, exhaustive);
+
+        self.update(|state| {
+            if state.pause_epoch == epoch {
+                state.forced_gc_requests = state.forced_gc_requests.saturating_sub(1);
+            }
+        });
+        self.apply();
+    }
+
+    /// Called after mmtk-core has stopped every mutator and before GC work starts.
+    pub fn on_pause_start(&self) {
+        self.update(|state| state.in_pause = true);
+    }
+
+    /// Called after mmtk-core sets `NotInGC` at the end of a pause and before mutators resume,
+    /// so no mutator can request another pause before the disabled state is applied.
+    pub fn on_pause_end(&self) {
+        self.update(|state| {
+            state.in_pause = false;
+            state.forced_gc_requests = 0;
+            state.pause_epoch = state.pause_epoch.wrapping_add(1);
+        });
+        self.apply();
+    }
+
+    pub fn on_collection_initialized(&self) {
+        self.apply();
+    }
+}

gc/mmtk/src/lib.rs

@@ -26,6 +26,7 @@ pub mod active_plan;
 pub mod api;
 pub mod binding;
 pub mod collection;
+pub mod gc_disable;
 pub mod heap;
 pub mod object_model;
 pub mod pinning_registry;