Support options in `ZJIT.enable`

[chancancode]

Transferred from ruby#15219 (comment)

When implementing ruby#15219, support for options were left out on the original PR.

I still intend to do this at some point. (Don't let me block you though, if you'd like to tackle this before I do, go for it.)

The code is rather straightforward, roughly this is the entire patch:

diff --git a/zjit.c b/zjit.c
index 9560d88130..c175d7f1c0 100644
--- a/zjit.c
+++ b/zjit.c
@@ -311,7 +311,7 @@ rb_zjit_writebarrier_check_immediate(VALUE recv, VALUE val)
 }
 
 // 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_enable(rb_execution_context_t *ec, VALUE self, VALUE gen_stats, VALUE print_stats, VALUE mem_size, VALUE call_threshold);
 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);
diff --git a/zjit.rb b/zjit.rb
index e2aa55f764..8bc60f15f7 100644
--- a/zjit.rb
+++ b/zjit.rb
@@ -27,8 +27,16 @@ def enabled?
     Primitive.cexpr! 'RBOOL(rb_zjit_enabled_p)'
   end
 
-  # Enable ZJIT compilation.
-  def enable
+  # Enable ZJIT compilation. `stats` option decides whether to enable ZJIT stats or not.
+  # Optional `mem_size` and `call_threshold` can be provided to override default configuration.
+  #
+  # * `stats`:
+  #     * `false`: Don't enable stats.
+  #     * `true`: Enable stats. Print stats at exit.
+  #     * `:quiet`: Enable stats. Do not print stats at exit.
+  # * `mem_size`: Max amount of memory that ZJIT can use in MiB (default: 128).
+  # * `call_threshold`: Number of calls to trigger JIT (default: 30).
+  def enable(stats: false, mem_size: nil, call_threshold: nil)
     return false if enabled?
 
     if Primitive.cexpr! 'RBOOL(rb_yjit_enabled_p)'
@@ -36,7 +44,19 @@ def enable
       return false
     end
 
-    Primitive.rb_zjit_enable
+    if mem_size
+      raise ArgumentError, "mem_size must be a Integer" unless mem_size.is_a?(Integer)
+      raise ArgumentError, "mem_size must be between 1 and 2048 MB" unless (1..2048).include?(mem_size)
+    end
+
+    if call_threshold
+      raise ArgumentError, "call_threshold must be a Integer" unless call_threshold.is_a?(Integer)
+      raise ArgumentError, "call_threshold must be a positive integer" unless call_threshold.positive?
+    end
+
+    quiet = stats == :quiet
+    at_exit { print_stats } if stats && !quiet
+    Primitive.rb_zjit_enable(stats, !quiet, mem_size, call_threshold)
   end
 
   # Check if `--zjit-trace-exits` is used
diff --git a/zjit/src/options.rs b/zjit/src/options.rs
index 40b49146b7..3b2f06e28b 100644
--- a/zjit/src/options.rs
+++ b/zjit/src/options.rs
@@ -435,7 +435,7 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
 }
 
 /// Update rb_zjit_profile_threshold based on rb_zjit_call_threshold and options.num_profiles
-fn update_profile_threshold() {
+pub(crate) fn update_profile_threshold() {
     if unsafe { rb_zjit_call_threshold == 1 } {
         // If --zjit-call-threshold=1, never rewrite ISEQs to profile instructions.
         unsafe { rb_zjit_profile_threshold = 0; }
diff --git a/zjit/src/state.rs b/zjit/src/state.rs
index a807be3f12..53d5c43499 100644
--- a/zjit/src/state.rs
+++ b/zjit/src/state.rs
@@ -5,7 +5,7 @@ use crate::cruby::{self, rb_bug_panic_hook, rb_vm_insn_count, src_loc, EcPtr, Qn
 use crate::cruby_methods;
 use crate::invariants::Invariants;
 use crate::asm::CodeBlock;
-use crate::options::{get_option, rb_zjit_prepare_options};
+use crate::options::{OPTIONS, get_option, rb_zjit_call_threshold, rb_zjit_prepare_options, update_profile_threshold};
 use crate::stats::{Counters, InsnCounters, SideExitLocations};
 use crate::virtualmem::CodePtr;
 use std::collections::HashMap;
@@ -346,10 +346,33 @@ fn zjit_enable() {
 
 /// 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 {
+pub extern "C" fn rb_zjit_enable(_ec: EcPtr, _self: VALUE, gen_stats: VALUE, print_stats: VALUE, mem_size: VALUE, call_threshold: VALUE) -> VALUE {
     with_vm_lock(src_loc!(), || {
         // Options would not have been initialized during boot if no flags were specified
         rb_zjit_prepare_options();
+        let options = unsafe { OPTIONS.as_mut().unwrap() };
+
+        // Handle stats option
+        if gen_stats.test() {
+            options.stats = true;
+            options.print_stats = print_stats.test();
+        }
+
+        // Handle mem_size option
+        if !mem_size.nil_p() {
+            let mem_size_mb = mem_size.as_isize() >> 1;
+            let mem_size_bytes = mem_size_mb * 1024 * 1024;
+            options.mem_bytes = mem_size_bytes as usize;
+        }
+
+        // Handle call_threshold option
+        if !call_threshold.nil_p() {
+            let threshold = call_threshold.as_isize() >> 1;
+            unsafe {
+                rb_zjit_call_threshold = threshold as u64;
+            }
+            update_profile_threshold();
+        }
 
         // Initialize and enable ZJIT
         zjit_enable();

But while writing tests, I discovered more bugs in the existing implementation, including in YJIT, that warrants reconsidering how we approach enabling the JIT.

I enumerated some in #883. I think those aren't necessarily a hard blocker, but is related to the architecture issue.

Another issue I found is that we sometimes register the exit hooks twice. For example, on 3.4.7, --yjit-disable --yjit-stats prints the stats twice at exit:

• Once here: https://github.com/ruby/ruby/blob/3c4cda10eb80fa38f2b27d4774e3e3e82bf190c9/yjit.rb#L261-L264
• Then here again: https://github.com/ruby/ruby/blob/3c4cda10eb80fa38f2b27d4774e3e3e82bf190c9/yjit.rb#L66

Both JITs were originally developed with the "either enable at boot or not" model, with the lazy enable branch added later. It seems like some of these paths didn't get properly reconciled, and the same patterns get carried over to ZJIT.

The current state of ZJIT as of 4.0.0 is that you can't override the options lazily, but that at least gives you somewhat consistent/coherent results for exit hooks, etc.

Personally, I'd like to reconcile the two enable paths in ZJIT (so that enabling at boot behaves much more similar to calling ZJIT.enable as the first statement) before adding the overrides to ZJIT, but that's just my personal preference.