YJIT: Add `--yjit-pause` and `RubyVM::YJIT.resume` (#7609)

* YJIT: Add --yjit-pause and RubyVM::YJIT.resume

This allows booting YJIT in a suspended state. We chose to add a new
command line option as opposed to simply allowing YJIT.resume to work
without any command line option because it allows for combining with
YJIT tuning command line options. It also simpifies implementation.

Paired with Kokubun and Maxime.

* Update yjit.rb

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

---------

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>
Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

test/ruby/test_yjit.rb

@@ -51,6 +51,29 @@ class TestYJIT < Test::Unit::TestCase
     #assert_in_out_err('--yjit-call-threshold=', '', [], /--yjit-call-threshold needs an argument/)
   end
 
+  def test_starting_paused
+    program = <<~RUBY
+      def not_compiled = nil
+      def will_compile = nil
+      def compiled_counts = RubyVM::YJIT.runtime_stats[:compiled_iseq_count]
+      counts = []
+      not_compiled
+      counts << compiled_counts
+
+      RubyVM::YJIT.resume
+
+      will_compile
+      counts << compiled_counts
+
+      if counts[0] == 0 && counts[1] > 0
+        p :ok
+      end
+    RUBY
+    assert_in_out_err(%w[--yjit-pause --yjit-stats --yjit-call-threshold=1], program, success: true) do |stdout, stderr|
+      assert_equal([":ok"], stdout)
+    end
+  end
+
   def test_yjit_stats_and_v_no_error
     _stdout, stderr, _status = EnvUtil.invoke_ruby(%w(-v --yjit-stats), '', true, true)
     refute_includes(stderr, "NoMethodError")

vm.c

@@ -376,7 +376,7 @@ jit_compile(rb_execution_context_t *ec)
     // Increment the ISEQ's call counter
     const rb_iseq_t *iseq = ec->cfp->iseq;
     struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
-    bool yjit_enabled = rb_yjit_enabled_p();
+    bool yjit_enabled = rb_yjit_compile_new_iseqs();
     if (yjit_enabled || rb_rjit_call_p) {
         body->total_calls++;
     }

yjit.c

@@ -1114,6 +1114,7 @@ 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);
+VALUE rb_yjit_resume(rb_execution_context_t *ec, VALUE self);
 
 // Preprocessed yjit.rb generated during build
 #include "yjit.rbinc"

yjit.h

@@ -26,6 +26,7 @@
 
 // Expose these as declarations since we are building YJIT.
 bool rb_yjit_enabled_p(void);
+bool rb_yjit_compile_new_iseqs(void);
 unsigned rb_yjit_call_threshold(void);
 void rb_yjit_invalidate_all_method_lookup_assumptions(void);
 void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme);
@@ -48,6 +49,7 @@ void rb_yjit_tracing_invalidate_all(void);
 // In these builds, YJIT could never be turned on. Provide dummy implementations.
 
 static inline bool rb_yjit_enabled_p(void) { return false; }
+static inline bool rb_yjit_compile_new_iseqs(void) { return false; }
 static inline unsigned rb_yjit_call_threshold(void) { return UINT_MAX; }
 static inline void rb_yjit_invalidate_all_method_lookup_assumptions(void) {}
 static inline void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme) {}

yjit.rb

@@ -29,6 +29,11 @@ module RubyVM::YJIT
     Primitive.rb_yjit_reset_stats_bang
   end
 
+  # Resume YJIT compilation after paused on startup with --yjit-pause
+  def self.resume
+    Primitive.rb_yjit_resume
+  end
+
   # If --yjit-trace-exits is enabled parse the hashes from
   # Primitive.rb_yjit_get_exit_locations into a format readable
   # by Stackprof. This will allow us to find the exact location of a

yjit/src/options.rs

@@ -28,6 +28,10 @@ pub struct Options {
     // Trace locations of exits
     pub gen_trace_exits: bool,
 
+    // Whether to start YJIT in paused state (initialize YJIT but don't
+    // compile anything)
+    pub pause: bool,
+
     /// Dump compiled and executed instructions for debugging
     pub dump_insns: bool,
 
@@ -50,6 +54,7 @@ pub static mut OPTIONS: Options = Options {
     max_versions: 4,
     gen_stats: false,
     gen_trace_exits: false,
+    pause: false,
     dump_insns: false,
     dump_disasm: None,
     verify_ctx: false,
@@ -132,6 +137,10 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             }
         },
 
+        ("pause", "") => unsafe {
+            OPTIONS.pause = true;
+        },
+
         ("dump-disasm", _) => match opt_val.to_string().as_str() {
             "" => unsafe { OPTIONS.dump_disasm = Some(DumpDisasm::Stdout) },
             directory => {

yjit/src/yjit.rs

@@ -15,6 +15,9 @@ use std::sync::atomic::{AtomicBool, Ordering};
 /// See [rb_yjit_enabled_p]
 static YJIT_ENABLED: AtomicBool = AtomicBool::new(false);
 
+/// When false, we don't compile new iseqs, but might still service existing branch stubs.
+static COMPILE_NEW_ISEQS: AtomicBool = AtomicBool::new(false);
+
 /// Parse one command-line option.
 /// This is called from ruby.c
 #[no_mangle]
@@ -32,6 +35,11 @@ pub extern "C" fn rb_yjit_enabled_p() -> raw::c_int {
     YJIT_ENABLED.load(Ordering::Acquire).into()
 }
 
+#[no_mangle]
+pub extern "C" fn rb_yjit_compile_new_iseqs() -> bool {
+    COMPILE_NEW_ISEQS.load(Ordering::Acquire).into()
+}
+
 /// Like rb_yjit_enabled_p, but for Rust code.
 pub fn yjit_enabled_p() -> bool {
     YJIT_ENABLED.load(Ordering::Acquire)
@@ -60,6 +68,8 @@ pub extern "C" fn rb_yjit_init_rust() {
 
         // YJIT enabled and initialized successfully
         YJIT_ENABLED.store(true, Ordering::Release);
+
+        COMPILE_NEW_ISEQS.store(!get_option!(pause), Ordering::Release);
     });
 
     if let Err(_) = result {
@@ -134,6 +144,15 @@ pub extern "C" fn rb_yjit_code_gc(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
     Qnil
 }
 
+#[no_mangle]
+pub extern "C" fn rb_yjit_resume(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+    if yjit_enabled_p() {
+        COMPILE_NEW_ISEQS.store(true, Ordering::Release);
+    }
+
+    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 {