YJIT: Disable code GC (#8865)

Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>
Co-authored-by: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>

2 files changed, 30 insertions, 6 deletions

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index a091272470..61661e1a2c 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -2234,7 +2234,9 @@ pub fn gen_entry_point(iseq: IseqPtr, ec: EcPtr, jit_exception: bool) -> Option<
         // Compilation failed
         None => {
             // Trigger code GC. This entry point will be recompiled later.
-            cb.code_gc(ocb);
+            if get_option!(code_gc) {
+                cb.code_gc(ocb);
+            }
             return None;
         }
 
@@ -2300,7 +2302,9 @@ c_callable! {
                     .unwrap_or_else(|| {
                         // Trigger code GC (e.g. no space).
                         // This entry point will be recompiled later.
-                        cb.code_gc(ocb);
+                        if get_option!(code_gc) {
+                            cb.code_gc(ocb);
+                        }
                         CodegenGlobals::get_stub_exit_code().raw_ptr(cb)
                     });
 
@@ -2548,6 +2552,11 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
         // So we do it here instead.
         rb_set_cfp_sp(cfp, reconned_sp);
 
+        // Bail if code GC is disabled and we've already run out of spaces.
+        if !get_option!(code_gc) && (cb.has_dropped_bytes() || ocb.unwrap().has_dropped_bytes()) {
+            return CodegenGlobals::get_stub_exit_code().raw_ptr(cb);
+        }
+
         // Bail if we're about to run out of native stack space.
         // We've just reconstructed interpreter state.
         if rb_ec_stack_check(ec as _) != 0 {
@@ -2564,7 +2573,6 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
     if block.is_none() {
         let branch_old_shape = branch.gen_fn.get_shape();
 
-
         // If the new block can be generated right after the branch (at cb->write_pos)
         if cb.get_write_ptr() == branch.end_addr.get() {
             // This branch should be terminating its block
@@ -2622,7 +2630,9 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
             // because incomplete code could be used when cb.dropped_bytes is flipped
             // by code GC. So this place, after all compilation, is the safest place
             // to hook code GC on branch_stub_hit.
-            cb.code_gc(ocb);
+            if get_option!(code_gc) {
+                cb.code_gc(ocb);
+            }
 
             // Failed to service the stub by generating a new block so now we
             // need to exit to the interpreter at the stubbed location. We are
diff --git a/yjit/src/options.rs b/yjit/src/options.rs
index 98c0548677..3980eae0ae 100644
--- a/yjit/src/options.rs
+++ b/yjit/src/options.rs
@@ -72,6 +72,9 @@ pub struct Options {
     /// Enable generating frame pointers (for x86. arm64 always does this)
     pub frame_pointer: bool,
 
+    /// Run code GC when exec_mem_size is reached.
+    pub code_gc: bool,
+
     /// Enable writing /tmp/perf-{pid}.map for Linux perf
     pub perf_map: bool,
 }
@@ -92,15 +95,17 @@ pub static mut OPTIONS: Options = Options {
     verify_ctx: false,
     dump_iseq_disasm: None,
     frame_pointer: false,
+    code_gc: false,
     perf_map: false,
 };
 
 /// YJIT option descriptions for `ruby --help`.
-static YJIT_OPTIONS: [(&str, &str); 8] = [
+static YJIT_OPTIONS: [(&str, &str); 9] = [
     ("--yjit-stats",                    "Enable collecting YJIT statistics"),
     ("--yjit-trace-exits",              "Record Ruby source location when exiting from generated code"),
     ("--yjit-trace-exits-sample-rate",  "Trace exit locations only every Nth occurrence"),
     ("--yjit-exec-mem-size=num",        "Size of executable memory block in MiB (default: 128)"),
+    ("--yjit-code-gc",                  "Run code GC when the code size reaches the limit"),
     ("--yjit-call-threshold=num",       "Number of calls to trigger JIT"),
     ("--yjit-cold-threshold=num",       "Global call after which ISEQs not compiled (default: 200K)"),
     ("--yjit-max-versions=num",         "Maximum number of versions per basic block (default: 4)"),
@@ -120,7 +125,12 @@ macro_rules! get_option {
     // Unsafe is ok here because options are initialized
     // once before any Ruby code executes
     ($option_name:ident) => {
-        unsafe { OPTIONS.$option_name }
+        {
+            // Make this a statement since attributes on expressions are experimental
+            #[allow(unused_unsafe)]
+            let ret = unsafe { OPTIONS.$option_name };
+            ret
+        }
     };
 }
 pub(crate) use get_option;
@@ -204,6 +214,10 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             }
         },
 
+        ("code-gc", _) => unsafe {
+            OPTIONS.code_gc = true;
+        },
+
         ("perf", _) => match opt_val {
             "" => unsafe {
                 OPTIONS.frame_pointer = true;