abby/kv_bench
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

adapt new mace 0.0.14

Browse Source
This commit is contained in:
abbycin 2025-10-12 17:09:55 +08:00
parent d9b969f644
commit 07baaae51e
Signed by: abby
GPG Key ID: B636E0F0307EF8EB
9 changed files with 377 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
.gitignore vendored
View File

@ -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

58
.vscode/launch.json vendored Normal file
View File

@ -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}"
}
]
}

11
Cargo.toml
View File

@ -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

2
heap_trace/.gitignore vendored Normal file
View File

@ -0,0 +1,2 @@
/target
Cargo.lock

7
heap_trace/Cargo.toml Normal file
View File

@ -0,0 +1,7 @@
[package]
name = "myalloc"
version = "0.1.0"
edition = "2024"
[dependencies]
backtrace = "0.3.76"

233
heap_trace/src/lib.rs Normal file
View File

@ -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<String> {
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<Mutex<HashMap<String, Status>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
static G_TRACE: LazyLock<Mutex<HashMap<u64, String>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
const META_LEN: usize = 8;
thread_local! {
static G_SELF: Cell<bool> = 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::<u64>());
let sz = real_size(&layout);
std::alloc::Layout::from_size_align(sz, align).unwrap()
}
fn write_hash(x: *mut u8, align: usize, s: Option<String>) -> *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::<u64>();
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::<u64>();
(p.read_unaligned(), p.cast::<u8>())
};
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: 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: 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));
}

1
scripts/init.sh
View File

@ -2,3 +2,4 @@
python3 -m venv .
./bin/pip3 install pandas matplotlib adjustText
rm -f .gitignore

53
scripts/mem_analyze.py Normal file
View File

@ -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}")

15
src/main.rs
View File

@ -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));
}