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

test/ruby/test_yjit.rb

@@ -51,6 +51,29 @@ def test_command_line_switches
     #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,10 @@ def self.reset_stats!
     Primitive.rb_yjit_reset_stats_bang
   end
 
+  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 {