ZJIT: add support for lazy RubyVM::ZJIT.enable

This implements Shopify#854:

- Splits boot-time and enable-time initialization,
  tracks progress with `InitializationState` enum

- Introduces `RubyVM::ZJIT.enable` Ruby method for
  enabling the JIT lazily, if not already enabled

- Introduces `--zjit-disable` flag, which can be
  used alongside the other `--zjit-*` flags but
  prevents enabling the JIT at boot time

- Adds ZJIT infra to support JIT hooks, but this
  is not currently exercised (Shopify/ruby#667)

Left for future enhancements:

- Support kwargs for overriding the CLI flags in
  `RubyVM::ZJIT.enable`

Closes Shopify#854

jit_hook.rb

@@ -3,9 +3,8 @@ class Module
   # This method is removed in jit_undef.rb.
   private def with_jit(&block) # :nodoc:
     # ZJIT currently doesn't compile Array#each properly, so it's disabled for now.
-    if defined?(RubyVM::ZJIT) && Primitive.rb_zjit_option_enabled_p && false # TODO: remove `&& false` (Shopify/ruby#667)
-      # We don't support lazily enabling ZJIT yet, so we can call the block right away.
-      block.call
+    if defined?(RubyVM::ZJIT) && false # TODO: remove `&& false` (Shopify/ruby#667)
+      RubyVM::ZJIT.send(:add_jit_hook, block)
     elsif defined?(RubyVM::YJIT)
       RubyVM::YJIT.send(:add_jit_hook, block)
     end

ruby.c

@@ -1842,10 +1842,8 @@ ruby_opt_init(ruby_cmdline_options_t *opt)
     rb_yjit_init(opt->yjit);
 #endif
 #if USE_ZJIT
-    if (opt->zjit) {
-        extern void rb_zjit_init(void);
-        rb_zjit_init();
-    }
+    extern void rb_zjit_init(bool);
+    rb_zjit_init(opt->zjit);
 #endif
 
     ruby_set_script_name(opt->script_name);
@@ -2368,6 +2366,12 @@ process_options(int argc, char **argv, ruby_cmdline_options_t *opt)
 #if USE_ZJIT
         if (!FEATURE_USED_P(opt->features, zjit) && env_var_truthy("RUBY_ZJIT_ENABLE")) {
             FEATURE_SET(opt->features, FEATURE_BIT(zjit));
+
+            // When the --zjit flag is specified, we would have call setup_zjit_options(""),
+            // which would have called rb_zjit_prepare_options() internally. This ensures we
+            // go through the same set up but with less overhead than setup_zjit_options("").
+            extern void rb_zjit_prepare_options();
+            rb_zjit_prepare_options();
         }
 #endif
     }
@@ -2383,10 +2387,9 @@ process_options(int argc, char **argv, ruby_cmdline_options_t *opt)
     }
 #endif
 #if USE_ZJIT
-    if (FEATURE_SET_P(opt->features, zjit) && !opt->zjit) {
-        extern void rb_zjit_prepare_options(void);
-        rb_zjit_prepare_options();
-        opt->zjit = true;
+    if (FEATURE_SET_P(opt->features, zjit)) {
+        bool rb_zjit_option_enable(void);
+        opt->zjit = rb_zjit_option_enable(); // set opt->zjit for Init_ruby_description() and calling rb_zjit_init()
     }
 #endif
 

test/ruby/test_zjit.rb

@@ -59,6 +59,44 @@ class TestZJIT < Test::Unit::TestCase
     end
   end
 
+  def test_zjit_enable
+    assert_separately([], <<~'RUBY')
+      refute_predicate RubyVM::ZJIT, :enabled?
+      refute_predicate RubyVM::ZJIT, :stats_enabled?
+      refute_includes RUBY_DESCRIPTION, "+ZJIT"
+
+      RubyVM::ZJIT.enable
+
+      assert_predicate RubyVM::ZJIT, :enabled?
+      refute_predicate RubyVM::ZJIT, :stats_enabled?
+      assert_includes RUBY_DESCRIPTION, "+ZJIT"
+    RUBY
+  end
+
+  def test_zjit_disable
+    assert_separately(["--zjit", "--zjit-disable"], <<~'RUBY')
+      refute_predicate RubyVM::ZJIT, :enabled?
+      refute_includes RUBY_DESCRIPTION, "+ZJIT"
+
+      RubyVM::ZJIT.enable
+
+      assert_predicate RubyVM::ZJIT, :enabled?
+      assert_includes RUBY_DESCRIPTION, "+ZJIT"
+    RUBY
+  end
+
+  def test_zjit_enable_respects_existing_options
+    assert_separately(['--zjit-disable', '--zjit-stats=quiet'], <<~RUBY)
+      refute_predicate RubyVM::ZJIT, :enabled?
+      assert_predicate RubyVM::ZJIT, :stats_enabled?
+
+      RubyVM::ZJIT.enable
+
+      assert_predicate RubyVM::ZJIT, :enabled?
+      assert_predicate RubyVM::ZJIT, :stats_enabled?
+    RUBY
+  end
+
   def test_call_itself
     assert_compiles '42', <<~RUBY, call_threshold: 2
       def test = 42.itself

version.c

@@ -276,6 +276,15 @@ ruby_set_yjit_description(void)
     define_ruby_description(YJIT_DESCRIPTION);
 }
 
+void
+ruby_set_zjit_description(void)
+{
+    VALUE mRuby = rb_path2class("Ruby");
+    rb_const_remove(rb_cObject, rb_intern("RUBY_DESCRIPTION"));
+    rb_const_remove(mRuby, rb_intern("DESCRIPTION"));
+    define_ruby_description(ZJIT_DESCRIPTION);
+}
+
 void
 ruby_show_version(void)
 {

zjit.c

@@ -305,6 +305,7 @@ rb_zjit_class_has_default_allocator(VALUE klass)
 VALUE rb_vm_get_untagged_block_handler(rb_control_frame_t *reg_cfp);
 
 // Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them.
+VALUE rb_zjit_enable(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self, VALUE target_key);
 VALUE rb_zjit_reset_stats_bang(rb_execution_context_t *ec, VALUE self);

zjit.rb

@@ -7,6 +7,8 @@
 # This module may not exist if ZJIT does not support the particular platform
 # for which CRuby is built.
 module RubyVM::ZJIT
+  # Blocks that are called when YJIT is enabled
+  @jit_hooks = []
   # Avoid calling a Ruby method here to avoid interfering with compilation tests
   if Primitive.rb_zjit_print_stats_p
     at_exit { print_stats }
@@ -22,6 +24,18 @@ class << RubyVM::ZJIT
     Primitive.cexpr! 'RBOOL(rb_zjit_enabled_p)'
   end
 
+  # Enable ZJIT compilation.
+  def enable
+    return false if enabled?
+
+    if Primitive.cexpr! 'RBOOL(rb_yjit_enabled_p)'
+      warn("Only one JIT can be enabled at the same time.")
+      return false
+    end
+
+    Primitive.rb_zjit_enable
+  end
+
   # Check if `--zjit-trace-exits` is used
   def trace_exit_locations_enabled?
     Primitive.rb_zjit_trace_exit_locations_enabled_p
@@ -234,6 +248,17 @@ class << RubyVM::ZJIT
   # :stopdoc:
   private
 
+  # Register a block to be called when ZJIT is enabled
+  def add_jit_hook(hook)
+    @jit_hooks << hook
+  end
+
+  # Run ZJIT hooks registered by `#with_jit`
+  def call_jit_hooks
+    @jit_hooks.each(&:call)
+    @jit_hooks.clear
+  end
+
   def print_counters(keys, buf:, stats:, right_align: false, base: nil)
     key_pad = keys.map { |key| key.to_s.sub(/_time_ns\z/, '_time').size }.max + 1
     key_align = '-' unless right_align

zjit/src/options.rs

@@ -45,7 +45,7 @@ pub struct Options {
     /// Number of times YARV instructions should be profiled.
     pub num_profiles: NumProfiles,
 
-    /// Enable YJIT statsitics
+    /// Enable ZJIT statistics
     pub stats: bool,
 
     /// Print stats on exit (when stats is also true)
@@ -54,6 +54,10 @@ pub struct Options {
     /// Enable debug logging
     pub debug: bool,
 
+    // Whether to enable JIT at boot. This option prevents other
+    // ZJIT tuning options from enabling ZJIT at boot.
+    pub disable: bool,
+
     /// Turn off the HIR optimizer
     pub disable_hir_opt: bool,
 
@@ -97,6 +101,7 @@ impl Default for Options {
             stats: false,
             print_stats: false,
             debug: false,
+            disable: false,
             disable_hir_opt: false,
             dump_hir_init: None,
             dump_hir_opt: None,
@@ -123,6 +128,8 @@ pub const ZJIT_OPTIONS: &[(&str, &str)] = &[
     ("--zjit-num-profiles=num",
                      "Number of profiled calls before JIT (default: 5)."),
     ("--zjit-stats[=quiet]", "Enable collecting ZJIT statistics (=quiet to suppress output)."),
+    ("--zjit-disable",
+                     "Disable ZJIT for lazily enabling it with RubyVM::ZJIT.enable."),
     ("--zjit-perf",  "Dump ISEQ symbols into /tmp/perf-{}.map for Linux perf."),
     ("--zjit-log-compiled-iseqs=path",
                      "Log compiled ISEQs to the file. The file will be truncated."),
@@ -175,7 +182,7 @@ const DUMP_LIR_ALL: &[DumpLIR] = &[
     DumpLIR::scratch_split,
 ];
 
-/// Mamximum value for --zjit-mem-size/--zjit-exec-mem-size in MiB.
+/// Maximum value for --zjit-mem-size/--zjit-exec-mem-size in MiB.
 /// We set 1TiB just to avoid overflow. We could make it smaller.
 const MAX_MEM_MIB: usize = 1024 * 1024;
 
@@ -319,6 +326,8 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
 
         ("debug", "") => options.debug = true,
 
+        ("disable", "") => options.disable = true,
+
         ("disable-hir-opt", "") => options.disable_hir_opt = true,
 
         // --zjit-dump-hir dumps the actual input to the codegen, which is currently the same as --zjit-dump-hir-opt.
@@ -442,15 +451,13 @@ macro_rules! debug {
 }
 pub(crate) use debug;
 
-/// Return Qtrue if --zjit* has been specified. For the `#with_jit` hook,
-/// this becomes Qtrue before ZJIT is actually initialized and enabled.
+/// Return true if ZJIT should be enabled at boot.
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_option_enabled_p(_ec: EcPtr, _self: VALUE) -> VALUE {
-    // If any --zjit* option is specified, OPTIONS becomes Some.
-    if unsafe { OPTIONS.is_some() } {
-        Qtrue
+pub extern "C" fn rb_zjit_option_enable() -> bool {
+    if unsafe { OPTIONS.as_ref() }.is_some_and(|opts| !opts.disable) {
+        true
     } else {
-        Qfalse
+        false
     }
 }
 

zjit/src/state.rs

@@ -1,11 +1,11 @@
 //! Runtime state of ZJIT.
 
 use crate::codegen::{gen_entry_trampoline, gen_exit_trampoline, gen_exit_trampoline_with_counter, gen_function_stub_hit_trampoline};
-use crate::cruby::{self, rb_bug_panic_hook, rb_vm_insn_count, EcPtr, Qnil, rb_vm_insn_addr2opcode, rb_profile_frames, VALUE, VM_INSTRUCTION_SIZE, size_t, rb_gc_mark};
+use crate::cruby::{self, rb_bug_panic_hook, rb_vm_insn_count, src_loc, EcPtr, Qnil, Qtrue, rb_vm_insn_addr2opcode, rb_profile_frames, VALUE, VM_INSTRUCTION_SIZE, size_t, rb_gc_mark, with_vm_lock};
 use crate::cruby_methods;
 use crate::invariants::Invariants;
 use crate::asm::CodeBlock;
-use crate::options::get_option;
+use crate::options::{get_option, rb_zjit_prepare_options};
 use crate::stats::{Counters, InsnCounters, SideExitLocations};
 use crate::virtualmem::CodePtr;
 use std::collections::HashMap;
@@ -63,12 +63,42 @@ pub struct ZJITState {
     exit_locations: Option<SideExitLocations>,
 }
 
+/// Tracks the initialization progress
+enum InitializationState {
+    Uninitialized,
+
+    /// At boot time, rb_zjit_init will be called regardless of whether
+    /// ZJIT is enabled, in this phase we initialize any states that must
+    /// be captured at during boot.
+    Initialized(cruby_methods::Annotations),
+
+    /// When ZJIT is enabled, either during boot with `--zjit`, or lazily
+    /// at a later time with `RubyVM::ZJIT.enable`, we perform the rest
+    /// of the initialization steps and produce the `ZJITState` instance.
+    Enabled(ZJITState),
+
+    /// Indicates that ZJITState::init has panicked. Should never be
+    /// encountered in practice since we abort immediately when that
+    /// happens.
+    Panicked,
+}
+
 /// Private singleton instance of the codegen globals
-static mut ZJIT_STATE: Option<ZJITState> = None;
+static mut ZJIT_STATE: InitializationState = InitializationState::Uninitialized;
 
 impl ZJITState {
     /// Initialize the ZJIT globals. Return the address of the JIT entry trampoline.
     pub fn init() -> *const u8 {
+        use InitializationState::*;
+
+        let initialization_state = unsafe {
+            std::mem::replace(&mut ZJIT_STATE, Panicked)
+        };
+
+        let Initialized(method_annotations) = initialization_state else {
+            panic!("rb_zjit_init was never called");
+        };
+
         let mut cb = {
             use crate::options::*;
             use crate::virtualmem::*;
@@ -99,7 +129,7 @@ impl ZJITState {
             send_fallback_counters: [0; VM_INSTRUCTION_SIZE as usize],
             invariants: Invariants::default(),
             assert_compiles: false,
-            method_annotations: cruby_methods::init(),
+            method_annotations,
             exit_trampoline,
             function_stub_hit_trampoline,
             exit_trampoline_with_counter: exit_trampoline,
@@ -107,7 +137,7 @@ impl ZJITState {
             not_annotated_frame_cfunc_counter_pointers: HashMap::new(),
             exit_locations,
         };
-        unsafe { ZJIT_STATE = Some(zjit_state); }
+        unsafe { ZJIT_STATE = Enabled(zjit_state); }
 
         // With --zjit-stats, use a different trampoline on function stub exits
         // to count exit_compilation_failure. Note that the trampoline code depends
@@ -123,12 +153,16 @@ impl ZJITState {
 
     /// Return true if zjit_state has been initialized
     pub fn has_instance() -> bool {
-        unsafe { ZJIT_STATE.as_mut().is_some() }
+        matches!(unsafe { &ZJIT_STATE }, InitializationState::Enabled(_))
     }
 
     /// Get a mutable reference to the codegen globals instance
     fn get_instance() -> &'static mut ZJITState {
-        unsafe { ZJIT_STATE.as_mut().unwrap() }
+        if let InitializationState::Enabled(instance) = unsafe { &mut ZJIT_STATE } {
+            instance
+        } else {
+            panic!("ZJITState::get_instance called when ZJIT is not enabled")
+        }
     }
 
     /// Get a mutable reference to the inline code block
@@ -249,14 +283,39 @@ impl ZJITState {
     }
 }
 
-/// Initialize ZJIT
+/// Initialize ZJIT at boot. This is called even if ZJIT is disabled.
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_init() {
+pub extern "C" fn rb_zjit_init(zjit_enabled: bool) {
+    use InitializationState::*;
+
+    debug_assert!(
+        matches!(unsafe { &ZJIT_STATE }, Uninitialized),
+        "rb_zjit_init should only be called once during boot",
+    );
+
+    // Initialize IDs and method annotations.
+    // cruby_methods::init() must be called at boot,
+    // as cmes could have been re-defined after boot.
+    cruby::ids::init();
+
+    let method_annotations = cruby_methods::init();
+
+    unsafe { ZJIT_STATE = Initialized(method_annotations); }
+
+    // If --zjit, enable ZJIT immediately
+    if zjit_enabled {
+        zjit_enable();
+    }
+}
+
+/// Enable ZJIT compilation.
+fn zjit_enable() {
+    // TODO: call RubyVM::ZJIT::call_jit_hooks here
+
     // Catch panics to avoid UB for unwinding into C frames.
     // See https://doc.rust-lang.org/nomicon/exception-safety.html
     let result = std::panic::catch_unwind(|| {
         // Initialize ZJIT states
-        cruby::ids::init();
         let zjit_entry = ZJITState::init();
 
         // Install a panic hook for ZJIT
@@ -271,11 +330,33 @@ pub extern "C" fn rb_zjit_init() {
     });
 
     if result.is_err() {
-        println!("ZJIT: zjit_init() panicked. Aborting.");
+        println!("ZJIT: zjit_enable() panicked. Aborting.");
         std::process::abort();
     }
 }
 
+/// Enable ZJIT compilation, returning Qtrue if ZJIT was previously disabled
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_enable(_ec: EcPtr, _self: VALUE) -> VALUE {
+    with_vm_lock(src_loc!(), || {
+        // Options would not have been initialized during boot if no flags were specified
+        rb_zjit_prepare_options();
+
+        // Initialize and enable ZJIT
+        zjit_enable();
+
+        // Add "+ZJIT" to RUBY_DESCRIPTION
+        unsafe {
+            unsafe extern "C" {
+                fn ruby_set_zjit_description();
+            }
+            ruby_set_zjit_description();
+        }
+
+        Qtrue
+    })
+}
+
 /// Assert that any future ZJIT compilation will return a function pointer (not fail to compile)
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_assert_compiles(_ec: EcPtr, _self: VALUE) -> VALUE {