YJIT: Disable code GC (#8865)

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

.github/workflows/yjit-macos.yml

@@ -63,7 +63,7 @@ jobs:
             yjit_opts: '--yjit'
           - test_task: 'check'
             configure: '--enable-yjit=dev'
-            yjit_opts: '--yjit-call-threshold=1 --yjit-verify-ctx'
+            yjit_opts: '--yjit-call-threshold=1 --yjit-verify-ctx --yjit-code-gc'
       fail-fast: false
 
     env:

.github/workflows/yjit-ubuntu.yml

@@ -99,7 +99,7 @@ jobs:
 
           - test_task: 'check'
             configure: '--enable-yjit=dev'
-            yjit_opts: '--yjit-call-threshold=1 --yjit-verify-ctx'
+            yjit_opts: '--yjit-call-threshold=1 --yjit-verify-ctx --yjit-code-gc'
 
           - test_task: 'test-bundled-gems'
             configure: '--enable-yjit=dev'

test/ruby/test_yjit.rb

@@ -1058,8 +1058,30 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_disable_code_gc_with_many_iseqs
+    assert_compiles(code_gc_helpers + <<~'RUBY', exits: :any, result: :ok, mem_size: 1, code_gc: false)
+      fiber = Fiber.new {
+        # Loop to call the same basic block again after Fiber.yield
+        while true
+          Fiber.yield(nil.to_i)
+        end
+      }
+
+      return :not_paged1 unless add_pages(250) # use some pages
+      return :broken_resume1 if fiber.resume != 0 # leave an on-stack code as well
+
+      add_pages(2000) # use a whole lot of pages to run out of 1MiB
+      return :broken_resume2 if fiber.resume != 0 # on-stack code should be callable
+
+      code_gc_count = RubyVM::YJIT.runtime_stats[:code_gc_count]
+      return :"code_gc_#{code_gc_count}" if code_gc_count != 0
+
+      :ok
+    RUBY
+  end
+
   def test_code_gc_with_many_iseqs
-    assert_compiles(code_gc_helpers + <<~'RUBY', exits: :any, result: :ok, mem_size: 1)
+    assert_compiles(code_gc_helpers + <<~'RUBY', exits: :any, result: :ok, mem_size: 1, code_gc: true)
       fiber = Fiber.new {
         # Loop to call the same basic block again after Fiber.yield
         while true
@@ -1425,7 +1447,7 @@ class TestYJIT < Test::Unit::TestCase
   end
 
   ANY = Object.new
-  def assert_compiles(test_script, insns: [], call_threshold: 1, stdout: nil, exits: {}, result: ANY, frozen_string_literal: nil, mem_size: nil)
+  def assert_compiles(test_script, insns: [], call_threshold: 1, stdout: nil, exits: {}, result: ANY, frozen_string_literal: nil, mem_size: nil, code_gc: false)
     reset_stats = <<~RUBY
       RubyVM::YJIT.runtime_stats
       RubyVM::YJIT.reset_stats!
@@ -1459,7 +1481,7 @@ class TestYJIT < Test::Unit::TestCase
       #{write_results}
     RUBY
 
-    status, out, err, stats = eval_with_jit(script, call_threshold:, mem_size:)
+    status, out, err, stats = eval_with_jit(script, call_threshold:, mem_size:, code_gc:)
 
     assert status.success?, "exited with status #{status.to_i}, stderr:\n#{err}"
 
@@ -1520,13 +1542,14 @@ class TestYJIT < Test::Unit::TestCase
     s.chars.map { |c| c.ascii_only? ? c : "\\u%x" % c.codepoints[0] }.join
   end
 
-  def eval_with_jit(script, call_threshold: 1, timeout: 1000, mem_size: nil)
+  def eval_with_jit(script, call_threshold: 1, timeout: 1000, mem_size: nil, code_gc: false)
     args = [
       "--disable-gems",
       "--yjit-call-threshold=#{call_threshold}",
       "--yjit-stats=quiet"
     ]
     args << "--yjit-exec-mem-size=#{mem_size}" if mem_size
+    args << "--yjit-code-gc" if code_gc
     args << "-e" << script_shell_encode(script)
     stats_r, stats_w = IO.pipe
     # Separate thread so we don't deadlock when

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

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;