YJIT: Instrument global allocations on stats build

yjit.rb

@@ -262,6 +262,7 @@ def _print_stats
       $stderr.puts "inline_code_size:      " + ("%10d" % stats[:inline_code_size])
       $stderr.puts "outlined_code_size:    " + ("%10d" % stats[:outlined_code_size])
       $stderr.puts "freed_code_size:       " + ("%10d" % stats[:freed_code_size])
+      $stderr.puts "yjit_alloc_size:       " + ("%10d" % stats[:yjit_alloc_size]) if stats.key?(:yjit_alloc_size)
       $stderr.puts "live_page_count:       " + ("%10d" % stats[:live_page_count])
       $stderr.puts "freed_page_count:      " + ("%10d" % stats[:freed_page_count])
       $stderr.puts "code_gc_count:         " + ("%10d" % stats[:code_gc_count])

yjit/Cargo.toml

@@ -16,12 +16,13 @@ crate-type = ["staticlib"]
 # No required dependencies to simplify build process. TODO: Link to yet to be
 # written rationale. Optional For development and testing purposes
 capstone = { version = "0.10.0", optional = true }
+stats_alloc = { version = "0.1.10", optional = true }
 
 [features]
 # NOTE: Development builds select a set of these via configure.ac
 # For debugging, `make V=1` shows exact cargo invocation.
 disasm = ["capstone"]
-stats = []
+stats = ["stats_alloc"]
 
 [profile.dev]
 opt-level = 0

yjit/src/stats.rs

@@ -8,6 +8,14 @@ use crate::cruby::*;
 use crate::options::*;
 use crate::yjit::yjit_enabled_p;
 
+// stats_alloc is a middleware to instrument global allocations in Rust.
+#[cfg(feature="stats")]
+#[global_allocator]
+static GLOBAL_ALLOCATOR: &stats_alloc::StatsAlloc<std::alloc::System> = &stats_alloc::INSTRUMENTED_SYSTEM;
+
+#[cfg(feature="stats")]
+static mut GLOBAL_REGION: Option<stats_alloc::Region<std::alloc::System>> = None;
+
 // YJIT exit counts for each instruction type
 const VM_INSTRUCTION_SIZE_USIZE:usize = VM_INSTRUCTION_SIZE as usize;
 static mut EXIT_OP_COUNT: [u64; VM_INSTRUCTION_SIZE_USIZE] = [0; VM_INSTRUCTION_SIZE_USIZE];
@@ -396,6 +404,12 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
 
         // Code GC count
         hash_aset_usize!(hash, "code_gc_count", CodegenGlobals::get_code_gc_count());
+
+        // Rust global allocations in bytes
+        #[cfg(feature="stats")]
+        if let Some(size) = global_allocation_size() {
+            hash_aset_usize!(hash, "yjit_alloc_size", size);
+        }
     }
 
     // If we're not generating stats, the hash is done
@@ -606,3 +620,21 @@ pub extern "C" fn rb_yjit_count_side_exit_op(exit_pc: *const VALUE) -> *const VA
     // This function must return exit_pc!
     return exit_pc;
 }
+
+// Enable instrumenting global allocations in Rust.
+pub fn init_global_allocator() {
+    #[cfg(feature="stats")]
+    unsafe { GLOBAL_REGION = Some(stats_alloc::Region::new(GLOBAL_ALLOCATOR)) };
+}
+
+// Get the size of global allocations in Rust.
+fn global_allocation_size() -> Option<usize> {
+    #[cfg(feature="stats")]
+    unsafe {
+        if let Some(region) = GLOBAL_REGION.as_mut() {
+            let stats = region.change();
+            return Some(stats.bytes_allocated.saturating_sub(stats.bytes_deallocated));
+        }
+    }
+    None
+}

yjit/src/yjit.rs

@@ -4,6 +4,7 @@ use crate::cruby::*;
 use crate::invariants::*;
 use crate::options::*;
 use crate::stats::YjitExitLocations;
+use crate::stats::init_global_allocator;
 
 use std::os::raw;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -48,6 +49,8 @@ pub extern "C" fn rb_yjit_init_rust() {
     // TODO: need to make sure that command-line options have been
     // initialized by CRuby
 
+    init_global_allocator();
+
     // Catch panics to avoid UB for unwinding into C frames.
     // See https://doc.rust-lang.org/nomicon/exception-safety.html
     // TODO: set a panic handler so the we don't print a message