YJIT: GC and recompile all code pages (ruby#6406)

when it fails to allocate a new page.

Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>

Author: 2 people

cont.c

@@ -69,7 +69,7 @@ static VALUE rb_cFiberPool;
 #define FIBER_POOL_ALLOCATION_FREE
 #endif
 
-#define jit_cont_enabled mjit_enabled // To be used by YJIT later
+#define jit_cont_enabled (mjit_enabled || rb_yjit_enabled_p())
 
 enum context_type {
     CONTINUATION_CONTEXT = 0,

test/ruby/test_yjit.rb

@@ -825,12 +825,126 @@ def foo
     RUBY
   end
 
+  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
+      return :not_compiled1 unless compiles { nil } # should be JITable again
+
+      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]
+      return :"code_gc_#{code_gc_count}" if code_gc_count && code_gc_count != 2
+
+      :ok
+    RUBY
+  end
+
+  def test_on_stack_code_gc_call
+    assert_compiles(code_gc_helpers + <<~'RUBY', exits: :any, result: :ok)
+      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(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
+      return :broken_resume2 if fiber.resume != 0 # The code should be still callable
+
+      code_gc_count = RubyVM::YJIT.runtime_stats[:code_gc_count]
+      return :"code_gc_#{code_gc_count}" if code_gc_count && code_gc_count != 1
+
+      :ok
+    RUBY
+  end
+
+  def test_on_stack_code_gc_twice
+    assert_compiles(code_gc_helpers + <<~'RUBY', exits: :any, result: :ok)
+      fiber = Fiber.new {
+        # Loop to call the same basic block again after Fiber.yield
+        while Fiber.yield(nil.to_i); end
+      }
+
+      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
+
+      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
+
+      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
+
+      return :not_paged4 unless add_pages(100) # check everything still works
+
+      code_gc_count = RubyVM::YJIT.runtime_stats[:code_gc_count]
+      return :"code_gc_#{code_gc_count}" if code_gc_count && code_gc_count != 3
+
+      :ok
+    RUBY
+  end
+
+  def test_code_gc_with_many_iseqs
+    assert_compiles(code_gc_helpers + <<~'RUBY', exits: :any, result: :ok, mem_size: 1)
+      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(500) # 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 && code_gc_count == 0
+
+      :ok
+    RUBY
+  end
+
+  private
+
+  def code_gc_helpers
+    <<~'RUBY'
+      def compiles(&block)
+        failures = RubyVM::YJIT.runtime_stats[:compilation_failure]
+        block.call
+        failures == RubyVM::YJIT.runtime_stats[:compilation_failure]
+      end
+
+      def add_pages(num_jits)
+        pages = RubyVM::YJIT.runtime_stats[:compiled_page_count]
+        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
+
   def assert_no_exits(script)
     assert_compiles(script)
   end
 
   ANY = Object.new
-  def assert_compiles(test_script, insns: [], call_threshold: 1, stdout: nil, exits: {}, result: ANY, frozen_string_literal: nil)
+  def assert_compiles(test_script, insns: [], call_threshold: 1, stdout: nil, exits: {}, result: ANY, frozen_string_literal: nil, mem_size: nil)
     reset_stats = <<~RUBY
       RubyVM::YJIT.runtime_stats
       RubyVM::YJIT.reset_stats!
@@ -864,7 +978,7 @@ def collect_insns(iseq)
       #{write_results}
     RUBY
 
-    status, out, err, stats = eval_with_jit(script, call_threshold: call_threshold)
+    status, out, err, stats = eval_with_jit(script, call_threshold:, mem_size:)
 
     assert status.success?, "exited with status #{status.to_i}, stderr:\n#{err}"
 
@@ -918,12 +1032,13 @@ def script_shell_encode(s)
     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)
+  def eval_with_jit(script, call_threshold: 1, timeout: 1000, mem_size: nil)
     args = [
       "--disable-gems",
       "--yjit-call-threshold=#{call_threshold}",
       "--yjit-stats"
     ]
+    args << "--yjit-exec-mem-size=#{mem_size}" if mem_size
     args << "-e" << script_shell_encode(script)
     stats_r, stats_w = IO.pipe
     out, err, status = EnvUtil.invoke_ruby(args,

yjit.c

@@ -27,6 +27,7 @@
 #include "probes_helper.h"
 #include "iseq.h"
 #include "ruby/debug.h"
+#include "internal/cont.h"
 
 // For mmapp(), sysconf()
 #ifndef _WIN32
@@ -65,10 +66,7 @@ STATIC_ASSERT(pointer_tagging_scheme, USE_FLONUM);
 bool
 rb_yjit_mark_writable(void *mem_block, uint32_t mem_size)
 {
-    if (mprotect(mem_block, mem_size, PROT_READ | PROT_WRITE)) {
-        return false;
-    }
-    return true;
+    return mprotect(mem_block, mem_size, PROT_READ | PROT_WRITE) == 0;
 }
 
 void
@@ -85,6 +83,20 @@ rb_yjit_mark_executable(void *mem_block, uint32_t mem_size)
     }
 }
 
+// Free the specified memory block.
+bool
+rb_yjit_mark_unused(void *mem_block, uint32_t mem_size)
+{
+    // On Linux, you need to use madvise MADV_DONTNEED to free memory.
+    // We might not need to call this on macOS, but it's not really documented.
+    // We generally prefer to do the same thing on both to ease testing too.
+    madvise(mem_block, mem_size, MADV_DONTNEED);
+
+    // On macOS, mprotect PROT_NONE seems to reduce RSS.
+    // We also call this on Linux to avoid executing unused pages.
+    return mprotect(mem_block, mem_size, PROT_NONE) == 0;
+}
+
 // `start` is inclusive and `end` is exclusive.
 void
 rb_yjit_icache_invalidate(void *start, void *end)
@@ -387,6 +399,9 @@ rb_iseq_reset_jit_func(const rb_iseq_t *iseq)
 {
     RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(iseq, imemo_iseq));
     iseq->body->jit_func = NULL;
+    // Enable re-compiling this ISEQ. Event when it's invalidated for TracePoint,
+    // we'd like to re-compile ISEQs that haven't been converted to trace_* insns.
+    iseq->body->total_calls = 0;
 }
 
 // Get the PC for a given index in an iseq

yjit.rb

@@ -212,13 +212,17 @@ def _print_stats
       $stderr.puts "bindings_allocations:  " + ("%10d" % stats[:binding_allocations])
       $stderr.puts "bindings_set:          " + ("%10d" % stats[:binding_set])
       $stderr.puts "compilation_failure:   " + ("%10d" % compilation_failure) if compilation_failure != 0
-      $stderr.puts "compiled_iseq_count:   " + ("%10d" % stats[:compiled_iseq_count])
       $stderr.puts "compiled_block_count:  " + ("%10d" % stats[:compiled_block_count])
+      $stderr.puts "compiled_iseq_count:   " + ("%10d" % stats[:compiled_iseq_count])
+      $stderr.puts "compiled_page_count:   " + ("%10d" % stats[:compiled_page_count])
       $stderr.puts "freed_iseq_count:      " + ("%10d" % stats[:freed_iseq_count])
+      $stderr.puts "freed_page_count:      " + ("%10d" % stats[:freed_page_count])
       $stderr.puts "invalidation_count:    " + ("%10d" % stats[:invalidation_count])
       $stderr.puts "constant_state_bumps:  " + ("%10d" % stats[:constant_state_bumps])
       $stderr.puts "inline_code_size:      " + ("%10d" % stats[:inline_code_size])
       $stderr.puts "outlined_code_size:    " + ("%10d" % stats[:outlined_code_size])
+      $stderr.puts "freed_code_size:       " + ("%10d" % stats[:freed_code_size])
+      $stderr.puts "code_gc_count:         " + ("%10d" % stats[:code_gc_count])
       $stderr.puts "num_gc_obj_refs:       " + ("%10d" % stats[:num_gc_obj_refs])
 
       $stderr.puts "total_exit_count:      " + ("%10d" % total_exits)

yjit/bindgen/src/main.rs

@@ -263,6 +263,7 @@ fn main() {
         .allowlist_function("rb_yjit_reserve_addr_space")
         .allowlist_function("rb_yjit_mark_writable")
         .allowlist_function("rb_yjit_mark_executable")
+        .allowlist_function("rb_yjit_mark_unused")
         .allowlist_function("rb_yjit_get_page_size")
         .allowlist_function("rb_leaf_invokebuiltin_iseq_p")
         .allowlist_function("rb_leaf_builtin_function")
@@ -297,6 +298,9 @@ fn main() {
         // From internal/compile.h
         .allowlist_function("rb_vm_insn_decode")
 
+        // from internal/cont.h
+        .allowlist_function("rb_jit_cont_each_iseq")
+
         // From iseq.h
         .allowlist_function("rb_vm_insn_addr2opcode")
         .allowlist_function("rb_iseqw_to_iseq")