YJIT: Make yjit_alloc_size available by default (#8426)

3 files changed, 37 insertions, 22 deletions

diff --git a/yjit/Cargo.lock b/yjit/Cargo.lock
index 08ae1bd487..e9a59cb771 100644
--- a/yjit/Cargo.lock
+++ b/yjit/Cargo.lock
@@ -35,15 +35,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "21a41fed9d98f27ab1c6d161da622a4fa35e8a54a8adc24bbf3ddd0ef70b0e50"
 
 [[package]]
-name = "stats_alloc"
-version = "0.1.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5c0e04424e733e69714ca1bbb9204c1a57f09f5493439520f9f68c132ad25eec"
-
-[[package]]
 name = "yjit"
 version = "0.1.0"
 dependencies = [
  "capstone",
- "stats_alloc",
 ]
diff --git a/yjit/Cargo.toml b/yjit/Cargo.toml
index 8cd593e4ee..2194402cdd 100644
--- a/yjit/Cargo.toml
+++ b/yjit/Cargo.toml
@@ -16,13 +16,12 @@ 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_alloc"]
+stats = []
 
 [profile.dev]
 opt-level = 0
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index e52611ed67..7eae0cba21 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -3,6 +3,8 @@
 
 #![allow(dead_code)] // Counters are only used with the stats features
 
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::time::Instant;
 
 use crate::codegen::CodegenGlobals;
@@ -12,13 +14,42 @@ 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")]
+/// A middleware to count Rust-allocated bytes as yjit_alloc_size.
 #[global_allocator]
-static GLOBAL_ALLOCATOR: &stats_alloc::StatsAlloc<std::alloc::System> = &stats_alloc::INSTRUMENTED_SYSTEM;
+static GLOBAL_ALLOCATOR: StatsAlloc = StatsAlloc { alloc_size: AtomicUsize::new(0) };
+
+pub struct StatsAlloc {
+    alloc_size: AtomicUsize,
+}
+
+unsafe impl GlobalAlloc for StatsAlloc {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        self.alloc_size.fetch_add(layout.size(), Ordering::SeqCst);
+        System.alloc(layout)
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        self.alloc_size.fetch_sub(layout.size(), Ordering::SeqCst);
+        System.dealloc(ptr, layout)
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        self.alloc_size.fetch_add(layout.size(), Ordering::SeqCst);
+        System.alloc_zeroed(layout)
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        if new_size > layout.size() {
+            self.alloc_size.fetch_add(new_size - layout.size(), Ordering::SeqCst);
+        } else if new_size < layout.size() {
+            self.alloc_size.fetch_sub(layout.size() - new_size, Ordering::SeqCst);
+        }
+        System.realloc(ptr, layout, new_size)
+    }
+}
 
 // YJIT exit counts for each instruction type
-const VM_INSTRUCTION_SIZE_USIZE:usize = VM_INSTRUCTION_SIZE as usize;
+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];
 
 /// Global state needed for collecting backtraces of exits
@@ -592,8 +623,7 @@ fn rb_yjit_gen_stats_dict(context: bool) -> VALUE {
         hash_aset_usize!(hash, "code_region_size", cb.mapped_region_size());
 
         // Rust global allocations in bytes
-        #[cfg(feature="stats")]
-        hash_aset_usize!(hash, "yjit_alloc_size", global_allocation_size());
+        hash_aset_usize!(hash, "yjit_alloc_size", GLOBAL_ALLOCATOR.alloc_size.load(Ordering::SeqCst));
 
         // `context` is true at RubyVM::YJIT._print_stats for --yjit-stats. It's false by default
         // for RubyVM::YJIT.runtime_stats because counting all Contexts could be expensive.
@@ -837,13 +867,6 @@ pub extern "C" fn rb_yjit_count_side_exit_op(exit_pc: *const VALUE) -> *const VA
     return exit_pc;
 }
 
-// Get the size of global allocations in Rust.
-#[cfg(feature="stats")]
-fn global_allocation_size() -> usize {
-    let stats = GLOBAL_ALLOCATOR.stats();
-    stats.bytes_allocated.saturating_sub(stats.bytes_deallocated)
-}
-
 /// Measure the time taken by func() and add that to yjit_compile_time.
 pub fn with_compile_time<F, R>(func: F) -> R where F: FnOnce() -> R {
     let start = Instant::now();