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 <[email protected]>

---------

Co-authored-by: Alan Wu <[email protected]>
Co-authored-by: Takashi Kokubun <[email protected]>

2 files changed, 28 insertions, 0 deletions

diff --git a/yjit/src/options.rs b/yjit/src/options.rs
index 9a8205d933..dfae06d1e7 100644
--- a/yjit/src/options.rs
+++ b/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 => {
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index e133049ee8..10bddcc600 100644
--- a/yjit/src/yjit.rs
+++ b/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 {