diff --git a/.gitignore b/.gitignore index a9e4dbf..40096e6 100644 --- a/.gitignore +++ b/.gitignore @@ -3,9 +3,12 @@ /scripts/include /scripts/lib /scripts/lib64 -/scripts/shares/ +/scripts/share/ /scripts/.gitignore *.png *.csv pyvenv.cfg Cargo.lock +free.txt +alloc.txt +*.zst diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..34214f9 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,58 @@ +{ + // Use IntelliSense to learn about possible attributes. + // Hover to view descriptions of existing attributes. + // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 + "version": "0.2.0", + "configurations": [ + { + "type": "lldb", + "request": "launch", + "name": "Debug executable 'kv_bench'", + "cargo": { + "args": [ + "build", + "--bin=kv_bench", + "--package=kv_bench" + ], + "filter": { + "name": "kv_bench", + "kind": "bin" + } + }, + "args": [ + "--path", + "/home/abby/mace_bench", + "--threads", + "1", + "--iterations", + "100000", + "--mode", + "insert", + "--key-size", + "1024", + "--value-size", + "16" + ], + "cwd": "${workspaceFolder}" + }, + { + "type": "lldb", + "request": "launch", + "name": "Debug unit tests in executable 'kv_bench'", + "cargo": { + "args": [ + "test", + "--no-run", + "--bin=kv_bench", + "--package=kv_bench" + ], + "filter": { + "name": "kv_bench", + "kind": "bin" + } + }, + "args": [], + "cwd": "${workspaceFolder}" + } + ] +} \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index 36ac058..89a96c8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,11 +5,16 @@ edition = "2024" [dependencies] mace = { git = "https://github.com/abbycin/mace" } -clap = { version = "4.5.42", features = ["derive"] } +clap = { version = "4.5.48", features = ["derive"] } rand = "0.9.2" log = "0.4.22" -coreid = { path = "coreid"} -logger = { path = "logger"} +coreid = { path = "coreid" } +logger = { path = "logger" } +myalloc = { path = "heap_trace" } + +[features] +default = [] +custom_alloc = [] [profile.release] lto = true diff --git a/heap_trace/.gitignore b/heap_trace/.gitignore new file mode 100644 index 0000000..869df07 --- /dev/null +++ b/heap_trace/.gitignore @@ -0,0 +1,2 @@ +/target +Cargo.lock \ No newline at end of file diff --git a/heap_trace/Cargo.toml b/heap_trace/Cargo.toml new file mode 100644 index 0000000..d9931bc --- /dev/null +++ b/heap_trace/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "myalloc" +version = "0.1.0" +edition = "2024" + +[dependencies] +backtrace = "0.3.76" diff --git a/heap_trace/src/lib.rs b/heap_trace/src/lib.rs new file mode 100644 index 0000000..4cbfbbf --- /dev/null +++ b/heap_trace/src/lib.rs @@ -0,0 +1,233 @@ +use std::{ + alloc::{GlobalAlloc, System}, + cell::Cell, + collections::{HashMap, hash_map::Entry}, + fmt::Display, + hash::{DefaultHasher, Hash, Hasher}, + ptr, + sync::{LazyLock, Mutex, atomic::AtomicBool}, +}; + +pub struct MyAlloc; + +fn trace(size: usize, is_alloc: bool) -> Option { + let mut key = String::new(); + backtrace::trace(|f| { + backtrace::resolve_frame(f, |sym| { + if let Some(filename) = sym.filename() + && let Some(line) = sym.lineno() + { + if let Some(name) = filename.to_str() + && name.contains("mace") + { + if name.len() > 10 { + // sometime name maybe empty + let x = format!("{}:{}\n", name, line); + key.extend(x.chars().into_iter()); + } + } + } + }); + true + }); + if !key.is_empty() { + let mut lk = G_MAP.lock().unwrap(); + let tmp = key.clone(); + match lk.entry(tmp) { + Entry::Vacant(v) => { + if is_alloc { + v.insert(Status { + nr_alloc: 1, + alloc_size: size, + nr_free: 0, + free_size: 0, + }); + } else { + v.insert(Status { + nr_alloc: 0, + alloc_size: 0, + nr_free: 1, + free_size: size, + }); + } + } + Entry::Occupied(mut o) => { + let s = o.get_mut(); + if is_alloc { + s.nr_alloc += 1; + s.alloc_size += size; + } else { + s.nr_free += 1; + s.free_size += size; + } + } + } + Some(key) + } else { + None + } +} + +#[derive(Debug)] +pub struct Status { + nr_alloc: usize, + alloc_size: usize, + nr_free: usize, + free_size: usize, +} + +impl Display for Status { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_fmt(format_args!("{:?}", self)) + } +} + +static G_STOP: AtomicBool = AtomicBool::new(false); + +static G_MAP: LazyLock>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +static G_TRACE: LazyLock>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +const META_LEN: usize = 8; + +thread_local! { + static G_SELF: Cell = const { Cell::new(false) }; +} + +const fn real_size(layout: &std::alloc::Layout) -> usize { + if META_LEN > layout.align() { + META_LEN.checked_add(layout.size()).unwrap() + } else { + layout.align().checked_add(layout.size()).unwrap() + } +} + +fn new_layout(layout: std::alloc::Layout) -> std::alloc::Layout { + let align = layout.align().max(align_of::()); + let sz = real_size(&layout); + std::alloc::Layout::from_size_align(sz, align).unwrap() +} + +fn write_hash(x: *mut u8, align: usize, s: Option) -> *mut u8 { + let r = unsafe { x.add(META_LEN.max(align)) }; + if !G_SELF.with(|x| x.get()) { + G_SELF.with(|x| x.set(true)); + let p = x.cast::(); + if let Some(s) = s { + let mut stat = DefaultHasher::new(); + s.hash(&mut stat); + let h = stat.finish(); + unsafe { p.write_unaligned(h) }; + let mut lk = G_TRACE.lock().unwrap(); + lk.insert(h, s); + } else { + unsafe { p.write_unaligned(u64::MAX) }; + } + G_SELF.with(|x| x.set(false)); + } + r +} + +fn read_hash(x: *mut u8, align: usize) -> *mut u8 { + let (h, p) = unsafe { + let p = x.sub(META_LEN.max(align)).cast::(); + (p.read_unaligned(), p.cast::()) + }; + + if h == u64::MAX { + return p; + } + if !G_SELF.with(|x| x.get()) { + G_SELF.with(|x| x.set(true)); + let mut lk = G_TRACE.lock().unwrap(); + lk.remove(&h); + G_SELF.with(|x| x.set(false)); + } + p +} + +unsafe impl GlobalAlloc for MyAlloc { + unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 { + let s = if !G_SELF.with(|x| x.get()) && !G_STOP.load(std::sync::atomic::Ordering::Acquire) { + G_SELF.with(|x| x.set(true)); + let x = trace(layout.size(), true); + G_SELF.with(|x| x.set(false)); + x + } else { + None + }; + + let new = new_layout(layout); + let x = unsafe { System.alloc(new) }; + write_hash(x, new.align(), s) + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) { + if !G_SELF.with(|x| x.get()) && !G_STOP.load(std::sync::atomic::Ordering::Acquire) { + G_SELF.with(|x| x.set(true)); + trace(layout.size(), false); + G_SELF.with(|x| x.set(false)); + } + let new = new_layout(layout); + let p = read_hash(ptr, new.align()); + unsafe { System.dealloc(p, new) }; + } + + unsafe fn alloc_zeroed(&self, layout: std::alloc::Layout) -> *mut u8 { + let p = unsafe { self.alloc(layout) }; + if !p.is_null() { + unsafe { ptr::write_bytes(p, 0, layout.size()) }; + } + p + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 { + let s = if !G_SELF.with(|x| x.get()) && !G_STOP.load(std::sync::atomic::Ordering::Acquire) { + G_SELF.with(|x| x.set(true)); + let x = trace(layout.size(), true); + G_SELF.with(|x| x.set(false)); + x + } else { + None + }; + + unsafe { + let old_layout = new_layout(layout); + let raw = ptr.sub(META_LEN.max(old_layout.align())); + let new_total_size = META_LEN + new_size; + + let new_raw = System.realloc(raw, old_layout, new_total_size); + if new_raw.is_null() { + return new_raw; + } + write_hash(new_raw, old_layout.align(), s) + } + } +} + +pub fn print_filtered_trace(f: F) +where + F: Fn(&str, &Status), +{ + G_STOP.store(true, std::sync::atomic::Ordering::Release); + let lk = G_MAP.lock().unwrap(); + let t = G_TRACE.lock().unwrap(); + + for (_, v) in t.iter() { + if let Some(s) = lk.get(v) { + f(v, s); + } + } +} + +pub fn print_all_trace(f: F) +where + F: Fn(&str, &Status), +{ + G_STOP.store(true, std::sync::atomic::Ordering::Release); + let lk = G_MAP.lock().unwrap(); + + lk.iter().for_each(|(k, v)| f(k, v)); +} diff --git a/scripts/init.sh b/scripts/init.sh index 0585fd6..881ade4 100755 --- a/scripts/init.sh +++ b/scripts/init.sh @@ -2,3 +2,4 @@ python3 -m venv . ./bin/pip3 install pandas matplotlib adjustText +rm -f .gitignore diff --git a/scripts/mem_analyze.py b/scripts/mem_analyze.py new file mode 100644 index 0000000..2eb1aa2 --- /dev/null +++ b/scripts/mem_analyze.py @@ -0,0 +1,53 @@ +#!/usr/bin/python3 + +import sys + +assert(len(sys.argv) == 2) + +f = open(sys.argv[1]) +lines = [] + +allocs = [] + +while True: + line = f.readline() + if len(line) < 10: + break + if line.find('INFO') != -1: + continue + pos = line.find('Status') + if pos < 0: + pos = line.find('mace') + if pos != 0: + lines.append(line[pos:]) + else: + lines.append(line) + else: + cleaned = line[pos+6:].strip().strip('{}') + pairs = cleaned.split(',') + tl = [tuple(pair.split(': ')) for pair in pairs] + tl = [(k.strip(), int(v)) for k, v in tl] + allocs.append((tl, ''.join(lines))) + lines.clear() + +# sort by alloc_size +allocs.sort(key=lambda x: x[0][1][1], reverse=True) + +with open('alloc.txt', 'w') as o: + for x in allocs: + o.write(f'{x[0]}\n{x[1]}\n') +# sort by free_size +allocs.sort(key=lambda x: x[0][3][1], reverse=True) + +with open('free.txt', 'w') as o: + for x in allocs: + o.write(f'{x[0]}\n{x[1]}\n') + +alloc_size = 0 +free_size = 0 + +for x in allocs: + alloc_size += x[0][1][1] + free_size += x[0][3][1] + +print(f"total_alloc {alloc_size} total_free {free_size}") diff --git a/src/main.rs b/src/main.rs index 08adc6d..52b6854 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,8 @@ use clap::Parser; #[cfg(target_os = "linux")] use logger::Logger; use mace::{Mace, Options}; +#[cfg(feature = "custom_alloc")] +use myalloc::{MyAlloc, print_filtered_trace}; use rand::prelude::*; use std::path::Path; use std::process::exit; @@ -9,6 +11,10 @@ use std::sync::Arc; use std::thread::JoinHandle; use std::time::Instant; +#[cfg(feature = "custom_alloc")] +#[global_allocator] +static GLOBAL: MyAlloc = MyAlloc; + #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] struct Args { @@ -40,7 +46,7 @@ struct Args { fn main() { #[cfg(target_os = "linux")] { - Logger::init().add_file("/tmp/x.log", true); + Logger::init().add_file("/Data/x.log", true); log::set_max_level(log::LevelFilter::Info); } let args = Args::parse(); @@ -71,11 +77,11 @@ fn main() { let mut opt = Options::new(path); opt.sync_on_write = false; opt.tmp_store = args.mode != "get"; - opt.gc_timeout = 1000 * 60; // make sure GC will not work let mut saved = opt.clone(); saved.tmp_store = false; // opt.cache_capacity = 3 << 30; // this is very important for large key-value store let mut db = Mace::new(opt.validate().unwrap()).unwrap(); + db.disable_gc(); let mut rng = rand::rng(); let value = Arc::new(vec![b'0'; args.value_size]); @@ -100,8 +106,6 @@ fn main() { } }); pre_tx.commit().unwrap(); - drop(pre_tx); - log::info!("====="); drop(db); // re-open db saved.tmp_store = true; @@ -206,4 +210,7 @@ fn main() { ops, duration.as_millis() ); + drop(db); + #[cfg(feature = "custom_alloc")] + print_filtered_trace(|x, y| log::info!("{}{}", x, y)); }