examples/check_real_memory: use a hand written buddy allocator instead of jemalloc (#500)

* examples/check_real_memory: use a hand written buddy allocator instead of jemalloc

* Use pastey instead of paste

* Use buddy-alloc

Author: mohanson

.github/workflows/develop.yml

@@ -11,7 +11,7 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       # https://github.com/EmbarkStudios/cargo-deny-action
-      - uses: EmbarkStudios/cargo-deny-action@v1
+      - uses: EmbarkStudios/cargo-deny-action@v2
         with:
           arguments: --all-features
           command: check advisories licenses sources bans

Cargo.toml

@@ -40,10 +40,7 @@ cc = "1.0"
 argparse = "0.2"
 criterion = "0.5.1"
 proptest = "1.5.0"
-
-[target.'cfg(not(target_os = "windows"))'.dev-dependencies]
-jemallocator = "0.5.0"
-jemalloc-ctl = "0.5.0"
+buddy-alloc = "0.6.0"
 
 [[bench]]
 name = "bits_benchmark"

definitions/Cargo.toml

@@ -18,4 +18,4 @@ name = "generate_asm_constants"
 path = "src/generate_asm_constants.rs"
 
 [dependencies]
-paste = "1.0.12"
+paste = { package = "pastey", version = "0.2.1" }

examples/check_real_memory.rs

@@ -1,8 +1,6 @@
 // This example is mainly to test whether there is memory overflow.
-// Under linux, we choose to use smem, which can monitor memory changes more accurately
 
-use ckb_vm::{Bytes, FlatMemory, SparseMemory, run_with_memory};
-use std::process::{Command, id};
+use ckb_vm::{Bytes, DEFAULT_MEMORY_SIZE, SparseMemory, run_with_memory};
 
 #[cfg(has_asm)]
 use ckb_vm::{
@@ -13,162 +11,97 @@ use ckb_vm::{
     },
 };
 
-#[cfg(has_asm)]
-use std::thread;
+use buddy_alloc::{BuddyAllocParam, buddy_alloc::BuddyAlloc};
+use std::alloc::GlobalAlloc;
+use std::alloc::Layout;
+use std::cell::RefCell;
 
-static BIN_PATH_BUFFER: &'static [u8] = include_bytes!("../tests/programs/alloc_many");
-static BIN_NAME: &str = "alloc_many";
-
-#[cfg(not(target_os = "windows"))]
-#[global_allocator]
-static GLOBAL: jemallocator::Jemalloc = jemallocator::Jemalloc;
+const HEAP_SIZE: usize = 16 * 1024 * 1024;
+const LEAF_SIZE: usize = 64;
+#[repr(align(64))]
+struct Heap<const S: usize>([u8; S]);
+static mut HEAP: Heap<HEAP_SIZE> = Heap([0u8; HEAP_SIZE]);
 
-static G_CHECK_LOOP: usize = 10;
-
-fn get_current_memory_linux() -> usize {
-    let output = String::from_utf8(
-        Command::new("smem")
-            .arg("-P")
-            .arg("check_real_memory")   // current process name
-            .arg("-c")
-            .arg("pid uss")
-            .output()
-            .expect("run ps failed")
-            .stdout,
-    )
-    .unwrap();
+pub struct NonThreadsafeAlloc {
+    buddy_alloc_param: BuddyAllocParam,
+    inner_buddy_alloc: RefCell<Option<BuddyAlloc>>,
+}
 
-    let outputs = output.split("\n").collect::<Vec<&str>>();
-    for i in 1..outputs.len() {
-        let mut has_pid = false;
-        let mut memory_size: u32 = 0;
-        for d in outputs[i].split(" ").collect::<Vec<&str>>() {
-            if d == " " || d == "" {
-                continue;
-            }
-            let val: u32 = d.parse().unwrap();
-            if !has_pid {
-                if val != id() {
-                    continue;
-                }
-                has_pid = true;
-            } else {
-                memory_size = val;
-            }
-        }
-        if memory_size != 0 {
-            return memory_size as usize;
+impl NonThreadsafeAlloc {
+    /// see BuddyAlloc::new
+    pub const fn new(buddy_alloc_param: BuddyAllocParam) -> Self {
+        NonThreadsafeAlloc {
+            inner_buddy_alloc: RefCell::new(None),
+            buddy_alloc_param,
         }
     }
 
-    0
-}
-
-fn get_current_memory() -> usize {
-    if !cfg!(target_os = "linux") {
-        get_current_memory_linux()
-    } else {
-        let pid = format!("{}", id());
-        let output = String::from_utf8(
-            Command::new("ps")
-                .arg("-p")
-                .arg(pid)
-                .arg("-o")
-                .arg("rss")
-                .output()
-                .expect("run ps failed")
-                .stdout,
-        )
-        .unwrap();
-
-        let output = output.split("\n").collect::<Vec<&str>>();
+    unsafe fn with_buddy_alloc<R, F: FnOnce(&mut BuddyAlloc) -> R>(&self, f: F) -> R {
+        let mut inner = self.inner_buddy_alloc.borrow_mut();
+        let alloc = inner.get_or_insert_with(|| unsafe { BuddyAlloc::new(self.buddy_alloc_param) });
+        f(alloc)
+    }
 
-        let memory_size = output[1].replace(" ", "");
-        memory_size.parse().unwrap()
+    fn used(&self) -> usize {
+        HEAP_SIZE
+            - self
+                .inner_buddy_alloc
+                .borrow()
+                .as_ref()
+                .unwrap()
+                .available_bytes()
     }
 }
 
-struct MemoryOverflow {
-    #[cfg(not(target_os = "windows"))]
-    base_allocated: usize,
+unsafe impl GlobalAlloc for NonThreadsafeAlloc {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let bytes = layout.size();
+        unsafe { self.with_buddy_alloc(|alloc| alloc.malloc(bytes)) }
+    }
 
-    #[cfg(not(target_os = "windows"))]
-    base_resident: usize,
+    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
+        unsafe { self.with_buddy_alloc(|alloc| alloc.free(ptr)) };
+    }
 }
 
-impl MemoryOverflow {
-    pub fn new() -> Self {
-        Self {
-            #[cfg(not(target_os = "windows"))]
-            base_allocated: jemalloc_ctl::stats::allocated::read().unwrap(),
-
-            #[cfg(not(target_os = "windows"))]
-            base_resident: jemalloc_ctl::stats::resident::read().unwrap(),
-        }
-    }
+unsafe impl Sync for NonThreadsafeAlloc {}
 
-    pub fn check(&self) {
-        #[cfg(not(target_os = "windows"))]
-        assert!(jemalloc_ctl::stats::allocated::read().unwrap() <= self.base_allocated);
+#[allow(static_mut_refs)]
+#[global_allocator]
+static ALLOC: NonThreadsafeAlloc =
+    unsafe { NonThreadsafeAlloc::new(BuddyAllocParam::new(HEAP.0.as_ptr(), HEAP_SIZE, LEAF_SIZE)) };
 
-        #[cfg(not(target_os = "windows"))]
-        assert!(jemalloc_ctl::stats::resident::read().unwrap() <= self.base_resident);
-    }
-}
+static BIN_PATH_BUFFER: &'static [u8] = include_bytes!("../tests/programs/alloc_many");
+static BIN_NAME: &str = "alloc_many";
+static G_CHECK_LOOP: usize = 10;
 
-fn check_interpreter(memory_size: usize) -> Result<(), ()> {
-    println!(
-        "Check interpreter memory overflow, ckb-vm memory size: {}",
-        memory_size
-    );
-    println!("Base memory: {}", get_current_memory());
+fn check_interpreter() -> Result<(), Box<dyn std::error::Error>> {
+    println!("Check interpreter: init");
+    println!("Check interpreter: base memory used {}", ALLOC.used());
     for _ in 0..G_CHECK_LOOP {
         let result = run_with_memory::<u64, SparseMemory<u64>>(
             &Bytes::from(BIN_PATH_BUFFER),
             &vec![Bytes::from(BIN_NAME)],
-            memory_size,
+            DEFAULT_MEMORY_SIZE,
         );
         assert!(result.is_ok());
         assert_eq!(result.unwrap(), 0);
-        println!("Current memory: {}", get_current_memory());
+        println!("Check interpreter: step memory used {}", ALLOC.used());
     }
-    println!("End of check");
-    Ok(())
-}
-
-fn check_falt(memory_size: usize) -> Result<(), ()> {
-    println!(
-        "Check falt memory overflow, ckb-vm memory size: {}",
-        memory_size
-    );
-    println!("Base memory: {}", get_current_memory());
-    for _ in 0..G_CHECK_LOOP {
-        let result = run_with_memory::<u64, FlatMemory<u64>>(
-            &Bytes::from(BIN_PATH_BUFFER),
-            &vec![Bytes::from(BIN_NAME)],
-            memory_size,
-        );
-        assert!(result.is_ok());
-        assert_eq!(result.unwrap(), 0);
-        println!("Current memory: {}", get_current_memory());
-    }
-    println!("End of check");
+    println!("Check interpreter: done memory used {}", ALLOC.used());
     Ok(())
 }
 
 #[cfg(has_asm)]
-fn check_asm(memory_size: usize) -> Result<(), ()> {
-    println!(
-        "Check asm memory overflow, ckb-vm memory size: {}",
-        memory_size
-    );
-    println!("Base memory: {}", get_current_memory());
+fn check_asm() -> Result<(), Box<dyn std::error::Error>> {
+    println!("Check asm: init",);
+    println!("Check asm: base memory used {}", ALLOC.used());
     for _ in 0..G_CHECK_LOOP {
         let asm_core = <AsmCoreMachine as SupportMachine>::new_with_memory(
             ISA_IMC,
             VERSION0,
             u64::MAX,
-            memory_size,
+            DEFAULT_MEMORY_SIZE,
         );
         let core = AsmDefaultMachineBuilder::new(asm_core).build();
         let mut machine = AsmMachine::new(core);
@@ -181,84 +114,15 @@ fn check_asm(memory_size: usize) -> Result<(), ()> {
         let result = machine.run();
         assert!(result.is_ok());
         assert_eq!(result.unwrap(), 0);
-
-        println!("Current memory: {}", get_current_memory());
+        println!("Check asm: step memory used {}", ALLOC.used());
     }
-    println!("End of check");
+    println!("Check asm: done memory used {}", ALLOC.used());
     Ok(())
 }
 
-#[cfg(has_asm)]
-fn check_asm_in_thread(memory_size: usize) -> Result<(), ()> {
-    println!(
-        "Check asm in thread memory overflow, ckb-vm memory size: {}",
-        memory_size
-    );
-    println!("Base memory: {}", get_current_memory());
-    for _ in 0..G_CHECK_LOOP {
-        let asm_core = <AsmCoreMachine as SupportMachine>::new_with_memory(
-            ISA_IMC,
-            VERSION0,
-            u64::MAX,
-            memory_size,
-        );
-        let core = AsmDefaultMachineBuilder::new(asm_core).build();
-        let mut machine = AsmMachine::new(core);
-        machine
-            .load_program(
-                &Bytes::from(BIN_PATH_BUFFER),
-                [Ok(Bytes::from(BIN_NAME))].into_iter(),
-            )
-            .unwrap();
-        let thread_join_handle = thread::spawn(move || {
-            let result = machine.run();
-            assert!(result.is_ok());
-            assert_eq!(result.unwrap(), 0);
-        });
-        thread_join_handle.join().unwrap();
-        println!("Current memory: {}", get_current_memory());
-    }
-    println!("End of check");
-    Ok(())
-}
-
-fn test_memory(memory_size: usize) -> Result<(), ()> {
-    if check_interpreter(memory_size).is_err() {
-        return Err(());
-    }
-    if check_falt(memory_size).is_err() {
-        return Err(());
-    }
-
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    check_interpreter()?;
     #[cfg(has_asm)]
-    if check_asm(memory_size).is_err() {
-        return Err(());
-    }
-
-    #[cfg(has_asm)]
-    if check_asm_in_thread(memory_size).is_err() {
-        return Err(());
-    }
+    check_asm()?;
     Ok(())
 }
-
-fn main() {
-    #[cfg(not(target_os = "windows"))]
-    jemalloc_ctl::epoch::advance().unwrap();
-
-    let memory_overflow = MemoryOverflow::new();
-
-    let memory_size = 1024 * 1024 * 4;
-    if test_memory(memory_size).is_err() {
-        panic!("run testcase failed");
-    }
-
-    memory_overflow.check();
-
-    let memory_size = 1024 * 1024 * 2;
-    if test_memory(memory_size).is_err() {
-        panic!("run testcase failed");
-    }
-
-    memory_overflow.check();
-}