YJIT: Add RubyVM::YJIT.code_gc

k0kubun · k0kubun · commit c706b917b160 · 2022-10-27T17:23:56.000-07:00
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
@@ -830,10 +830,10 @@ def foo
   def test_code_gc
     assert_compiles(code_gc_helpers + <<~'RUBY', exits: :any, result: :ok)
       return :not_paged unless add_pages(100) # prepare freeable pages
-      code_gc # first code GC
+      RubyVM::YJIT.code_gc # first code GC
       return :not_compiled1 unless compiles { nil } # should be JITable again
 
-      code_gc # second code GC
+      RubyVM::YJIT.code_gc # second code GC
       return :not_compiled2 unless compiles { nil } # should be JITable again
 
       code_gc_count = RubyVM::YJIT.runtime_stats[:code_gc_count]
@@ -854,7 +854,7 @@ def test_on_stack_code_gc_call
 
       return :not_paged1 unless add_pages(400) # go to a page without initial ocb code
       return :broken_resume1 if fiber.resume != 0 # JIT the fiber
-      code_gc # first code GC, which should not free the fiber page
+      RubyVM::YJIT.code_gc # first code GC, which should not free the fiber page
       return :broken_resume2 if fiber.resume != 0 # The code should be still callable
 
       code_gc_count = RubyVM::YJIT.runtime_stats[:code_gc_count]
@@ -873,19 +873,19 @@ def test_on_stack_code_gc_twice
 
       return :not_paged1 unless add_pages(400) # go to a page without initial ocb code
       return :broken_resume1 if fiber.resume(true) != 0 # JIT the fiber
-      code_gc # first code GC, which should not free the fiber page
+      RubyVM::YJIT.code_gc # first code GC, which should not free the fiber page
 
       return :not_paged2 unless add_pages(300) # add some stuff to be freed
       # Not calling fiber.resume here to test the case that the YJIT payload loses some
       # information at the previous code GC. The payload should still be there, and
       # thus we could know the fiber ISEQ is still on stack on this second code GC.
-      code_gc # second code GC, which should still not free the fiber page
+      RubyVM::YJIT.code_gc # second code GC, which should still not free the fiber page
 
       return :not_paged3 unless add_pages(200) # attempt to overwrite the fiber page (it shouldn't)
       return :broken_resume2 if fiber.resume(true) != 0 # The fiber code should be still fine
 
       return :broken_resume3 if fiber.resume(false) != nil # terminate the fiber
-      code_gc # third code GC, freeing a page that used to be on stack
+      RubyVM::YJIT.code_gc # third code GC, freeing a page that used to be on stack
 
       return :not_paged4 unless add_pages(100) # check everything still works
 
@@ -933,11 +933,6 @@ def add_pages(num_jits)
         num_jits.times { return false unless eval('compiles { nil.to_i }') }
         pages.nil? || pages < RubyVM::YJIT.runtime_stats[:compiled_page_count]
       end
-
-      def code_gc
-        RubyVM::YJIT.simulate_oom! # bump write_pos
-        eval('proc { nil }.call') # trigger code GC
-      end
     RUBY
   end
 
diff --git a/yjit.c b/yjit.c
@@ -1053,6 +1053,7 @@ VALUE rb_yjit_get_stats(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_reset_stats_bang(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_disasm_iseq(rb_execution_context_t *ec, VALUE self, VALUE iseq);
 VALUE rb_yjit_insns_compiled(rb_execution_context_t *ec, VALUE self, VALUE iseq);
+VALUE rb_yjit_code_gc(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_simulate_oom_bang(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_get_exit_locations(rb_execution_context_t *ec, VALUE self);
 
diff --git a/yjit.rb b/yjit.rb
@@ -162,6 +162,11 @@ def self.insns_compiled(iseq)
     end
   end
 
+  # Free and recompile all existing JIT code
+  def self.code_gc
+    Primitive.rb_yjit_code_gc
+  end
+
   def self.simulate_oom!
     Primitive.rb_yjit_simulate_oom_bang
   end
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
@@ -79,6 +79,18 @@ pub extern "C" fn rb_yjit_iseq_gen_entry_point(iseq: IseqPtr, ec: EcPtr) -> *con
     }
 }
 
+/// Free and recompile all existing JIT code
+#[no_mangle]
+pub extern "C" fn rb_yjit_code_gc(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+    if !yjit_enabled_p() {
+        return Qnil;
+    }
+
+    let cb = CodegenGlobals::get_inline_cb();
+    cb.code_gc();
+    Qnil
+}
+
 /// Simulate a situation where we are out of executable memory
 #[no_mangle]
 pub extern "C" fn rb_yjit_simulate_oom_bang(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
