ZJIT: Add --zjit-exec-mem-size (#14175)

* ZJIT: Add --zjit-exec-mem-size

* Add a comment about the limit

3 files changed, 38 insertions, 12 deletions

diff --git a/zjit/src/options.rs b/zjit/src/options.rs
index 92f56b8916..07584c9b99 100644
--- a/zjit/src/options.rs
+++ b/zjit/src/options.rs
@@ -22,6 +22,10 @@ pub static mut OPTIONS: Option<Options> = None;
 
 #[derive(Clone, Debug)]
 pub struct Options {
+    /// Hard limit of the executable memory block to allocate in bytes.
+    /// Note that the command line argument is expressed in MiB and not bytes.
+    pub exec_mem_bytes: usize,
+
     /// Number of times YARV instructions should be profiled.
     pub num_profiles: u8,
 
@@ -58,6 +62,7 @@ pub struct Options {
 impl Default for Options {
     fn default() -> Self {
         Options {
+            exec_mem_bytes: 64 * 1024 * 1024,
             num_profiles: 1,
             stats: false,
             debug: false,
@@ -74,12 +79,18 @@ impl Default for Options {
 }
 
 /// `ruby --help` descriptions for user-facing options. Do not add options for ZJIT developers.
-/// Note that --help allows only 80 chars per line, including indentation.    80-char limit --> |
+/// Note that --help allows only 80 chars per line, including indentation, and it also puts the
+/// description in a separate line if the option name is too long.  80-char limit --> | (any character beyond this `|` column fails the test)
 pub const ZJIT_OPTIONS: &'static [(&str, &str)] = &[
-    ("--zjit-call-threshold=num", "Number of calls to trigger JIT (default: 2)."),
-    ("--zjit-num-profiles=num",   "Number of profiled calls before JIT (default: 1, max: 255)."),
-    ("--zjit-stats",              "Enable collecting ZJIT statistics."),
-    ("--zjit-perf",               "Dump ISEQ symbols into /tmp/perf-{}.map for Linux perf."),
+    // TODO: Hide --zjit-exec-mem-size from ZJIT_OPTIONS once we add --zjit-mem-size (Shopify/ruby#686)
+    ("--zjit-exec-mem-size=num",
+                     "Size of executable memory block in MiB (default: 64)."),
+    ("--zjit-call-threshold=num",
+                     "Number of calls to trigger JIT (default: 2)."),
+    ("--zjit-num-profiles=num",
+                     "Number of profiled calls before JIT (default: 1, max: 255)."),
+    ("--zjit-stats", "Enable collecting ZJIT statistics."),
+    ("--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."),
 ];
@@ -163,6 +174,20 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
     match (opt_name, opt_val) {
         ("", "") => {}, // Simply --zjit
 
+        ("mem-size", _) => match opt_val.parse::<usize>() {
+            Ok(n) => {
+                // Reject 0 or too large values that could overflow.
+                // The upper bound is 1 TiB but we could make it smaller.
+                if n == 0 || n > 1024 * 1024 {
+                    return None
+                }
+
+                // Convert from MiB to bytes internally for convenience
+                options.exec_mem_bytes = n * 1024 * 1024;
+            }
+            Err(_) => return None,
+        },
+
         ("call-threshold", _) => match opt_val.parse() {
             Ok(n) => {
                 unsafe { rb_zjit_call_threshold = n; }
diff --git a/zjit/src/state.rs b/zjit/src/state.rs
index 79be91fd85..dca04b7a72 100644
--- a/zjit/src/state.rs
+++ b/zjit/src/state.rs
@@ -48,7 +48,7 @@ impl ZJITState {
             use crate::cruby::*;
             use crate::options::*;
 
-            let exec_mem_size: usize = 64 * 1024 * 1024; // TODO: implement the option
+            let exec_mem_bytes: usize = get_option!(exec_mem_bytes);
             let virt_block: *mut u8 = unsafe { rb_zjit_reserve_addr_space(64 * 1024 * 1024) };
 
             // Memory protection syscalls need page-aligned addresses, so check it here. Assuming
@@ -73,8 +73,8 @@ impl ZJITState {
                 crate::virtualmem::sys::SystemAllocator {},
                 page_size,
                 NonNull::new(virt_block).unwrap(),
-                exec_mem_size,
-                64 * 1024 * 1024, // TODO: support the option
+                exec_mem_bytes,
+                exec_mem_bytes, // TODO: change this to --zjit-mem-size (Shopify/ruby#686)
             );
             let mem_block = Rc::new(RefCell::new(mem_block));
 
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
index 5b39ecdf4b..fa8b741eea 100644
--- a/zjit/src/stats.rs
+++ b/zjit/src/stats.rs
@@ -70,10 +70,6 @@ fn incr_counter(counter: Counter, amount: u64) {
     unsafe { *ptr += amount; }
 }
 
-pub fn zjit_alloc_size() -> usize {
-    0 // TODO: report the actual memory usage
-}
-
 /// Return a Hash object that contains ZJIT statistics
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE) -> VALUE {
@@ -113,3 +109,8 @@ pub fn with_time_stat<F, R>(counter: Counter, func: F) -> R where F: FnOnce() ->
     incr_counter(counter, nanos as u64);
     ret
 }
+
+/// The number of bytes ZJIT has allocated on the Rust heap.
+pub fn zjit_alloc_size() -> usize {
+    0 // TODO: report the actual memory usage to support --zjit-mem-size (Shopify/ruby#686)
+}